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 f4466e0..42d73e6 100644
--- a/src/com/android/customization/model/theme/custom/CustomThemeManager.java
+++ b/src/com/android/customization/model/theme/custom/CustomThemeManager.java
@@ -35,7 +35,7 @@
 public class CustomThemeManager implements CustomizationManager<ThemeComponentOption> {
 
     private static final String TAG = "CustomThemeManager";
-    private static final String EXTRA_CUSTOM_THEME_OPTION = "custom_theme_option";
+    private static final String KEY_STATE_CURRENT_SELECTION = "CustomThemeManager.currentSelection";
 
     private final CustomTheme mOriginalTheme;
     private CustomTheme.Builder mBuilder;
@@ -85,14 +85,14 @@
     public void saveCustomTheme(Context context, Bundle savedInstanceState) {
         CustomTheme customTheme =
                 buildPartialCustomTheme(context, /* id= */ null, /* title= */ null);
-        savedInstanceState.putString(EXTRA_CUSTOM_THEME_OPTION,
+        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(EXTRA_CUSTOM_THEME_OPTION);
+        String packages = savedInstanceState.getString(KEY_STATE_CURRENT_SELECTION);
         if (!TextUtils.isEmpty(packages)) {
             try {
                 mBuilder = themeBundleProvider.parseCustomTheme(packages);
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/ThemeFragment.java b/src/com/android/customization/picker/theme/ThemeFragment.java
index 1c3b4c1..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);
@@ -313,10 +321,11 @@
                     mSelectedTheme = findDefaultThemeBundle(options);
                 }
                 mOptionsController.setSelectedOption(mSelectedTheme);
-                // Set selected option above will show BottomActionBar when entering the tab. But
-                // it should not show when entering the tab. But it's visible for previously
-                // selected theme.
-                if (mSelectedTheme != previouslySelectedTheme) {
+                boolean bottomActionBarVisibility = savedInstanceState != null
+                        && savedInstanceState.getBoolean(KEY_STATE_BOTTOM_ACTION_BAR_VISIBILITY);
+                if (bottomActionBarVisibility) {
+                    mBottomActionBar.show();
+                } else {
                     mBottomActionBar.hide();
                 }
             }