diff --git a/res/drawable/color_chip.xml b/res/drawable/color_chip.xml
index bc09992..11301b1 100644
--- a/res/drawable/color_chip.xml
+++ b/res/drawable/color_chip.xml
@@ -18,10 +18,10 @@
 <selector xmlns:android="http://schemas.android.com/apk/res/android">
     <item
         android:state_activated="true"
-        android:drawable="@drawable/color_chip_filled" />
-    <item
-        android:state_activated="false"
         android:drawable="@drawable/color_chip_hollow" />
     <item
-        android:drawable="@drawable/color_chip_hollow"/>
+        android:state_activated="false"
+        android:drawable="@drawable/color_chip_filled" />
+    <item
+        android:drawable="@drawable/color_chip_filled"/>
 </selector>
\ No newline at end of file
diff --git a/res/drawable/color_chip_hollow.xml b/res/drawable/color_chip_hollow.xml
index c0eab74..699dc60 100644
--- a/res/drawable/color_chip_hollow.xml
+++ b/res/drawable/color_chip_hollow.xml
@@ -15,10 +15,27 @@
      See the License for the specific language governing permissions and
      limitations under the License.
 -->
-<shape xmlns:android="http://schemas.android.com/apk/res/android"
-        android:shape="ring"
-        android:innerRadius="10dp"
-        android:thickness="6dp"
-        android:useLevel="false">
-    <solid android:color="@android:color/black"/>
-</shape>
+<layer-list xmlns:android="http://schemas.android.com/apk/res/android">
+    <item
+        android:height="@dimen/component_color_chip_size"
+        android:width="@dimen/component_color_chip_size"
+        android:gravity="center">
+        <shape
+            android:shape="oval">
+            <solid android:color="@android:color/black" />
+        </shape>
+    </item>
+    <item
+        android:height="@dimen/component_color_chip_container_size"
+        android:width="@dimen/component_color_chip_container_size"
+        android:gravity="center">
+        <shape
+            android:shape="ring"
+            android:innerRadius="24dp"
+            android:thickness="@dimen/option_border_width"
+            android:useLevel="false">
+        <solid android:color="@android:color/black"/>
+        </shape>
+    </item>
+</layer-list>
+
diff --git a/res/drawable/option_border_edge_custom.xml b/res/drawable/option_border_edge_custom.xml
index 1863c5d..6d633c3 100644
--- a/res/drawable/option_border_edge_custom.xml
+++ b/res/drawable/option_border_edge_custom.xml
@@ -26,8 +26,7 @@
     <item android:state_activated="false">
         <shape android:shape="rectangle">
             <stroke
-                android:color="@color/black_24_alpha"
-                android:alpha="0.24"
+                android:color="@color/option_border_color"
                 android:width="@dimen/option_border_width" />
             <corners android:radius="?android:dialogCornerRadius" />
         </shape>
diff --git a/res/layout/fragment_custom_theme_component.xml b/res/layout/fragment_custom_theme_component.xml
index cf80192..fb3e5b3 100644
--- a/res/layout/fragment_custom_theme_component.xml
+++ b/res/layout/fragment_custom_theme_component.xml
@@ -33,13 +33,13 @@
             android:layout_height="0dp"
             android:background="?android:colorPrimary"
             app:layout_constrainedHeight="true"
-            app:layout_constraintBottom_toTopOf="@+id/guideline"
+            app:layout_constraintBottom_toTopOf="@+id/divider"
             app:layout_constraintEnd_toEndOf="parent"
             app:layout_constraintHeight_max="@dimen/preview_pager_max_height"
-            app:layout_constraintHorizontal_bias="0.5"
             app:layout_constraintStart_toStartOf="parent"
             app:layout_constraintTop_toTopOf="parent"
-            app:layout_constraintVertical_bias="0.0">
+            app:layout_constraintVertical_bias="0.0"
+            app:layout_constraintHeight_percent="@dimen/preview_pager_maximum_height_ratio">
 
             <include
                 android:id="@+id/component_preview_content"
@@ -50,6 +50,7 @@
                 android:layout_marginHorizontal="@dimen/preview_page_horizontal_margin"
                 android:layout_marginTop="@dimen/preview_page_top_margin"/>
         </FrameLayout>
+
         <View
             android:id="@+id/divider"
             android:layout_width="match_parent"
@@ -58,43 +59,31 @@
             app:layout_constraintEnd_toEndOf="parent"
             app:layout_constraintStart_toStartOf="parent"
             app:layout_constraintTop_toBottomOf="@+id/component_preview_container"/>
-
-        <androidx.constraintlayout.widget.Guideline
-            android:id="@+id/guideline"
-            android:layout_width="wrap_content"
-            android:layout_height="wrap_content"
-            android:orientation="horizontal"
-            app:layout_constraintGuide_percent=".7"/>
-
-
+        
         <TextView
             android:id="@+id/component_options_title"
             android:layout_width="match_parent"
             android:layout_height="wrap_content"
-            android:layout_marginTop="10dp"
-            android:layout_marginStart="10dp"
-            android:layout_marginEnd="10dp"
-            android:layout_marginBottom="10dp"
+            android:layout_marginVertical="18dp"
+            android:layout_marginHorizontal="16dp"
             android:textAlignment="center"
             android:textAppearance="@style/TitleTextAppearance"
             android:textSize="@dimen/component_options_title_size"
             app:layout_constraintEnd_toEndOf="parent"
             app:layout_constraintStart_toStartOf="parent"
             app:layout_constraintTop_toBottomOf="@+id/component_preview_container"
-            app:layout_constraintBottom_toBottomOf="parent"
-            app:layout_constraintVertical_bias=".15"/>
+            app:layout_constraintBottom_toTopOf="@+id/options_container"/>
 
         <androidx.recyclerview.widget.RecyclerView
             android:id="@+id/options_container"
             android:layout_width="match_parent"
-            android:layout_height="@dimen/options_container_height"
+            android:layout_height="0dp"
             android:layout_gravity="center_horizontal"
             app:layout_constraintBottom_toBottomOf="parent"
             app:layout_constraintEnd_toEndOf="parent"
-            app:layout_constraintHorizontal_bias="0.5"
             app:layout_constraintStart_toStartOf="parent"
             app:layout_constraintTop_toBottomOf="@+id/component_options_title"
-            app:layout_constraintVertical_bias=".25"/>
+            app:layout_constraintVertical_bias="1"/>
 
     </androidx.constraintlayout.widget.ConstraintLayout>
 </LinearLayout>
diff --git a/res/layout/fragment_custom_theme_name.xml b/res/layout/fragment_custom_theme_name.xml
index 27a8f19..ab3c459 100644
--- a/res/layout/fragment_custom_theme_name.xml
+++ b/res/layout/fragment_custom_theme_name.xml
@@ -36,55 +36,51 @@
             android:clipToPadding="false"
             android:background="@color/fullscreen_preview_background"
             app:layout_constrainedHeight="true"
-            app:layout_constraintBottom_toTopOf="@+id/guideline"
+            app:layout_constraintBottom_toTopOf="@+id/component_options_title"
             app:layout_constraintEnd_toEndOf="parent"
             app:layout_constraintHeight_max="@dimen/preview_pager_max_height"
-            app:layout_constraintHorizontal_bias="0.5"
             app:layout_constraintStart_toStartOf="parent"
             app:layout_constraintTop_toTopOf="parent"
-            app:layout_constraintVertical_bias="0.0">
+            app:layout_constraintVertical_bias="0.0"
+            app:layout_constraintHeight_percent="@dimen/preview_pager_maximum_height_ratio">
+
             <include layout="@layout/theme_preview_card"/>
         </FrameLayout>
 
-        <androidx.constraintlayout.widget.Guideline
-            android:id="@+id/guideline"
-            android:layout_width="wrap_content"
-            android:layout_height="wrap_content"
-            android:orientation="horizontal"
-            app:layout_constraintGuide_percent=".7"/>
-
         <TextView
             android:id="@+id/component_options_title"
             android:layout_width="match_parent"
             android:layout_height="wrap_content"
-            android:layout_marginTop="10dp"
-            android:layout_marginStart="10dp"
-            android:layout_marginEnd="10dp"
-            android:layout_marginBottom="10dp"
+            android:layout_marginVertical="18dp"
+            android:layout_marginHorizontal="16dp"
             android:textAlignment="center"
             android:textAppearance="@style/TitleTextAppearance"
             android:textSize="@dimen/component_options_title_size"
             app:layout_constraintEnd_toEndOf="parent"
             app:layout_constraintStart_toStartOf="parent"
             app:layout_constraintTop_toBottomOf="@+id/component_preview_container"
-            app:layout_constraintBottom_toBottomOf="parent"
-            app:layout_constraintVertical_bias=".15"/>
+            app:layout_constraintBottom_toTopOf="@+id/edit_text_container"/>
 
-        <EditText
-            android:id="@+id/custom_theme_name"
-            style="@style/CustomThemeNameEditText"
-            android:layout_width="wrap_content"
-            android:layout_height="wrap_content"
-            android:importantForAutofill="no"
-            android:minWidth="300dp"
-            app:layout_constraintBottom_toBottomOf="parent"
+        <FrameLayout
+            android:id="@+id/edit_text_container"
+            android:layout_width="match_parent"
+            android:layout_height="0dp"
             app:layout_constraintEnd_toEndOf="parent"
-            app:layout_constraintHorizontal_bias="0.5"
             app:layout_constraintStart_toStartOf="parent"
             app:layout_constraintTop_toBottomOf="@+id/component_options_title"
-            app:layout_constraintVertical_bias=".25"/>
+            app:layout_constraintBottom_toBottomOf="parent"
+            app:layout_constraintVertical_bias="1">
 
-
+            <EditText
+                android:id="@+id/custom_theme_name"
+                android:layout_width="wrap_content"
+                android:layout_height="wrap_content"
+                android:layout_marginVertical="16dp"
+                android:layout_marginHorizontal="16dp"
+                android:layout_gravity="center|top"
+                android:importantForAutofill="no"
+                android:minWidth="300dp"
+                style="@style/CustomThemeNameEditText"/>
+        </FrameLayout>
     </androidx.constraintlayout.widget.ConstraintLayout>
-
 </LinearLayout>
diff --git a/res/layout/theme_color_option.xml b/res/layout/theme_color_option.xml
index 33e4668..9a3c4c4 100644
--- a/res/layout/theme_color_option.xml
+++ b/res/layout/theme_color_option.xml
@@ -15,15 +15,17 @@
      limitations under the License.
 -->
 <FrameLayout xmlns:android="http://schemas.android.com/apk/res/android"
-    android:layout_width="match_parent"
+    android:layout_width="wrap_content"
     android:layout_height="wrap_content"
     android:layout_gravity="center"
-    android:padding="8dp">
+    android:layout_marginTop="28dp"
+    android:layout_marginHorizontal="@dimen/component_options_margin_horizontal">
 
     <ImageView
         android:id="@+id/option_tile"
         android:layout_width="@dimen/component_color_chip_container_size"
         android:layout_height="@dimen/component_color_chip_container_size"
         android:layout_gravity="center"
+        android:scaleType="center"
         android:src="@drawable/color_chip"/>
 </FrameLayout>
diff --git a/res/layout/theme_font_option.xml b/res/layout/theme_font_option.xml
index 48dcaa2..dea4f78 100644
--- a/res/layout/theme_font_option.xml
+++ b/res/layout/theme_font_option.xml
@@ -19,13 +19,6 @@
               android:layout_height="wrap_content"
               android:orientation="vertical">
 
-    <TextView
-        android:id="@+id/option_label"
-        android:layout_width="wrap_content"
-        android:layout_height="wrap_content"
-        android:layout_gravity="center_horizontal"
-        android:layout_marginBottom="@dimen/theme_option_label_margin"
-        android:textAppearance="@style/OptionTitleTextAppearance"/>
     <FrameLayout
         android:id="@+id/option_tile"
         android:layout_width="@dimen/option_tile_width"
@@ -33,6 +26,7 @@
         android:layout_gravity="center_horizontal"
         android:paddingHorizontal="@dimen/option_tile_padding_horizontal"
         android:paddingVertical="@dimen/option_tile_padding_vertical"
+        android:layout_marginHorizontal="@dimen/component_options_margin_horizontal"
         android:background="@drawable/option_border">
         <TextView
             android:id="@+id/thumbnail_text"
diff --git a/res/layout/theme_icon_option.xml b/res/layout/theme_icon_option.xml
index 62925ad..292b8cd 100644
--- a/res/layout/theme_icon_option.xml
+++ b/res/layout/theme_icon_option.xml
@@ -16,16 +16,9 @@
 -->
 <LinearLayout xmlns:android="http://schemas.android.com/apk/res/android"
               android:layout_width="wrap_content"
-              android:layout_height="match_parent"
+              android:layout_height="wrap_content"
               android:orientation="vertical">
 
-    <TextView
-        android:id="@+id/option_label"
-        android:layout_width="wrap_content"
-        android:layout_height="wrap_content"
-        android:layout_gravity="center_horizontal"
-        android:layout_marginBottom="@dimen/theme_option_label_margin"
-        android:textAppearance="@style/OptionTitleTextAppearance"/>
     <FrameLayout
         android:id="@+id/option_tile"
         android:layout_width="@dimen/option_tile_width"
@@ -33,6 +26,7 @@
         android:layout_gravity="center_horizontal"
         android:paddingHorizontal="@dimen/option_tile_padding_horizontal"
         android:paddingVertical="@dimen/option_tile_padding_vertical"
+        android:layout_marginHorizontal="@dimen/component_options_margin_horizontal"
         android:background="@drawable/option_border">
         <ImageView
             android:id="@+id/option_icon"
diff --git a/res/layout/theme_shape_option.xml b/res/layout/theme_shape_option.xml
index 94e31ea..c5682c0 100644
--- a/res/layout/theme_shape_option.xml
+++ b/res/layout/theme_shape_option.xml
@@ -16,15 +16,14 @@
 -->
 <LinearLayout xmlns:android="http://schemas.android.com/apk/res/android"
               android:layout_width="wrap_content"
-              android:layout_height="match_parent"
+              android:layout_height="wrap_content"
               android:orientation="vertical">
     <FrameLayout
         android:id="@+id/option_tile"
         android:layout_width="@dimen/option_tile_width"
         android:layout_height="@dimen/option_tile_width"
         android:layout_gravity="center"
-        android:layout_marginTop="16dp"
-        android:padding="2dp">
+        android:layout_marginHorizontal="@dimen/component_options_margin_horizontal">
         <ImageView
             android:id="@+id/shape_thumbnail"
             android:layout_width="match_parent"
diff --git a/res/values-land/bools.xml b/res/values-land/bools.xml
deleted file mode 100644
index 409538d..0000000
--- a/res/values-land/bools.xml
+++ /dev/null
@@ -1,20 +0,0 @@
-<?xml version="1.0" encoding="utf-8"?>
-<!--
-
-     Copyright (C) 2019 The Android Open Source Project
-
-     Licensed under the Apache License, Version 2.0 (the "License");
-     you may not use this file except in compliance with the License.
-     You may obtain a copy of the License at
-
-          http://www.apache.org/licenses/LICENSE-2.0
-
-     Unless required by applicable law or agreed to in writing, software
-     distributed under the License is distributed on an "AS IS" BASIS,
-     WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-     See the License for the specific language governing permissions and
-     limitations under the License.
--->
-<resources>
-    <bool name="use_grid_for_options">true</bool>
-</resources>
\ No newline at end of file
diff --git a/res/values-night/colors.xml b/res/values-night/colors.xml
index feaf26f..c6b2a3f 100644
--- a/res/values-night/colors.xml
+++ b/res/values-night/colors.xml
@@ -24,8 +24,6 @@
     <color name="switch_track_tint">#171717</color>
     <color name="switch_thumb_tint">#bdbdbd</color>
 
-    <color name="theme_preview_workspace_shadow_color">@android:color/transparent</color>
-
     <color name="tile_enabled_icon_color">#2d2d2d</color>
     <color name="tile_disabled_icon_color">@color/light_grey</color>
     <color name="tile_disabled_background_color">#353535</color>
diff --git a/res/values/bools.xml b/res/values/bools.xml
deleted file mode 100644
index 409538d..0000000
--- a/res/values/bools.xml
+++ /dev/null
@@ -1,20 +0,0 @@
-<?xml version="1.0" encoding="utf-8"?>
-<!--
-
-     Copyright (C) 2019 The Android Open Source Project
-
-     Licensed under the Apache License, Version 2.0 (the "License");
-     you may not use this file except in compliance with the License.
-     You may obtain a copy of the License at
-
-          http://www.apache.org/licenses/LICENSE-2.0
-
-     Unless required by applicable law or agreed to in writing, software
-     distributed under the License is distributed on an "AS IS" BASIS,
-     WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-     See the License for the specific language governing permissions and
-     limitations under the License.
--->
-<resources>
-    <bool name="use_grid_for_options">true</bool>
-</resources>
\ No newline at end of file
diff --git a/res/values/colors.xml b/res/values/colors.xml
index d4ad0e1..82d02b0 100644
--- a/res/values/colors.xml
+++ b/res/values/colors.xml
@@ -31,7 +31,8 @@
     <color name="clockface_preview_background">@android:color/black</color>
 
     <color name="theme_preview_icon_color">@color/google_grey700</color>
-    <color name="theme_preview_workspace_shadow_color">#B0000000</color>
+    <color name="theme_preview_workspace_shadow_color_dark">#B0000000</color>
+    <color name="theme_preview_workspace_shadow_color_transparent">@android:color/transparent</color>
 
     <color name="text_color_dark">#2d2d2d</color>
     <color name="text_color_light">@color/material_white_text</color>
diff --git a/res/values/dimens.xml b/res/values/dimens.xml
index 461b39e..4f78920 100644
--- a/res/values/dimens.xml
+++ b/res/values/dimens.xml
@@ -99,15 +99,16 @@
     <dimen name="custom_theme_nav_height">56dp</dimen>
 
     <!-- Note, using dp instead of sp as this is just the font thumbnail, not text -->
-    <dimen name="font_comonent_option_thumbnail_size">32dp</dimen>
+    <dimen name="font_comonent_option_thumbnail_size">36dp</dimen>
 
     <dimen name="component_icon_thumb_size">40dp</dimen>
 
-    <dimen name="component_color_chip_size">32dp</dimen>
-    <dimen name="component_color_chip_container_size">33dp</dimen>
+    <dimen name="component_color_chip_size">40dp</dimen>
+    <dimen name="component_color_chip_container_size">52dp</dimen>
 
     <dimen name="component_shape_thumb_size">72dp</dimen>
-    <dimen name="component_options_title_size">18sp</dimen>
+    <dimen name="component_options_title_size">20sp</dimen>
+    <dimen name="component_options_margin_horizontal">4dp</dimen>
 
     <!-- For a corner radius of this size or larger, we'll preview a rounded qsb widget. -->
     <dimen name="roundCornerThreshold">16dp</dimen>
diff --git a/res/values/strings.xml b/res/values/strings.xml
index f1cd169..d80f74d 100755
--- a/res/values/strings.xml
+++ b/res/values/strings.xml
@@ -30,10 +30,6 @@
         the home screen. [CHAR LIMIT=15] -->
     <string name="grid_title">Grid</string>
 
-    <!-- Title of a section of the customization picker where the user can select a Wallpaper.
-        [CHAR LIMIT=19] -->
-    <string name="wallpaper_title">Wallpaper</string>
-
     <!-- Label for a button that allows the user to apply the currently selected Theme.
         [CHAR LIMIT=20] -->
     <string name="apply_theme_btn">Apply</string>
diff --git a/src/com/android/customization/model/grid/GridOption.java b/src/com/android/customization/model/grid/GridOption.java
index 2f32ce1..43afee4 100644
--- a/src/com/android/customization/model/grid/GridOption.java
+++ b/src/com/android/customization/model/grid/GridOption.java
@@ -20,9 +20,12 @@
 import android.net.Uri;
 import android.os.Parcel;
 import android.os.Parcelable;
+import android.text.TextUtils;
 import android.view.View;
 import android.widget.ImageView;
 
+import androidx.annotation.Nullable;
+
 import com.android.customization.model.CustomizationManager;
 import com.android.customization.model.CustomizationOption;
 import com.android.customization.widget.GridTileDrawable;
@@ -31,7 +34,6 @@
 /**
  * Represents a grid layout option available in the current launcher.
  */
-// TODO(chihhangchuang): Consider moving Parcelable into CustomizationOption.
 public class GridOption implements CustomizationOption<GridOption>, Parcelable {
     public static final Creator<GridOption> CREATOR = new Creator<GridOption>() {
         @Override
@@ -101,6 +103,21 @@
     }
 
     @Override
+    public boolean equals(@Nullable Object obj) {
+        if (this == obj) {
+            return true;
+        }
+
+        if (obj instanceof GridOption) {
+            GridOption other = (GridOption) obj;
+            return TextUtils.equals(this.name, other.name)
+                    && this.cols == other.cols
+                    && this.rows == other.rows;
+        }
+        return false;
+    }
+
+    @Override
     public int getLayoutResId() {
         return R.layout.grid_option;
     }
diff --git a/src/com/android/customization/model/theme/ThemeBundle.java b/src/com/android/customization/model/theme/ThemeBundle.java
index 8cb2865..43b07ff 100644
--- a/src/com/android/customization/model/theme/ThemeBundle.java
+++ b/src/com/android/customization/model/theme/ThemeBundle.java
@@ -62,6 +62,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
+import java.util.stream.Collectors;
 
 /**
  * Represents a Theme component available in the system as a "persona" bundle.
@@ -87,7 +88,7 @@
         mTitle = title;
         mIsDefault = isDefault;
         mPreviewInfo = previewInfo;
-        mPackagesByCategory = Collections.unmodifiableMap(overlayPackages);
+        mPackagesByCategory = Collections.unmodifiableMap(removeNullValues(overlayPackages));
     }
 
     @Override
@@ -214,6 +215,13 @@
         }
     }
 
+    private Map<String, String> removeNullValues(Map<String, String> map) {
+        return map.entrySet()
+                .stream()
+                .filter(entry -> entry.getValue() != null)
+                .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue));
+    }
+
     protected CharSequence getContentDescription(Context context) {
         if (mContentDescription == null) {
             CharSequence defaultName = context.getString(R.string.default_theme_title);
diff --git a/src/com/android/customization/model/theme/custom/CustomThemeManager.java b/src/com/android/customization/model/theme/custom/CustomThemeManager.java
index 7b9b67c..42d73e6 100644
--- a/src/com/android/customization/model/theme/custom/CustomThemeManager.java
+++ b/src/com/android/customization/model/theme/custom/CustomThemeManager.java
@@ -16,20 +16,29 @@
 package com.android.customization.model.theme.custom;
 
 import android.content.Context;
+import android.os.Bundle;
+import android.text.TextUtils;
+import android.util.Log;
 
 import androidx.annotation.Nullable;
 
 import com.android.customization.model.CustomizationManager;
 import com.android.customization.model.theme.ThemeBundle.PreviewInfo;
+import com.android.customization.model.theme.ThemeBundleProvider;
 import com.android.customization.model.theme.ThemeManager;
 import com.android.customization.model.theme.custom.CustomTheme.Builder;
 
+import org.json.JSONException;
+
 import java.util.Map;
 
 public class CustomThemeManager implements CustomizationManager<ThemeComponentOption> {
 
+    private static final String TAG = "CustomThemeManager";
+    private static final String KEY_STATE_CURRENT_SELECTION = "CustomThemeManager.currentSelection";
+
     private final CustomTheme mOriginalTheme;
-    private final CustomTheme.Builder mBuilder;
+    private CustomTheme.Builder mBuilder;
 
     private CustomThemeManager(Map<String, String> overlayPackages,
             @Nullable CustomTheme originalTheme) {
@@ -72,6 +81,29 @@
         return mBuilder.createPreviewInfo(context);
     }
 
+    /** Saves the custom theme selections while system config changes. */
+    public void saveCustomTheme(Context context, Bundle savedInstanceState) {
+        CustomTheme customTheme =
+                buildPartialCustomTheme(context, /* id= */ null, /* title= */ null);
+        savedInstanceState.putString(KEY_STATE_CURRENT_SELECTION,
+                customTheme.getSerializedPackages());
+    }
+
+    /** Reads the saved custom theme after system config changed. */
+    public void readCustomTheme(ThemeBundleProvider themeBundleProvider,
+                                Bundle savedInstanceState) {
+        String packages = savedInstanceState.getString(KEY_STATE_CURRENT_SELECTION);
+        if (!TextUtils.isEmpty(packages)) {
+            try {
+                mBuilder = themeBundleProvider.parseCustomTheme(packages);
+            } catch (JSONException e) {
+                Log.w(TAG, "Couldn't parse provided custom theme.");
+            }
+        } else {
+            Log.w(TAG, "No custom theme being restored.");
+        }
+    }
+
     public static CustomThemeManager create(
             @Nullable CustomTheme customTheme, ThemeManager themeManager) {
         if (customTheme != null && customTheme.isDefined()) {
diff --git a/src/com/android/customization/picker/WallpaperPreviewer.java b/src/com/android/customization/picker/WallpaperPreviewer.java
index 0185925..8f69d59 100644
--- a/src/com/android/customization/picker/WallpaperPreviewer.java
+++ b/src/com/android/customization/picker/WallpaperPreviewer.java
@@ -174,7 +174,7 @@
 
         mWallpaperConnection.setVisibility(true);
         mHomePreview.post(() -> {
-            if (!mWallpaperConnection.connect()) {
+            if (mWallpaperConnection != null && !mWallpaperConnection.connect()) {
                 mWallpaperConnection = null;
                 LiveTileOverlay.INSTANCE.detach(mHomePreview.getOverlay());
             }
diff --git a/src/com/android/customization/picker/grid/GridFragment.java b/src/com/android/customization/picker/grid/GridFragment.java
index a8f6eda..05f535c 100644
--- a/src/com/android/customization/picker/grid/GridFragment.java
+++ b/src/com/android/customization/picker/grid/GridFragment.java
@@ -64,6 +64,9 @@
 public class GridFragment extends AppbarFragment {
 
     private static final int FULL_PREVIEW_REQUEST_CODE = 1000;
+    private static final String KEY_STATE_SELECTED_OPTION = "GridFragment.selectedOption";
+    private static final String KEY_STATE_BOTTOM_ACTION_BAR_VISIBILITY =
+            "GridFragment.bottomActionBarVisibility";
 
     private static final String TAG = "GridFragment";
 
@@ -98,33 +101,16 @@
     private final Callback mApplyGridCallback = new Callback() {
         @Override
         public void onSuccess() {
-            mGridManager.fetchOptions(new OptionsFetchedListener<GridOption>() {
-                @Override
-                public void onOptionsLoaded(List<GridOption> options) {
-                    mOptionsController.resetOptions(options);
-                    mSelectedOption = getSelectedOption(options);
-                    mReloadOptionsAfterApplying = true;
-                    // It will trigger OptionSelectedListener#onOptionSelected.
-                    mOptionsController.setSelectedOption(mSelectedOption);
-                    Toast.makeText(getContext(), R.string.applied_grid_msg, Toast.LENGTH_SHORT)
-                            .show();
-                    // Since we disabled it when clicked apply button.
-                    mBottomActionBar.enableActions();
-                    mBottomActionBar.hide();
-                }
-
-                @Override
-                public void onError(@Nullable Throwable throwable) {
-                    if (throwable != null) {
-                        Log.e(TAG, "Error loading grid options", throwable);
-                    }
-                    showError();
-                }
-            }, true);
+            Toast.makeText(getContext(), R.string.applied_grid_msg, Toast.LENGTH_SHORT).show();
+            getActivity().overridePendingTransition(R.anim.fade_in, R.anim.fade_out);
+            getActivity().finish();
         }
 
         @Override
         public void onError(@Nullable Throwable throwable) {
+            // Since we disabled it when clicked apply button.
+            mBottomActionBar.enableActions();
+            mBottomActionBar.hide();
             //TODO(chihhangchuang): handle
         }
     };
@@ -152,7 +138,7 @@
 
         // Clear memory cache whenever grid fragment view is being loaded.
         Glide.get(getContext()).clearMemory();
-        setUpOptions();
+        setUpOptions(savedInstanceState);
 
         ImageView wallpaperPreviewImage = view.findViewById(R.id.wallpaper_preview_image);
         wallpaperPreviewImage.setOnClickListener(v -> showFullPreview());
@@ -188,6 +174,18 @@
     }
 
     @Override
+    public void onSaveInstanceState(@NonNull Bundle outState) {
+        super.onSaveInstanceState(outState);
+        if (mSelectedOption != null) {
+            outState.putParcelable(KEY_STATE_SELECTED_OPTION, mSelectedOption);
+        }
+        if (mBottomActionBar != null) {
+            outState.putBoolean(KEY_STATE_BOTTOM_ACTION_BAR_VISIBILITY,
+                    mBottomActionBar.isVisible());
+        }
+    }
+
+    @Override
     public void onActivityResult(int requestCode, int resultCode, @Nullable Intent data) {
         super.onActivityResult(requestCode, resultCode, data);
         if (requestCode == FULL_PREVIEW_REQUEST_CODE && resultCode == RESULT_OK) {
@@ -212,12 +210,15 @@
         if (mGridOptionPreviewer != null) {
             mGridOptionPreviewer.release();
         }
+        if (getContext() == null) {
+            return;
+        }
         mGridOptionPreviewer = new GridOptionPreviewer(
                 getContext(), mGridManager, mGridPreviewContainer);
         mGridOptionPreviewer.setGridOption(mSelectedOption, mGridManager.usesSurfaceView());
     }
 
-    private void setUpOptions() {
+    private void setUpOptions(@Nullable Bundle savedInstanceState) {
         hideError();
         mLoading.show();
         mGridManager.fetchOptions(new OptionsFetchedListener<GridOption>() {
@@ -225,7 +226,6 @@
             public void onOptionsLoaded(List<GridOption> options) {
                 mLoading.hide();
                 mOptionsController = new OptionSelectorController<>(mOptionsContainer, options);
-
                 mOptionsController.addListener(selected -> {
                     mSelectedOption = (GridOption) selected;
                     if (mReloadOptionsAfterApplying) {
@@ -237,7 +237,23 @@
                     updatePreview();
                 });
                 mOptionsController.initOptions(mGridManager);
-                mSelectedOption = getSelectedOption(options);
+
+                GridOption previouslySelectedOption = null;
+                if (savedInstanceState != null) {
+                    previouslySelectedOption = findEquivalent(
+                            options, savedInstanceState.getParcelable(KEY_STATE_SELECTED_OPTION));
+                }
+                mSelectedOption = previouslySelectedOption != null
+                        ? previouslySelectedOption
+                        : getActiveOption(options);
+                mOptionsController.setSelectedOption(mSelectedOption);
+                boolean bottomActionBarVisibility = savedInstanceState != null
+                        && savedInstanceState.getBoolean(KEY_STATE_BOTTOM_ACTION_BAR_VISIBILITY);
+                if (bottomActionBarVisibility) {
+                    mBottomActionBar.show();
+                } else {
+                    mBottomActionBar.hide();
+                }
                 updatePreview();
             }
 
@@ -251,7 +267,7 @@
         }, false);
     }
 
-    private GridOption getSelectedOption(List<GridOption> options) {
+    private GridOption getActiveOption(List<GridOption> options) {
         return options.stream()
                 .filter(option -> option.isActive(mGridManager))
                 .findAny()
@@ -259,6 +275,14 @@
                 .orElse(options.get(0));
     }
 
+    @Nullable
+    private GridOption findEquivalent(List<GridOption> options, GridOption target) {
+        return options.stream()
+                .filter(option -> option.equals(target))
+                .findAny()
+                .orElse(null);
+    }
+
     private void hideError() {
         mContent.setVisibility(View.VISIBLE);
         mError.setVisibility(View.GONE);
diff --git a/src/com/android/customization/picker/theme/CustomThemeActivity.java b/src/com/android/customization/picker/theme/CustomThemeActivity.java
index a5ce222..fc56337 100644
--- a/src/com/android/customization/picker/theme/CustomThemeActivity.java
+++ b/src/com/android/customization/picker/theme/CustomThemeActivity.java
@@ -85,12 +85,12 @@
     protected void onCreate(Bundle savedInstanceState) {
         CustomizationInjector injector = (CustomizationInjector) InjectorProvider.getInjector();
         mUserEventLogger = (ThemesUserEventLogger) injector.getUserEventLogger(this);
+        ThemeBundleProvider themeProvider =
+                new DefaultThemeProvider(this, injector.getCustomizationPreferences(this));
         Intent intent = getIntent();
         CustomTheme customTheme = null;
         if (intent != null && intent.hasExtra(EXTRA_THEME_PACKAGES)
                 && intent.hasExtra(EXTRA_THEME_TITLE) && intent.hasExtra(EXTRA_THEME_ID)) {
-            ThemeBundleProvider themeProvider =
-                    new DefaultThemeProvider(this, injector.getCustomizationPreferences(this));
             mIsDefinedTheme = intent.getBooleanExtra(CREATE_NEW_THEME, true);
             try {
                 CustomTheme.Builder themeBuilder = themeProvider.parseCustomTheme(
@@ -112,6 +112,9 @@
                 mUserEventLogger);
         mThemeManager.fetchOptions(null, false);
         mCustomThemeManager = CustomThemeManager.create(customTheme, mThemeManager);
+        if (savedInstanceState != null) {
+            mCustomThemeManager.readCustomTheme(themeProvider, savedInstanceState);
+        }
 
         int currentStep = 0;
         if (savedInstanceState != null) {
@@ -138,6 +141,9 @@
     protected void onSaveInstanceState(Bundle outState) {
         super.onSaveInstanceState(outState);
         outState.putInt(KEY_STATE_CURRENT_STEP, mCurrentStep);
+        if (mCustomThemeManager != null) {
+            mCustomThemeManager.saveCustomTheme(this, outState);
+        }
     }
 
     private void navigateToStep(int i) {
@@ -355,8 +361,7 @@
             return CustomThemeComponentFragment.newInstance(
                     title,
                     position,
-                    titleResId,
-                    true);
+                    titleResId);
         }
     }
 
diff --git a/src/com/android/customization/picker/theme/CustomThemeComponentFragment.java b/src/com/android/customization/picker/theme/CustomThemeComponentFragment.java
index b2217aa..a6fdb1c 100644
--- a/src/com/android/customization/picker/theme/CustomThemeComponentFragment.java
+++ b/src/com/android/customization/picker/theme/CustomThemeComponentFragment.java
@@ -98,6 +98,9 @@
             mOptionsController.addListener(selected -> {
                 mSelectedOption = (ThemeComponentOption) selected;
                 bindPreview();
+                // Preview and apply. The selection will be kept whatever user goes to previous page
+                // or encounter system config changes, the current selection can be recovered.
+                mCustomThemeManager.apply(mSelectedOption, /* callback= */ null);
             });
             mOptionsController.initOptions(mCustomThemeManager);
 
diff --git a/src/com/android/customization/picker/theme/ThemeFragment.java b/src/com/android/customization/picker/theme/ThemeFragment.java
index 1f4917e..5c5dcdd 100644
--- a/src/com/android/customization/picker/theme/ThemeFragment.java
+++ b/src/com/android/customization/picker/theme/ThemeFragment.java
@@ -68,6 +68,8 @@
 
     private static final String TAG = "ThemeFragment";
     private static final String KEY_SELECTED_THEME = "ThemeFragment.SelectedThemeBundle";
+    private static final String KEY_STATE_BOTTOM_ACTION_BAR_VISIBILITY =
+            "ThemeFragment.bottomActionBarVisibility";
     private static final int FULL_PREVIEW_REQUEST_CODE = 1000;
 
     /**
@@ -177,6 +179,9 @@
     }
 
     private void updateThemePreviewColorPerWallpaper() {
+        if (getContext() == null) {
+            return;
+        }
         if (mCurrentHomeWallpaper != null && mWallpaperImage.getMeasuredWidth() > 0
                 && mWallpaperImage.getMeasuredHeight() > 0) {
             WallpaperColorsLoader.getWallpaperColors(
@@ -192,11 +197,9 @@
         mThemeManager.apply(mSelectedTheme, new Callback() {
             @Override
             public void onSuccess() {
-                // Since we disabled it when clicked apply button.
-                mBottomActionBar.enableActions();
-                mBottomActionBar.hide();
-                Toast.makeText(getContext(), R.string.applied_theme_msg,
-                        Toast.LENGTH_LONG).show();
+                Toast.makeText(getContext(), R.string.applied_theme_msg, Toast.LENGTH_LONG).show();
+                getActivity().overridePendingTransition(R.anim.fade_in, R.anim.fade_out);
+                getActivity().finish();
             }
 
             @Override
@@ -217,6 +220,10 @@
         if (mSelectedTheme != null && !mSelectedTheme.isActive(mThemeManager)) {
             outState.putString(KEY_SELECTED_THEME, mSelectedTheme.getSerializedPackages());
         }
+        if (mBottomActionBar != null) {
+            outState.putBoolean(KEY_STATE_BOTTOM_ACTION_BAR_VISIBILITY,
+                    mBottomActionBar.isVisible());
+        }
     }
 
     @Override
@@ -226,7 +233,8 @@
                 mSelectedTheme = null;
                 reloadOptions();
             } else if (resultCode == CustomThemeActivity.RESULT_THEME_APPLIED) {
-                reloadOptions();
+                getActivity().overridePendingTransition(R.anim.fade_in, R.anim.fade_out);
+                getActivity().finish();
             } else {
                 if (mSelectedTheme != null) {
                     mOptionsController.setSelectedOption(mSelectedTheme);
@@ -290,28 +298,36 @@
                     }
                 });
                 mOptionsController.initOptions(mThemeManager);
+
                 String previouslySelected = savedInstanceState != null
                         ? savedInstanceState.getString(KEY_SELECTED_THEME) : null;
+                ThemeBundle previouslySelectedTheme = null;
+                ThemeBundle activeTheme = null;
                 for (ThemeBundle theme : options) {
                     if (previouslySelected != null
                             && previouslySelected.equals(theme.getSerializedPackages())) {
-                        mSelectedTheme = theme;
-                    } else if (theme.isActive(mThemeManager)) {
-                        mSelectedTheme = theme;
-                        break;
+                        previouslySelectedTheme = theme;
+                    }
+                    if (theme.isActive(mThemeManager)) {
+                        activeTheme = theme;
                     }
                 }
+                mSelectedTheme = previouslySelectedTheme != null
+                        ? previouslySelectedTheme
+                        : activeTheme;
+
                 if (mSelectedTheme == null) {
                     // Select the default theme if there is no matching custom enabled theme
                     mSelectedTheme = findDefaultThemeBundle(options);
-                } else {
-                    // Only show show checkmark if we found a matching theme
-                    mOptionsController.setAppliedOption(mSelectedTheme);
                 }
                 mOptionsController.setSelectedOption(mSelectedTheme);
-                // Set selected option above will show BottomActionBar when entering the tab. But
-                // it should not show when entering the tab.
-                mBottomActionBar.hide();
+                boolean bottomActionBarVisibility = savedInstanceState != null
+                        && savedInstanceState.getBoolean(KEY_STATE_BOTTOM_ACTION_BAR_VISIBILITY);
+                if (bottomActionBarVisibility) {
+                    mBottomActionBar.show();
+                } else {
+                    mBottomActionBar.hide();
+                }
             }
             @Override
             public void onError(@Nullable Throwable throwable) {
@@ -335,9 +351,6 @@
             if (mSelectedTheme == null) {
                 // Select the default theme if there is no matching custom enabled theme
                 mSelectedTheme = findDefaultThemeBundle(options);
-            } else {
-                // Only show show checkmark if we found a matching theme
-                mOptionsController.setAppliedOption(mSelectedTheme);
             }
             mOptionsController.setSelectedOption(mSelectedTheme);
             // Set selected option above will show BottomActionBar,
diff --git a/src/com/android/customization/picker/theme/ThemeOptionPreviewer.java b/src/com/android/customization/picker/theme/ThemeOptionPreviewer.java
index afeb6ff..f24cf64 100644
--- a/src/com/android/customization/picker/theme/ThemeOptionPreviewer.java
+++ b/src/com/android/customization/picker/theme/ThemeOptionPreviewer.java
@@ -174,23 +174,41 @@
      * text) which will change its content color according to different wallpapers.
      */
     public void updateColorForLauncherWidgets(WallpaperColors colors) {
-        int color = mContext.getColor(
-                (colors.getColorHints() & WallpaperColors.HINT_SUPPORTS_DARK_TEXT) == 0
-                        ? R.color.text_color_light
-                        : R.color.text_color_dark);
+        boolean useLightTextColor =
+                (colors.getColorHints() & WallpaperColors.HINT_SUPPORTS_DARK_TEXT) == 0;
+        int textColor = mContext.getColor(useLightTextColor
+                ? R.color.text_color_light
+                : R.color.text_color_dark);
+        int textShadowColor = mContext.getColor(useLightTextColor
+                ? R.color.theme_preview_workspace_shadow_color_dark
+                : R.color.theme_preview_workspace_shadow_color_transparent);
         // Update the top status bar clock text color.
-        mStatusBarClock.setTextColor(color);
+        mStatusBarClock.setTextColor(textColor);
         // Update the top status bar icon color.
         ViewGroup iconsContainer = mContentView.findViewById(R.id.theme_preview_top_bar_icons);
         for (int i = 0; i < iconsContainer.getChildCount(); i++) {
             ((ImageView) iconsContainer.getChildAt(i))
-                    .setImageTintList(ColorStateList.valueOf(color));
+                    .setImageTintList(ColorStateList.valueOf(textColor));
         }
         // Update smart space date color.
-        ((TextView) mContentView.findViewById(R.id.smart_space_date)).setTextColor(color);
+        mSmartSpaceDate.setTextColor(textColor);
+        mSmartSpaceDate.setShadowLayer(
+                mContext.getResources().getDimension(
+                        R.dimen.preview_theme_smartspace_key_ambient_shadow_blur),
+                /* dx = */ 0,
+                /* dy = */ 0,
+                textShadowColor);
+
         // Update shape app icon name text color.
         for (int id : mShapeIconAppNameIds) {
-            ((TextView) mContentView.findViewById(id)).setTextColor(color);
+            TextView appName = mContentView.findViewById(id);
+            appName.setTextColor(textColor);
+            appName.setShadowLayer(
+                    mContext.getResources().getDimension(
+                            R.dimen.preview_theme_app_name_key_ambient_shadow_blur),
+                    /* dx = */ 0,
+                    /* dy = */ 0,
+                    textShadowColor);
         }
 
         mHasWallpaperColorSet = true;
@@ -229,23 +247,11 @@
     private void setHeadlineFont(Typeface headlineFont) {
         mStatusBarClock.setTypeface(headlineFont);
         mSmartSpaceDate.setTypeface(headlineFont);
-        mSmartSpaceDate.setShadowLayer(
-                mContext.getResources().getDimension(
-                        R.dimen.preview_theme_smartspace_key_ambient_shadow_blur),
-                /* dx = */ 0,
-                /* dy = */ 0,
-                mContext.getColor(R.color.theme_preview_workspace_shadow_color));
 
         // Update font of app names.
         for (int id : mShapeIconAppNameIds) {
             TextView appName = mContentView.findViewById(id);
             appName.setTypeface(headlineFont);
-            appName.setShadowLayer(
-                    mContext.getResources().getDimension(
-                            R.dimen.preview_theme_app_name_key_ambient_shadow_blur),
-                    /* dx = */ 0,
-                    /* dy = */ 0,
-                    mContext.getColor(R.color.theme_preview_workspace_shadow_color));
         }
 
         // Update font of color/icons section title.
diff --git a/src/com/android/customization/widget/OptionSelectorController.java b/src/com/android/customization/widget/OptionSelectorController.java
index 19dc34e..5b76b36 100644
--- a/src/com/android/customization/widget/OptionSelectorController.java
+++ b/src/com/android/customization/widget/OptionSelectorController.java
@@ -72,14 +72,14 @@
     private CustomizationOption mAppliedOption;
 
     public OptionSelectorController(RecyclerView container, List<T> options) {
-        this(container, options, false, true);
+        this(container, options, true, true);
     }
 
     public OptionSelectorController(RecyclerView container, List<T> options,
             boolean useGrid, boolean showCheckmark) {
         mContainer = container;
         mOptions = options;
-        mUseGrid = container.getResources().getBoolean(R.bool.use_grid_for_options) || useGrid;
+        mUseGrid = useGrid;
         mShowCheckmark = showCheckmark;
     }