diff --git a/src/com/android/customization/picker/CustomizationPickerActivity.java b/src/com/android/customization/picker/CustomizationPickerActivity.java
index 15ced54..c3b3057 100644
--- a/src/com/android/customization/picker/CustomizationPickerActivity.java
+++ b/src/com/android/customization/picker/CustomizationPickerActivity.java
@@ -65,6 +65,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 +74,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 +86,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 =
@@ -293,10 +296,9 @@
             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();
+        Fragment fragment = getSupportFragmentManager().findFragmentById(R.id.fragment_container);
+        if (fragment instanceof BottomActionBarFragment
+                && ((BottomActionBarFragment) fragment).onBackPressed()) {
             return;
         }
 
@@ -409,6 +411,11 @@
         finish();
     }
 
+    @Override
+    public BottomActionBar getBottomActionBar() {
+        return findViewById(R.id.bottom_actionbar);
+    }
+
     /**
      * 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..fb5af3f 100644
--- a/src/com/android/customization/picker/grid/GridFragment.java
+++ b/src/com/android/customization/picker/grid/GridFragment.java
@@ -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;
@@ -72,7 +72,7 @@
 /**
  * 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 +87,7 @@
 
     public static GridFragment newInstance(CharSequence title) {
         GridFragment fragment = new GridFragment();
-        fragment.setArguments(ToolbarFragment.createArguments(title));
+        fragment.setArguments(AppbarFragment.createArguments(title));
         return fragment;
     }
 
@@ -167,8 +167,12 @@
                 loadWallpaperBackground();
             }
         });
+        return view;
+    }
 
-        mBottomActionBar = getActivity().findViewById(R.id.bottom_actionbar);
+    @Override
+    protected void onBottomActionBarReady(BottomActionBar bottomActionBar) {
+        mBottomActionBar = bottomActionBar;
         mBottomActionBar.showActionsOnly(EnumSet.of(CANCEL, APPLY));
         mBottomActionBar.setActionClickListener(CANCEL, unused -> getActivity().onBackPressed());
         mBottomActionBar.setActionClickListener(APPLY, unused -> {
@@ -185,7 +189,6 @@
                 }
             });
         });
-        return view;
     }
 
     private void loadWallpaperBackground() {
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;
     }