diff --git a/res/drawable/ic_check_circle_filled_24px.xml b/res/drawable/ic_check_circle_filled_24px.xml
index cacf1d8..01d2091 100644
--- a/res/drawable/ic_check_circle_filled_24px.xml
+++ b/res/drawable/ic_check_circle_filled_24px.xml
@@ -19,7 +19,7 @@
     android:viewportWidth="24"
     android:viewportHeight="24">
   <path
-      android:fillColor="@color/selected_check_background_color"
+      android:fillColor="@color/accent_color"
       android:pathData="M12,2C6.5,2 2,6.5 2,12s4.5,10 10,10s10,-4.5 10,-10S17.5,2 12,2zM10,17l-4,-4l1.4,-1.4l2.6,2.6l6.6,-6.6L18,9L10,17z"/>
   <path
       android:pathData="m8.0085,14.9866 l-1.9939,-1.994 0.6892,-0.6889 0.6892,-0.6889 1.2925,1.2926c0.7109,0.711 1.3035,1.2926 1.3169,1.2926 0.0134,0 1.5034,-1.4789 3.3111,-3.2865l3.2866,-3.2865 0.689,0.689 0.689,0.689 -3.9878,3.9878 -3.9878,3.9878z"
diff --git a/res/layout-land/fragment_grid_picker.xml b/res/layout-land/fragment_grid_picker.xml
index 77330a8..6c2d8c5 100644
--- a/res/layout-land/fragment_grid_picker.xml
+++ b/res/layout-land/fragment_grid_picker.xml
@@ -39,34 +39,18 @@
             android:background="@color/secondary_color"
             app:card_style="screen_aspect_ratio"/>
 
-        <LinearLayout
+        <FrameLayout
             android:id="@+id/options_section"
             android:layout_width="0dp"
             android:layout_height="match_parent"
             android:layout_weight="1"
-            android:paddingVertical="10dp"
-            android:orientation="vertical">
+            android:paddingVertical="10dp">
 
             <androidx.recyclerview.widget.RecyclerView
                 android:id="@+id/options_container"
                 android:layout_width="match_parent"
-                android:layout_height="0dp"
-                android:layout_weight="1"/>
-
-            <RelativeLayout
-                android:layout_width="match_parent"
-                android:layout_height="wrap_content"
-                android:paddingHorizontal="10dp">
-                <Button
-                    android:id="@+id/apply_button"
-                    style="@style/ActionPrimaryButton"
-                    android:layout_width="wrap_content"
-                    android:layout_height="wrap_content"
-                    android:layout_alignParentEnd="true"
-                    android:text="@string/apply_btn"/>
-            </RelativeLayout>
-
-        </LinearLayout>
+                android:layout_height="match_parent"/>
+        </FrameLayout>
     </LinearLayout>
         <androidx.core.widget.ContentLoadingProgressBar
             android:id="@+id/loading_indicator"
diff --git a/res/layout/fragment_grid_picker.xml b/res/layout/fragment_grid_picker.xml
index ce571c9..5bf96e1 100644
--- a/res/layout/fragment_grid_picker.xml
+++ b/res/layout/fragment_grid_picker.xml
@@ -39,36 +39,39 @@
                 android:background="@color/preview_pager_background"
                 app:card_style="screen_aspect_ratio"
                 app:layout_constrainedHeight="true"
-                app:layout_constraintBottom_toTopOf="@id/options_container"
+                app:layout_constraintTop_toTopOf="parent"
+                app:layout_constraintBottom_toTopOf="@id/options_title"
+                app:layout_constraintStart_toStartOf="parent"
                 app:layout_constraintEnd_toEndOf="parent"
                 app:layout_constraintHeight_max="@dimen/preview_pager_max_height"
                 app:layout_constraintHeight_min="@dimen/preview_pager_min_height"
                 app:layout_constraintVertical_bias="0.0"
-                app:layout_constraintStart_toStartOf="parent"
-                app:layout_constraintTop_toTopOf="parent"
                 app:layout_constraintVertical_chainStyle="spread_inside"/>
 
+            <TextView
+                android:id="@+id/options_title"
+                android:layout_width="match_parent"
+                android:layout_height="wrap_content"
+                android:gravity="center"
+                android:paddingVertical="18dp"
+                android:lineHeight="24dp"
+                android:singleLine="true"
+                android:text="@string/grid_options_title"
+                android:textAppearance="@style/TitleTextAppearance"
+                app:layout_constraintTop_toBottomOf="@id/grid_preview_pager"
+                app:layout_constraintBottom_toTopOf="@id/options_container"
+                app:layout_constraintStart_toStartOf="parent"
+                app:layout_constraintEnd_toEndOf="parent"/>
+
             <androidx.recyclerview.widget.RecyclerView
                 android:id="@+id/options_container"
                 android:layout_width="match_parent"
                 android:layout_height="@dimen/options_container_height"
                 android:layout_gravity="center_horizontal"
-                android:layout_marginTop="10dp"
-                app:layout_constraintBottom_toTopOf="@id/placeholder"
+                app:layout_constraintTop_toBottomOf="@+id/options_title"
+                app:layout_constraintStart_toStartOf="parent"
                 app:layout_constraintEnd_toEndOf="parent"
-                app:layout_constraintStart_toStartOf="parent"
-                app:layout_constraintTop_toBottomOf="@+id/grid_preview_pager"
                 app:layout_constraintVertical_bias="1.0"/>
-
-            <Space
-                android:id="@+id/placeholder"
-                android:layout_width="match_parent"
-                android:layout_height="@dimen/min_taptarget_height"
-                app:layout_constraintBottom_toBottomOf="parent"
-                app:layout_constraintStart_toStartOf="parent"
-                app:layout_constraintTop_toBottomOf="@id/options_container"
-                app:layout_constraintVertical_bias="1.0"/>
-
         </androidx.constraintlayout.widget.ConstraintLayout>
 
         <androidx.core.widget.ContentLoadingProgressBar
diff --git a/res/layout/fragment_grid_scalable_picker.xml b/res/layout/fragment_grid_scalable_picker.xml
index de44923..6a3a77a 100644
--- a/res/layout/fragment_grid_scalable_picker.xml
+++ b/res/layout/fragment_grid_scalable_picker.xml
@@ -44,46 +44,29 @@
         app:behavior_peekHeight="@dimen/content_container_minimum_height"
         app:layout_behavior="@string/bottom_sheet_behavior">
 
-        <androidx.constraintlayout.widget.ConstraintLayout
+        <LinearLayout
             android:id="@+id/content_section"
             android:layout_width="match_parent"
-            android:layout_height="match_parent">
+            android:layout_height="match_parent"
+            android:orientation="vertical">
+
+            <TextView
+                android:id="@+id/options_title"
+                android:layout_width="match_parent"
+                android:layout_height="wrap_content"
+                android:gravity="center"
+                android:paddingVertical="18dp"
+                android:lineHeight="24dp"
+                android:singleLine="true"
+                android:text="@string/grid_options_title"
+                android:textAppearance="@style/TitleTextAppearance"/>
 
             <androidx.recyclerview.widget.RecyclerView
                 android:id="@+id/options_container"
                 android:layout_width="match_parent"
                 android:layout_height="@dimen/options_container_height"
-                android:layout_gravity="center_horizontal"
-                android:layout_marginTop="10dp"
-                app:layout_constraintBottom_toTopOf="@id/placeholder"
-                app:layout_constraintEnd_toEndOf="parent"
-                app:layout_constraintStart_toStartOf="parent"
-                app:layout_constraintTop_toTopOf="parent"
-                app:layout_constraintVertical_bias="1.0" />
-
-            <Space
-                android:id="@+id/placeholder"
-                android:layout_width="match_parent"
-                android:layout_height="@dimen/min_taptarget_height"
-                app:layout_constraintBottom_toTopOf="@id/apply_button"
-                app:layout_constraintStart_toStartOf="parent"
-                app:layout_constraintTop_toBottomOf="@id/options_container"
-                app:layout_constraintVertical_bias="1.0" />
-
-            <Button
-                android:id="@+id/apply_button"
-                style="@style/ActionPrimaryButton"
-                android:layout_width="wrap_content"
-                android:layout_height="wrap_content"
-                android:layout_gravity="end"
-                android:layout_marginVertical="10dp"
-                android:layout_marginEnd="10dp"
-                android:layout_weight="1"
-                android:text="@string/apply_theme_btn"
-                app:layout_constraintBottom_toBottomOf="parent"
-                app:layout_constraintEnd_toEndOf="parent" />
-
-        </androidx.constraintlayout.widget.ConstraintLayout>
+                android:layout_gravity="center_horizontal"/>
+        </LinearLayout>
     </FrameLayout>
 
     <FrameLayout
diff --git a/res/values/dimens.xml b/res/values/dimens.xml
index 49daee4..3c78967 100644
--- a/res/values/dimens.xml
+++ b/res/values/dimens.xml
@@ -21,15 +21,15 @@
     <dimen name="tip_dot_size">8dp</dimen>
     <dimen name="tip_dot_line_width">2dp</dimen>
 
-    <dimen name="check_size">20dp</dimen>
-    <dimen name="check_offset">-4dp</dimen>
+    <dimen name="check_size">16dp</dimen>
+    <dimen name="check_offset">4dp</dimen>
 
     <dimen name="component_preview_page_bottom_margin">16dp</dimen>
 
     <!-- Dimensions for the customization option tiles -->
-    <dimen name="options_container_height">104dp</dimen>
+    <dimen name="options_container_height">120dp</dimen>
     <dimen name="options_container_width">0dp</dimen>
-    <dimen name="option_tile_width">72dp</dimen>
+    <dimen name="option_tile_width">88dp</dimen>
     <dimen name="theme_option_icon_sample_height">18dp</dimen>
     <dimen name="theme_option_icon_sample_width">18dp</dimen>
     <dimen name="theme_option_shape_sample_height">16dp</dimen>
@@ -38,8 +38,8 @@
     <dimen name="theme_option_font_sample_width">52dp</dimen>
     <dimen name="theme_option_sample_padding">5dp</dimen>
     <!-- Note, using dp instead of sp as this is just the "+" symbol, not text -->
-    <dimen name="option_tile_padding_vertical">12dp</dimen>
-    <dimen name="option_tile_padding_horizontal">10dp</dimen>
+    <dimen name="option_tile_padding_vertical">16dp</dimen>
+    <dimen name="option_tile_padding_horizontal">16dp</dimen>
 
     <dimen name="option_bottom_margin">8dp</dimen>
     <!-- Note, using dp instead of sp as this text is more like a "snapshot" of the font -->
diff --git a/res/values/strings.xml b/res/values/strings.xml
index c940efd..c0aef7a 100755
--- a/res/values/strings.xml
+++ b/res/values/strings.xml
@@ -67,7 +67,7 @@
 
     <!-- Plus sign used to indicate that the user can add a custom theme -->
 
-    <!-- Name for the Android Theme that comes preset with the device [CHAR LIMIT=10]-->
+    <!-- Name for the Android Theme (Style/Wallpaper/Grid/Clock...) that comes preset with the device [CHAR LIMIT=10]-->
     <string name="default_theme_title">Default</string>
 
     <!-- Title of a page that shows the user the preview of a font selection [CHAR_LIMIT=20] -->
@@ -92,14 +92,20 @@
     <!-- Body text for previewing a font [CHAR LIMIT=160] -->
     <string name="font_card_body">Add your favorite fonts to every screen</string>
 
+    <!--Title for grid options [CHAR LIMIT=NONE] -->
+    <string name="grid_options_title">Choose a grid size</string>
+
     <!--Title for a grid option, describing the number of columns and rows, eg: 4x4 [CHAR LIMIT=10] -->
     <string name="grid_title_pattern"><xliff:g name="num_cols" example="1">%1$d</xliff:g>x<xliff:g name="num_rows" example="1">%2$d</xliff:g></string>
 
     <!-- Message shown when a theme has been applied successfully in the system [CHAR LIMIT=NONE] -->
-    <string name="applied_theme_msg">Style applied</string>
+    <string name="applied_theme_msg">Style set successfully</string>
 
     <!-- Message shown when a clock has been applied successfully in the system [CHAR LIMIT=NONE] -->
-    <string name="applied_clock_msg">Clock applied</string>
+    <string name="applied_clock_msg">Clock set successfully</string>
+
+    <!-- Message shown when a grid has been applied successfully in the system [CHAR LIMIT=NONE] -->
+    <string name="applied_grid_msg">Grid set successfully</string>
 
     <!-- Message shown when a theme couldn't be applied in the system because of an error
         [CHAR LIMIT=NONE] -->
diff --git a/src/com/android/customization/model/grid/GridOptionsManager.java b/src/com/android/customization/model/grid/GridOptionsManager.java
index b0e0787..84f5373 100644
--- a/src/com/android/customization/model/grid/GridOptionsManager.java
+++ b/src/com/android/customization/model/grid/GridOptionsManager.java
@@ -58,7 +58,7 @@
 
     @Override
     public void fetchOptions(OptionsFetchedListener<GridOption> callback, boolean reload) {
-        new FetchTask(mProvider, callback).execute();
+        new FetchTask(mProvider, callback, reload).execute();
     }
 
     /** See if using surface view to render grid options */
@@ -74,16 +74,18 @@
     private static class FetchTask extends AsyncTask<Void, Void, Pair<List<GridOption>, String>> {
         private final LauncherGridOptionsProvider mProvider;
         @Nullable private final OptionsFetchedListener<GridOption> mCallback;
+        private final boolean mReload;
 
         private FetchTask(@NonNull LauncherGridOptionsProvider provider,
-                @Nullable OptionsFetchedListener<GridOption> callback) {
+                @Nullable OptionsFetchedListener<GridOption> callback, boolean reload) {
             mCallback = callback;
             mProvider = provider;
+            mReload = reload;
         }
 
         @Override
         protected Pair<List<GridOption>, String> doInBackground(Void[] params) {
-            return mProvider.fetch(false);
+            return mProvider.fetch(mReload);
         }
 
         @Override
diff --git a/src/com/android/customization/model/grid/LauncherGridOptionsProvider.java b/src/com/android/customization/model/grid/LauncherGridOptionsProvider.java
index 70a8fc8..6a47f95 100644
--- a/src/com/android/customization/model/grid/LauncherGridOptionsProvider.java
+++ b/src/com/android/customization/model/grid/LauncherGridOptionsProvider.java
@@ -29,7 +29,6 @@
 import androidx.annotation.WorkerThread;
 
 import com.android.customization.model.ResourceConstants;
-import com.android.systemui.shared.system.SurfaceViewRequestUtils;
 import com.android.wallpaper.R;
 import com.android.wallpaper.util.PreviewUtils;
 
@@ -53,6 +52,9 @@
     private static final String COL_PREVIEW_COUNT = "preview_count";
     private static final String COL_IS_DEFAULT = "is_default";
 
+    // Normal gird size name
+    private static final String GRID_NAME_NORMAL = "normal";
+
     private static final String METADATA_KEY_PREVIEW_VERSION = "preview_version";
 
     private final Context mContext;
@@ -100,7 +102,9 @@
                 int cols = c.getInt(c.getColumnIndex(COL_COLS));
                 int previewCount = c.getInt(c.getColumnIndex(COL_PREVIEW_COUNT));
                 boolean isSet = Boolean.valueOf(c.getString(c.getColumnIndex(COL_IS_DEFAULT)));
-                String title = mContext.getString(R.string.grid_title_pattern, cols, rows);
+                String title = GRID_NAME_NORMAL.equals(name)
+                        ? mContext.getString(R.string.default_theme_title)
+                        : mContext.getString(R.string.grid_title_pattern, cols, rows);
                 mOptions.add(new GridOption(title, name, isSet, rows, cols,
                         mPreviewUtils.getUri(PREVIEW), previewCount, iconPath));
             }
@@ -116,7 +120,7 @@
      * Request rendering of home screen preview via Launcher to Wallpaper using SurfaceView
      * @param name      the grid option name
      * @param bundle    surface view request bundle generated from
-     *                  {@link SurfaceViewRequestUtils#createSurfaceBundle(SurfaceView)}.
+     *    {@link com.android.wallpaper.util.SurfaceViewUtils#createSurfaceViewRequest(SurfaceView)}.
      */
     Bundle renderPreview(String name, Bundle bundle) {
         bundle.putString("name", name);
diff --git a/src/com/android/customization/picker/CustomizationPickerActivity.java b/src/com/android/customization/picker/CustomizationPickerActivity.java
index 15ced54..f1e9a47 100644
--- a/src/com/android/customization/picker/CustomizationPickerActivity.java
+++ b/src/com/android/customization/picker/CustomizationPickerActivity.java
@@ -25,6 +25,7 @@
 import android.view.Gravity;
 import android.view.Menu;
 import android.view.MenuItem;
+import android.view.View;
 
 import androidx.annotation.IdRes;
 import androidx.annotation.NonNull;
@@ -65,6 +66,7 @@
 import com.android.wallpaper.module.Injector;
 import com.android.wallpaper.module.InjectorProvider;
 import com.android.wallpaper.module.UserEventLogger;
+import com.android.wallpaper.picker.BottomActionBarFragment;
 import com.android.wallpaper.picker.CategoryFragment;
 import com.android.wallpaper.picker.CategoryFragment.CategoryFragmentHost;
 import com.android.wallpaper.picker.MyPhotosStarter;
@@ -73,6 +75,7 @@
 import com.android.wallpaper.picker.WallpaperPickerDelegate;
 import com.android.wallpaper.picker.WallpapersUiContainer;
 import com.android.wallpaper.widget.BottomActionBar;
+import com.android.wallpaper.widget.BottomActionBar.BottomActionBarHost;
 
 import com.google.android.material.bottomnavigation.BottomNavigationView;
 
@@ -84,7 +87,8 @@
  *  Fragments providing customization options.
  */
 public class CustomizationPickerActivity extends FragmentActivity implements WallpapersUiContainer,
-        CategoryFragmentHost, ThemeFragmentHost, GridFragmentHost, ClockFragmentHost {
+        CategoryFragmentHost, ThemeFragmentHost, GridFragmentHost, ClockFragmentHost,
+        BottomActionBarHost {
 
     private static final String TAG = "CustomizationPickerActivity";
     @VisibleForTesting static final String WALLPAPER_FLAVOR_EXTRA =
@@ -98,6 +102,7 @@
 
     private static final Map<Integer, CustomizationSection> mSections = new HashMap<>();
     private CategoryFragment mWallpaperCategoryFragment;
+    private BottomActionBar mBottomActionBar;
 
     private boolean mWallpaperCategoryInitialized;
 
@@ -134,6 +139,15 @@
                                 ? R.id.nav_wallpaper : R.id.nav_theme);
             }
         }
+
+        mBottomActionBar = findViewById(R.id.bottom_actionbar);
+        mBottomActionBar.getViewTreeObserver().addOnGlobalLayoutListener(() -> {
+            // Only update the visibility of mBottomNav when mBottomActionBar visibility changes.
+            // Since the listener will be triggered by mBottomActionBar and its child views.
+            if (mBottomActionBar.getVisibility() == mBottomNav.getVisibility()) {
+                mBottomNav.setVisibility(mBottomActionBar.isVisible() ? View.GONE : View.VISIBLE);
+            }
+        });
     }
 
     @Override
@@ -288,15 +302,14 @@
 
     @Override
     public void onBackPressed() {
-        // For wallpaper tab, since it had child fragment.
-        if (mWallpaperCategoryFragment != null && mWallpaperCategoryFragment.popChildFragment()) {
+        Fragment fragment = getSupportFragmentManager().findFragmentById(R.id.fragment_container);
+        if (fragment instanceof BottomActionBarFragment
+                && ((BottomActionBarFragment) fragment).onBackPressed()) {
             return;
         }
 
-        // For other tabs without child fragment. Hide the BottomActionBar if back key is pressed.
-        BottomActionBar bottomActionBar = findViewById(R.id.bottom_actionbar);
-        if (bottomActionBar != null && bottomActionBar.isVisible()) {
-            bottomActionBar.hide();
+        // For wallpaper tab, since it had child fragment.
+        if (mWallpaperCategoryFragment != null && mWallpaperCategoryFragment.popChildFragment()) {
             return;
         }
 
@@ -409,6 +422,11 @@
         finish();
     }
 
+    @Override
+    public BottomActionBar getBottomActionBar() {
+        return mBottomActionBar;
+    }
+
     /**
      * Represents a section of the Picker (eg "ThemeBundle", "Clock", etc).
      * There should be a concrete subclass per available section, providing the corresponding
diff --git a/src/com/android/customization/picker/clock/ClockFragment.java b/src/com/android/customization/picker/clock/ClockFragment.java
index 48b4c0d..47fad5f 100644
--- a/src/com/android/customization/picker/clock/ClockFragment.java
+++ b/src/com/android/customization/picker/clock/ClockFragment.java
@@ -43,7 +43,7 @@
 import com.android.wallpaper.R;
 import com.android.wallpaper.asset.Asset;
 import com.android.wallpaper.module.InjectorProvider;
-import com.android.wallpaper.picker.ToolbarFragment;
+import com.android.wallpaper.picker.AppbarFragment;
 import com.android.wallpaper.widget.PreviewPager;
 
 import java.util.List;
@@ -51,7 +51,7 @@
 /**
  * Fragment that contains the main UI for selecting and applying a Clockface.
  */
-public class ClockFragment extends ToolbarFragment {
+public class ClockFragment extends AppbarFragment {
 
     private static final String TAG = "ClockFragment";
 
@@ -64,7 +64,7 @@
 
     public static ClockFragment newInstance(CharSequence title) {
         ClockFragment fragment = new ClockFragment();
-        fragment.setArguments(ToolbarFragment.createArguments(title));
+        fragment.setArguments(AppbarFragment.createArguments(title));
         return fragment;
     }
 
diff --git a/src/com/android/customization/picker/grid/GridFragment.java b/src/com/android/customization/picker/grid/GridFragment.java
index c9589ce..43eb82f 100644
--- a/src/com/android/customization/picker/grid/GridFragment.java
+++ b/src/com/android/customization/picker/grid/GridFragment.java
@@ -18,7 +18,6 @@
 import static android.content.res.Configuration.ORIENTATION_LANDSCAPE;
 
 import static com.android.wallpaper.widget.BottomActionBar.BottomAction.APPLY;
-import static com.android.wallpaper.widget.BottomActionBar.BottomAction.CANCEL;
 
 import android.app.Activity;
 import android.content.Context;
@@ -37,6 +36,7 @@
 import android.view.View.OnLayoutChangeListener;
 import android.view.ViewGroup;
 import android.widget.ImageView;
+import android.widget.Toast;
 
 import androidx.annotation.NonNull;
 import androidx.annotation.Nullable;
@@ -58,7 +58,7 @@
 import com.android.wallpaper.model.WallpaperInfo;
 import com.android.wallpaper.module.CurrentWallpaperInfoFactory;
 import com.android.wallpaper.module.InjectorProvider;
-import com.android.wallpaper.picker.ToolbarFragment;
+import com.android.wallpaper.picker.AppbarFragment;
 import com.android.wallpaper.util.SurfaceViewUtils;
 import com.android.wallpaper.widget.BottomActionBar;
 import com.android.wallpaper.widget.PreviewPager;
@@ -66,13 +66,12 @@
 import com.bumptech.glide.Glide;
 import com.bumptech.glide.request.RequestOptions;
 
-import java.util.EnumSet;
 import java.util.List;
 
 /**
  * Fragment that contains the UI for selecting and applying a GridOption.
  */
-public class GridFragment extends ToolbarFragment {
+public class GridFragment extends AppbarFragment {
 
     private static final int PREVIEW_FADE_DURATION_MS = 100;
 
@@ -87,7 +86,7 @@
 
     public static GridFragment newInstance(CharSequence title) {
         GridFragment fragment = new GridFragment();
-        fragment.setArguments(ToolbarFragment.createArguments(title));
+        fragment.setArguments(AppbarFragment.createArguments(title));
         return fragment;
     }
 
@@ -107,6 +106,41 @@
     private View mError;
     private BottomActionBar mBottomActionBar;
     private ThemesUserEventLogger mEventLogger;
+    private boolean mReloadOptionsAfterApplying;
+
+    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);
+        }
+
+        @Override
+        public void onError(@Nullable Throwable throwable) {
+            //TODO(chihhangchuang): handle
+        }
+    };
 
     @Override
     public void onAttach(Context context) {
@@ -167,25 +201,17 @@
                 loadWallpaperBackground();
             }
         });
+        return view;
+    }
 
-        mBottomActionBar = getActivity().findViewById(R.id.bottom_actionbar);
-        mBottomActionBar.showActionsOnly(EnumSet.of(CANCEL, APPLY));
-        mBottomActionBar.setActionClickListener(CANCEL, unused -> getActivity().onBackPressed());
+    @Override
+    protected void onBottomActionBarReady(BottomActionBar bottomActionBar) {
+        mBottomActionBar = bottomActionBar;
+        mBottomActionBar.showActionsOnly(APPLY);
         mBottomActionBar.setActionClickListener(APPLY, unused -> {
             mBottomActionBar.disableActions();
-            mGridManager.apply(mSelectedOption, new Callback() {
-                @Override
-                public void onSuccess() {
-                    getActivity().finish();
-                }
-
-                @Override
-                public void onError(@Nullable Throwable throwable) {
-                    //TODO(santie): handle
-                }
-            });
+            mGridManager.apply(mSelectedOption, mApplyGridCallback);
         });
-        return view;
     }
 
     private void loadWallpaperBackground() {
@@ -218,20 +244,16 @@
 
                 mOptionsController.addListener(selected -> {
                     mSelectedOption = (GridOption) selected;
+                    if (mReloadOptionsAfterApplying) {
+                        mReloadOptionsAfterApplying = false;
+                        return;
+                    }
                     mBottomActionBar.show();
                     mEventLogger.logGridSelected(mSelectedOption);
                     createAdapter();
                 });
                 mOptionsController.initOptions(mGridManager);
-                for (GridOption option : options) {
-                    if (option.isActive(mGridManager)) {
-                        mSelectedOption = option;
-                    }
-                }
-                // For development only, as there should always be a grid set.
-                if (mSelectedOption == null) {
-                    mSelectedOption = options.get(0);
-                }
+                mSelectedOption = getSelectedOption(options);
                 createAdapter();
             }
 
@@ -245,6 +267,14 @@
         }, false);
     }
 
+    private GridOption getSelectedOption(List<GridOption> options) {
+        return options.stream()
+                .filter(option -> option.isActive(mGridManager))
+                .findAny()
+                // For development only, as there should always be a grid set.
+                .orElse(options.get(0));
+    }
+
     private void hideError() {
         mContent.setVisibility(View.VISIBLE);
         mError.setVisibility(View.GONE);
diff --git a/src/com/android/customization/picker/theme/CustomThemeComponentFragment.java b/src/com/android/customization/picker/theme/CustomThemeComponentFragment.java
index f6471f0..b2217aa 100644
--- a/src/com/android/customization/picker/theme/CustomThemeComponentFragment.java
+++ b/src/com/android/customization/picker/theme/CustomThemeComponentFragment.java
@@ -15,27 +15,20 @@
  */
 package com.android.customization.picker.theme;
 
-import android.app.AlertDialog;
-import android.content.Context;
-import android.graphics.drawable.Drawable;
 import android.os.Bundle;
 import android.view.LayoutInflater;
-import android.view.MenuItem;
 import android.view.View;
 import android.view.ViewGroup;
-import android.widget.TextView;
 
 import androidx.annotation.NonNull;
 import androidx.annotation.Nullable;
-import androidx.annotation.StringRes;
 import androidx.recyclerview.widget.RecyclerView;
 
-import com.android.customization.model.theme.custom.CustomThemeManager;
 import com.android.customization.model.theme.custom.ThemeComponentOption;
 import com.android.customization.model.theme.custom.ThemeComponentOptionProvider;
 import com.android.customization.widget.OptionSelectorController;
 import com.android.wallpaper.R;
-import com.android.wallpaper.picker.ToolbarFragment;
+import com.android.wallpaper.picker.AppbarFragment;
 
 public class CustomThemeComponentFragment extends CustomThemeStepFragment {
     private static final String ARG_USE_GRID_LAYOUT = "CustomThemeComponentFragment.use_grid";;
@@ -48,7 +41,7 @@
     public static CustomThemeComponentFragment newInstance(CharSequence toolbarTitle, int position,
             int titleResId, boolean allowGridLayout) {
         CustomThemeComponentFragment fragment = new CustomThemeComponentFragment();
-        Bundle arguments = ToolbarFragment.createArguments(toolbarTitle);
+        Bundle arguments = AppbarFragment.createArguments(toolbarTitle);
         arguments.putInt(ARG_KEY_POSITION, position);
         arguments.putInt(ARG_KEY_TITLE_RES_ID, titleResId);
         arguments.putBoolean(ARG_USE_GRID_LAYOUT, allowGridLayout);
diff --git a/src/com/android/customization/picker/theme/CustomThemeNameFragment.java b/src/com/android/customization/picker/theme/CustomThemeNameFragment.java
index e21fc80..ed59eb1 100644
--- a/src/com/android/customization/picker/theme/CustomThemeNameFragment.java
+++ b/src/com/android/customization/picker/theme/CustomThemeNameFragment.java
@@ -42,14 +42,14 @@
 import com.android.wallpaper.asset.BitmapCachingAsset;
 import com.android.wallpaper.module.CurrentWallpaperInfoFactory;
 import com.android.wallpaper.module.InjectorProvider;
-import com.android.wallpaper.picker.ToolbarFragment;
+import com.android.wallpaper.picker.AppbarFragment;
 
 public class CustomThemeNameFragment extends CustomThemeStepFragment {
 
     public static CustomThemeNameFragment newInstance(CharSequence toolbarTitle, int position,
             int titleResId) {
         CustomThemeNameFragment fragment = new CustomThemeNameFragment();
-        Bundle arguments = ToolbarFragment.createArguments(toolbarTitle);
+        Bundle arguments = AppbarFragment.createArguments(toolbarTitle);
         arguments.putInt(ARG_KEY_POSITION, position);
         arguments.putInt(ARG_KEY_TITLE_RES_ID, titleResId);
         fragment.setArguments(arguments);
@@ -104,7 +104,6 @@
         mNameEditor = view.findViewById(R.id.custom_theme_name);
         mNameEditor.setText(mCustomThemeManager.getOriginalTheme().getTitle());
         bindCover(view.findViewById(R.id.component_preview_content));
-
         return view;
     }
 
diff --git a/src/com/android/customization/picker/theme/CustomThemeStepFragment.java b/src/com/android/customization/picker/theme/CustomThemeStepFragment.java
index ebf55e3..b05ebc4 100644
--- a/src/com/android/customization/picker/theme/CustomThemeStepFragment.java
+++ b/src/com/android/customization/picker/theme/CustomThemeStepFragment.java
@@ -18,9 +18,9 @@
 import com.android.customization.model.theme.custom.ThemeComponentOption;
 import com.android.customization.model.theme.custom.ThemeComponentOptionProvider;
 import com.android.wallpaper.R;
-import com.android.wallpaper.picker.ToolbarFragment;
+import com.android.wallpaper.picker.AppbarFragment;
 
-abstract class CustomThemeStepFragment extends ToolbarFragment {
+abstract class CustomThemeStepFragment extends AppbarFragment {
     protected static final String ARG_KEY_POSITION = "CustomThemeStepFragment.position";
     protected static final String ARG_KEY_TITLE_RES_ID = "CustomThemeStepFragment.title_res";
     protected CustomThemeComponentStepHost mHost;
diff --git a/src/com/android/customization/picker/theme/ThemeFragment.java b/src/com/android/customization/picker/theme/ThemeFragment.java
index c3f2a8d..b56c402 100644
--- a/src/com/android/customization/picker/theme/ThemeFragment.java
+++ b/src/com/android/customization/picker/theme/ThemeFragment.java
@@ -66,7 +66,7 @@
 import com.android.wallpaper.model.WallpaperInfo;
 import com.android.wallpaper.module.CurrentWallpaperInfoFactory;
 import com.android.wallpaper.module.InjectorProvider;
-import com.android.wallpaper.picker.ToolbarFragment;
+import com.android.wallpaper.picker.AppbarFragment;
 import com.android.wallpaper.widget.PreviewPager;
 
 import java.util.List;
@@ -74,7 +74,7 @@
 /**
  * Fragment that contains the main UI for selecting and applying a ThemeBundle.
  */
-public class ThemeFragment extends ToolbarFragment {
+public class ThemeFragment extends AppbarFragment {
 
     private static final String TAG = "ThemeFragment";
     private static final String KEY_SELECTED_THEME = "ThemeFragment.SelectedThemeBundle";
@@ -87,7 +87,7 @@
     }
     public static ThemeFragment newInstance(CharSequence title) {
         ThemeFragment fragment = new ThemeFragment();
-        fragment.setArguments(ToolbarFragment.createArguments(title));
+        fragment.setArguments(AppbarFragment.createArguments(title));
         return fragment;
     }