diff --git a/src/com/android/customization/model/grid/GridOptionViewModel.java b/src/com/android/customization/model/grid/GridOptionViewModel.java
new file mode 100644
index 0000000..33fa8e179
--- /dev/null
+++ b/src/com/android/customization/model/grid/GridOptionViewModel.java
@@ -0,0 +1,53 @@
+/*
+ * Copyright (C) 2021 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.
+ */
+package com.android.customization.model.grid;
+
+import androidx.lifecycle.SavedStateHandle;
+import androidx.lifecycle.ViewModel;
+
+/** The class to store status of the grid fragment view. */
+public class GridOptionViewModel extends ViewModel {
+    private static final String SELECTED_OPTION_KEY = "selected_option";
+    private static final String BOTTOM_ACTION_BAR_VISIBLE_KEY = "bottom_action_bar_visible";
+
+    private SavedStateHandle mState;
+
+    public GridOptionViewModel(SavedStateHandle savedStateHandle) {
+        mState = savedStateHandle;
+    }
+
+    /** Gets selected {@link GridOption} from {@link SavedStateHandle} */
+    public GridOption getSelectedOption() {
+        return mState.get(SELECTED_OPTION_KEY);
+    }
+
+    /** Sets selected {@link GridOption} to {@link SavedStateHandle} */
+    public void setSelectedOption(GridOption selectedOption) {
+        mState.set(SELECTED_OPTION_KEY, selectedOption);
+    }
+
+    /** Gets bottom action bar visible from {@link SavedStateHandle} */
+    public boolean getBottomActionBarVisible() {
+        return mState.contains(BOTTOM_ACTION_BAR_VISIBLE_KEY)
+                ? mState.get(BOTTOM_ACTION_BAR_VISIBLE_KEY)
+                : false;
+    }
+
+    /** Sets bottom action bar visible to {@link SavedStateHandle} */
+    public void setBottomActionBarVisible(boolean bottomActionBarVisible) {
+        mState.set(BOTTOM_ACTION_BAR_VISIBLE_KEY, bottomActionBarVisible);
+    }
+}
diff --git a/src/com/android/customization/picker/grid/GridFragment.java b/src/com/android/customization/picker/grid/GridFragment.java
index d554f72..e09bcbe 100644
--- a/src/com/android/customization/picker/grid/GridFragment.java
+++ b/src/com/android/customization/picker/grid/GridFragment.java
@@ -28,12 +28,14 @@
 import androidx.annotation.NonNull;
 import androidx.annotation.Nullable;
 import androidx.core.widget.ContentLoadingProgressBar;
+import androidx.lifecycle.ViewModelProvider;
 import androidx.recyclerview.widget.RecyclerView;
 
 import com.android.customization.model.CustomizationManager.Callback;
 import com.android.customization.model.CustomizationManager.OptionsFetchedListener;
 import com.android.customization.model.CustomizationOption;
 import com.android.customization.model.grid.GridOption;
+import com.android.customization.model.grid.GridOptionViewModel;
 import com.android.customization.model.grid.GridOptionsManager;
 import com.android.customization.module.ThemesUserEventLogger;
 import com.android.customization.picker.WallpaperPreviewer;
@@ -57,21 +59,18 @@
 public class GridFragment extends AppbarFragment {
 
     private static final String TAG = "GridFragment";
-    private static final String KEY_STATE_SELECTED_OPTION = "GridFragment.selectedOption";
-    private static final String KEY_STATE_BOTTOM_ACTION_BAR_VISIBLE =
-            "GridFragment.bottomActionBarVisible";
 
     private WallpaperInfo mHomeWallpaper;
     private RecyclerView mOptionsContainer;
     private OptionSelectorController<GridOption> mOptionsController;
     private GridOptionsManager mGridManager;
-    private GridOption mSelectedOption;
     private ContentLoadingProgressBar mLoading;
     private View mContent;
     private View mError;
     private BottomActionBar mBottomActionBar;
     private ThemesUserEventLogger mEventLogger;
     private GridOptionPreviewer mGridOptionPreviewer;
+    private GridOptionViewModel mGridOptionViewModel;
 
     private final Callback mApplyGridCallback = new Callback() {
         @Override
@@ -89,10 +88,18 @@
             // Since we disabled it when clicked apply button.
             mBottomActionBar.enableActions();
             mBottomActionBar.hide();
+            mGridOptionViewModel.setBottomActionBarVisible(false);
             //TODO(chihhangchuang): handle
         }
     };
 
+    @Override
+    public void onCreate(Bundle savedInstanceState) {
+        super.onCreate(savedInstanceState);
+        mGridOptionViewModel = new ViewModelProvider(requireActivity()).get(
+                GridOptionViewModel.class);
+    }
+
     @Nullable
     @Override
     public View onCreateView(@NonNull LayoutInflater inflater, @Nullable ViewGroup container,
@@ -121,7 +128,7 @@
         mGridManager = GridOptionsManager.getInstance(getContext());
         mEventLogger = (ThemesUserEventLogger) InjectorProvider.getInjector()
                 .getUserEventLogger(getContext());
-        setUpOptions(savedInstanceState);
+        setUpOptions();
 
         SurfaceView wallpaperSurface = view.findViewById(R.id.wallpaper_preview_surface);
         WallpaperPreviewer wallpaperPreviewer = new WallpaperPreviewer(getLifecycle(),
@@ -141,6 +148,13 @@
     }
 
     @Override
+    public boolean onBackPressed() {
+        mGridOptionViewModel.setSelectedOption(null);
+        mGridOptionViewModel.setBottomActionBarVisible(false);
+        return super.onBackPressed();
+    }
+
+    @Override
     public void onDestroy() {
         super.onDestroy();
         if (mGridOptionPreviewer != null) {
@@ -149,17 +163,6 @@
     }
 
     @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_VISIBLE, mBottomActionBar.isVisible());
-        }
-    }
-
-    @Override
     public CharSequence getDefaultTitle() {
         return getString(R.string.grid_title);
     }
@@ -169,7 +172,8 @@
         super.onBottomActionBarReady(bottomActionBar);
         mBottomActionBar = bottomActionBar;
         mBottomActionBar.showActionsOnly(APPLY_TEXT);
-        mBottomActionBar.setActionClickListener(APPLY_TEXT, v -> applyGridOption(mSelectedOption));
+        mBottomActionBar.setActionClickListener(APPLY_TEXT,
+                v -> applyGridOption(mGridOptionViewModel.getSelectedOption()));
     }
 
     private void applyGridOption(GridOption gridOption) {
@@ -177,7 +181,7 @@
         mGridManager.apply(gridOption, mApplyGridCallback);
     }
 
-    private void setUpOptions(@Nullable Bundle savedInstanceState) {
+    private void setUpOptions() {
         hideError();
         mLoading.show();
         mGridManager.fetchOptions(new OptionsFetchedListener<GridOption>() {
@@ -188,24 +192,21 @@
                         mOptionsContainer, options, /* useGrid= */ false,
                         CheckmarkStyle.CENTER_CHANGE_COLOR_WHEN_NOT_SELECTED);
                 mOptionsController.initOptions(mGridManager);
+                GridOption previouslySelectedOption = findEquivalent(options,
+                        mGridOptionViewModel.getSelectedOption());
+                mGridOptionViewModel.setSelectedOption(
+                        previouslySelectedOption != null
+                                ? previouslySelectedOption
+                                : getActiveOption(options));
 
-                // Find the selected Grid option.
-                GridOption previouslySelectedOption = null;
-                if (savedInstanceState != null) {
-                    previouslySelectedOption = findEquivalent(
-                            options, savedInstanceState.getParcelable(KEY_STATE_SELECTED_OPTION));
-                }
-                mSelectedOption = previouslySelectedOption != null
-                        ? previouslySelectedOption
-                        : getActiveOption(options);
-
-                mOptionsController.setSelectedOption(mSelectedOption);
-                onOptionSelected(mSelectedOption);
-                restoreBottomActionBarVisibility(savedInstanceState);
+                mOptionsController.setSelectedOption(mGridOptionViewModel.getSelectedOption());
+                onOptionSelected(mGridOptionViewModel.getSelectedOption());
+                restoreBottomActionBarVisibility();
 
                 mOptionsController.addListener(selectedOption -> {
                     onOptionSelected(selectedOption);
                     mBottomActionBar.show();
+                    mGridOptionViewModel.setBottomActionBarVisible(true);
                 });
             }
 
@@ -247,15 +248,13 @@
     }
 
     private void onOptionSelected(CustomizationOption selectedOption) {
-        mSelectedOption = (GridOption) selectedOption;
-        mEventLogger.logGridSelected(mSelectedOption);
-        mGridOptionPreviewer.setGridOption(mSelectedOption);
+        mGridOptionViewModel.setSelectedOption((GridOption) selectedOption);
+        mEventLogger.logGridSelected(mGridOptionViewModel.getSelectedOption());
+        mGridOptionPreviewer.setGridOption(mGridOptionViewModel.getSelectedOption());
     }
 
-    private void restoreBottomActionBarVisibility(@Nullable Bundle savedInstanceState) {
-        boolean isBottomActionBarVisible = savedInstanceState != null
-                && savedInstanceState.getBoolean(KEY_STATE_BOTTOM_ACTION_BAR_VISIBLE);
-        if (isBottomActionBarVisible) {
+    private void restoreBottomActionBarVisibility() {
+        if (mGridOptionViewModel.getBottomActionBarVisible()) {
             mBottomActionBar.show();
         } else {
             mBottomActionBar.hide();