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/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/grid/GridFragment.java b/src/com/android/customization/picker/grid/GridFragment.java
index e25e57c..b861e6a 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";
 
@@ -102,7 +105,8 @@
                 @Override
                 public void onOptionsLoaded(List<GridOption> options) {
                     mOptionsController.resetOptions(options);
-                    mSelectedOption = getSelectedOption(options);
+                    mSelectedOption = getActiveOption(options);
+                    mOptionsController.setAppliedOption(mSelectedOption);
                     mReloadOptionsAfterApplying = true;
                     // It will trigger OptionSelectedListener#onOptionSelected.
                     mOptionsController.setSelectedOption(mSelectedOption);
@@ -152,7 +156,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 +192,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) {
@@ -220,7 +236,7 @@
         mGridOptionPreviewer.setGridOption(mSelectedOption, mGridManager.usesSurfaceView());
     }
 
-    private void setUpOptions() {
+    private void setUpOptions(@Nullable Bundle savedInstanceState) {
         hideError();
         mLoading.show();
         mGridManager.fetchOptions(new OptionsFetchedListener<GridOption>() {
@@ -228,7 +244,6 @@
             public void onOptionsLoaded(List<GridOption> options) {
                 mLoading.hide();
                 mOptionsController = new OptionSelectorController<>(mOptionsContainer, options);
-
                 mOptionsController.addListener(selected -> {
                     mSelectedOption = (GridOption) selected;
                     if (mReloadOptionsAfterApplying) {
@@ -240,7 +255,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();
             }
 
@@ -254,7 +285,7 @@
         }, false);
     }
 
-    private GridOption getSelectedOption(List<GridOption> options) {
+    private GridOption getActiveOption(List<GridOption> options) {
         return options.stream()
                 .filter(option -> option.isActive(mGridManager))
                 .findAny()
@@ -262,6 +293,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 61f1949..64147b2 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;
 
     /**
@@ -220,6 +222,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
@@ -316,10 +322,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();
                 }
             }