diff --git a/quickstep/recents_ui_overrides/src/com/android/quickstep/views/OverviewActionsView.java b/quickstep/recents_ui_overrides/src/com/android/quickstep/views/OverviewActionsView.java
index 6e9ca23..d160686 100644
--- a/quickstep/recents_ui_overrides/src/com/android/quickstep/views/OverviewActionsView.java
+++ b/quickstep/recents_ui_overrides/src/com/android/quickstep/views/OverviewActionsView.java
@@ -17,7 +17,7 @@
 package com.android.quickstep.views;
 
 import static com.android.launcher3.config.FeatureFlags.ENABLE_OVERVIEW_ACTIONS;
-import static com.android.quickstep.SysUINavigationMode.getMode;
+import static com.android.quickstep.SysUINavigationMode.removeShelfFromOverview;
 
 import android.content.Context;
 import android.util.AttributeSet;
@@ -31,7 +31,6 @@
 import com.android.launcher3.R;
 import com.android.launcher3.util.MultiValueAlpha;
 import com.android.launcher3.util.MultiValueAlpha.AlphaProperty;
-import com.android.quickstep.SysUINavigationMode;
 import com.android.quickstep.TaskOverlayFactory.OverlayUICallbacks;
 
 import java.lang.annotation.Retention;
@@ -119,8 +118,7 @@
     protected void onAttachedToWindow() {
         super.onAttachedToWindow();
         updateHiddenFlags(HIDDEN_DISABLED_FEATURE, !ENABLE_OVERVIEW_ACTIONS.get());
-        updateHiddenFlags(HIDDEN_UNSUPPORTED_NAVIGATION,
-                getMode(getContext()) == SysUINavigationMode.Mode.TWO_BUTTONS);
+        updateHiddenFlags(HIDDEN_UNSUPPORTED_NAVIGATION, !removeShelfFromOverview(getContext()));
     }
 
     public void updateHiddenFlags(@ActionsHiddenFlags int visibilityFlags, boolean enable) {
diff --git a/quickstep/recents_ui_overrides/src/com/android/quickstep/views/RecentsView.java b/quickstep/recents_ui_overrides/src/com/android/quickstep/views/RecentsView.java
index fc42acf..6041917 100644
--- a/quickstep/recents_ui_overrides/src/com/android/quickstep/views/RecentsView.java
+++ b/quickstep/recents_ui_overrides/src/com/android/quickstep/views/RecentsView.java
@@ -117,7 +117,6 @@
 import com.android.quickstep.RecentsAnimationTargets;
 import com.android.quickstep.RecentsModel;
 import com.android.quickstep.RecentsModel.TaskVisualsChangeListener;
-import com.android.quickstep.SysUINavigationMode;
 import com.android.quickstep.SystemUiProxy;
 import com.android.quickstep.TaskThumbnailCache;
 import com.android.quickstep.TaskUtils;
@@ -802,14 +801,11 @@
     @Override
     public void setInsets(Rect insets) {
         mInsets.set(insets);
-        resetPaddingFromTaskSize();
-    }
-
-    private void resetPaddingFromTaskSize() {
         DeviceProfile dp = mActivity.getDeviceProfile();
         getTaskSize(dp, mTempRect);
         mTaskWidth = mTempRect.width();
         mTaskHeight = mTempRect.height();
+
         mTempRect.top -= mTaskTopMargin;
         setPadding(mTempRect.left - mInsets.left, mTempRect.top - mInsets.top,
                 dp.widthPx - mInsets.right - mTempRect.right,
@@ -1575,7 +1571,6 @@
             mActivity.getDragLayer().recreateControllers();
             mActionsView.updateHiddenFlags(HIDDEN_NON_ZERO_ROTATION,
                     touchRotation != 0 || launcherRotation != 0);
-            resetPaddingFromTaskSize();
             requestLayout();
         }
     }
diff --git a/quickstep/recents_ui_overrides/src/com/android/quickstep/views/TaskThumbnailView.java b/quickstep/recents_ui_overrides/src/com/android/quickstep/views/TaskThumbnailView.java
index 8dc41d0..e525842 100644
--- a/quickstep/recents_ui_overrides/src/com/android/quickstep/views/TaskThumbnailView.java
+++ b/quickstep/recents_ui_overrides/src/com/android/quickstep/views/TaskThumbnailView.java
@@ -50,7 +50,6 @@
 import com.android.quickstep.TaskOverlayFactory;
 import com.android.quickstep.TaskOverlayFactory.TaskOverlay;
 import com.android.quickstep.views.TaskView.FullscreenDrawParams;
-import com.android.quickstep.util.RecentsOrientedState;
 import com.android.systemui.plugins.OverviewScreenshotActions;
 import com.android.systemui.plugins.PluginListener;
 import com.android.systemui.shared.recents.model.Task;
diff --git a/quickstep/src/com/android/quickstep/SysUINavigationMode.java b/quickstep/src/com/android/quickstep/SysUINavigationMode.java
index 0569828..c715c93 100644
--- a/quickstep/src/com/android/quickstep/SysUINavigationMode.java
+++ b/quickstep/src/com/android/quickstep/SysUINavigationMode.java
@@ -24,9 +24,7 @@
 import android.content.res.Resources;
 import android.util.Log;
 
-import com.android.launcher3.Launcher;
 import com.android.launcher3.util.MainThreadInitializedObject;
-import com.android.quickstep.views.RecentsView;
 
 import java.io.PrintWriter;
 import java.util.ArrayList;
@@ -130,14 +128,10 @@
         }
     }
 
-    public static boolean removeShelfFromOverview(Launcher launcher) {
-        // The shelf is core to the two-button mode model, so we need to continue supporting it
-        // when in portrait.
-        if (getMode(launcher) != Mode.TWO_BUTTONS) {
-            return true;
-        }
-        RecentsView recentsView = launcher.getOverviewPanel();
-        return !recentsView.getPagedOrientationHandler().isLayoutNaturalToLauncher();
+    /** @return Whether we can remove the shelf from overview. */
+    public static boolean removeShelfFromOverview(Context context) {
+        // The shelf is core to the two-button mode model, so we need to continue supporting it.
+        return getMode(context) != Mode.TWO_BUTTONS;
     }
 
     public void dump(PrintWriter pw) {
diff --git a/quickstep/src/com/android/quickstep/util/LayoutUtils.java b/quickstep/src/com/android/quickstep/util/LayoutUtils.java
index a2471d8..4edf2fb 100644
--- a/quickstep/src/com/android/quickstep/util/LayoutUtils.java
+++ b/quickstep/src/com/android/quickstep/util/LayoutUtils.java
@@ -16,8 +16,7 @@
 package com.android.quickstep.util;
 
 import static com.android.launcher3.config.FeatureFlags.ENABLE_OVERVIEW_ACTIONS;
-import static com.android.quickstep.SysUINavigationMode.Mode.TWO_BUTTONS;
-import static com.android.quickstep.SysUINavigationMode.getMode;
+import static com.android.quickstep.SysUINavigationMode.removeShelfFromOverview;
 
 import static java.lang.annotation.RetentionPolicy.SOURCE;
 
@@ -61,7 +60,7 @@
         } else {
             Resources res = context.getResources();
 
-            if (showOverviewActions(context)) {
+            if (ENABLE_OVERVIEW_ACTIONS.get() && removeShelfFromOverview(context)) {
                 //TODO: this needs to account for the swipe gesture height and accessibility
                 // UI when shown.
                 extraSpace = res.getDimensionPixelSize(R.dimen.overview_actions_height);
@@ -77,7 +76,7 @@
 
     public static void calculateFallbackTaskSize(Context context, DeviceProfile dp, Rect outRect) {
         float extraSpace;
-        if (showOverviewActions(context)) {
+        if (ENABLE_OVERVIEW_ACTIONS.get() && removeShelfFromOverview(context)) {
             extraSpace = context.getResources()
                     .getDimensionPixelSize(R.dimen.overview_actions_height);
         } else {
@@ -92,7 +91,7 @@
         float taskWidth, taskHeight, paddingHorz;
         Resources res = context.getResources();
         Rect insets = dp.getInsets();
-        final boolean showLargeTaskSize = showOverviewActions(context);
+        final boolean overviewActionsEnabled = ENABLE_OVERVIEW_ACTIONS.get();
 
         if (dp.isMultiWindowMode) {
             if (multiWindowStrategy == MULTI_WINDOW_STRATEGY_HALF_SCREEN) {
@@ -122,7 +121,7 @@
             final int paddingResId;
             if (dp.isVerticalBarLayout()) {
                 paddingResId = R.dimen.landscape_task_card_horz_space;
-            } else if (showLargeTaskSize) {
+            } else if (overviewActionsEnabled && removeShelfFromOverview(context)) {
                 paddingResId = R.dimen.portrait_task_card_horz_space_big_overview;
             } else {
                 paddingResId = R.dimen.portrait_task_card_horz_space;
@@ -131,7 +130,7 @@
         }
 
         float topIconMargin = res.getDimension(R.dimen.task_thumbnail_top_margin);
-        float paddingVert = showLargeTaskSize
+        float paddingVert = overviewActionsEnabled && removeShelfFromOverview(context)
                 ? 0 : res.getDimension(R.dimen.task_card_vert_space);
 
         // Note this should be same as dp.availableWidthPx and dp.availableHeightPx unless
@@ -157,7 +156,7 @@
 
     public static int getShelfTrackingDistance(Context context, DeviceProfile dp) {
         // Track the bottom of the window.
-        if (showOverviewActions(context)) {
+        if (ENABLE_OVERVIEW_ACTIONS.get() && removeShelfFromOverview(context)) {
             Rect taskSize = new Rect();
             calculateLauncherTaskSize(context, dp, taskSize);
             return (dp.heightPx - taskSize.height()) / 2;
@@ -181,8 +180,4 @@
             return srcHeight / targetHeight;
         }
     }
-
-    private static boolean showOverviewActions(Context context) {
-        return ENABLE_OVERVIEW_ACTIONS.get() && getMode(context) != TWO_BUTTONS;
-    }
 }
diff --git a/quickstep/src/com/android/quickstep/util/RecentsOrientedState.java b/quickstep/src/com/android/quickstep/util/RecentsOrientedState.java
index c6384d3..74daeca 100644
--- a/quickstep/src/com/android/quickstep/util/RecentsOrientedState.java
+++ b/quickstep/src/com/android/quickstep/util/RecentsOrientedState.java
@@ -25,7 +25,6 @@
 import static com.android.launcher3.states.RotationHelper.ALLOW_ROTATION_PREFERENCE_KEY;
 import static com.android.launcher3.util.Executors.UI_HELPER_EXECUTOR;
 
-import static com.android.quickstep.SysUINavigationMode.Mode.TWO_BUTTONS;
 import static java.lang.annotation.RetentionPolicy.SOURCE;
 
 import android.content.ContentResolver;
@@ -50,7 +49,6 @@
 import com.android.launcher3.Utilities;
 import com.android.launcher3.testing.TestProtocol;
 import com.android.launcher3.touch.PagedOrientationHandler;
-import com.android.quickstep.SysUINavigationMode;
 
 import java.lang.annotation.Retention;
 import java.util.function.IntConsumer;
@@ -156,9 +154,7 @@
         if (isFixedRotationTransformEnabled(context)) {
             mFlags |= FLAG_MULTIPLE_ORIENTATION_SUPPORTED_BY_FLAG;
         }
-        // TODO(b/154665738): Determine if we need animation for 2 button mode or not
-        if (mOrientationListener.canDetectOrientation()
-                && SysUINavigationMode.getMode(context) != TWO_BUTTONS) {
+        if (mOrientationListener.canDetectOrientation()) {
             mFlags |= FLAG_ROTATION_WATCHER_SUPPORTED;
         }
     }
diff --git a/quickstep/src/com/android/quickstep/views/ShelfScrimView.java b/quickstep/src/com/android/quickstep/views/ShelfScrimView.java
index af77c62..f5498c9 100644
--- a/quickstep/src/com/android/quickstep/views/ShelfScrimView.java
+++ b/quickstep/src/com/android/quickstep/views/ShelfScrimView.java
@@ -165,7 +165,7 @@
             if ((OVERVIEW.getVisibleElements(mLauncher) & ALL_APPS_HEADER_EXTRA) == 0) {
                 mDragHandleProgress = 1;
                 if (FeatureFlags.ENABLE_OVERVIEW_ACTIONS.get()
-                        && SysUINavigationMode.removeShelfFromOverview(mLauncher)) {
+                        && SysUINavigationMode.removeShelfFromOverview(context)) {
                     // Fade in all apps background quickly to distinguish from swiping from nav bar.
                     mMidAlpha = Themes.getAttrInteger(context, R.attr.allAppsInterimScrimAlpha);
                     mMidProgress = OverviewState.getDefaultVerticalProgress(mLauncher);
diff --git a/src/com/android/launcher3/InsettableFrameLayout.java b/src/com/android/launcher3/InsettableFrameLayout.java
index 9a66d32..faa18b8 100644
--- a/src/com/android/launcher3/InsettableFrameLayout.java
+++ b/src/com/android/launcher3/InsettableFrameLayout.java
@@ -91,9 +91,6 @@
     @Override
     public void onViewAdded(View child) {
         super.onViewAdded(child);
-        if (!isAttachedToWindow()) {
-            return;
-        }
         setFrameLayoutChildInsets(child, mInsets, new Rect());
     }
 
diff --git a/src/com/android/launcher3/Launcher.java b/src/com/android/launcher3/Launcher.java
index 49723d5..070a392 100644
--- a/src/com/android/launcher3/Launcher.java
+++ b/src/com/android/launcher3/Launcher.java
@@ -18,6 +18,7 @@
 
 import static android.content.pm.ActivityInfo.CONFIG_ORIENTATION;
 import static android.content.pm.ActivityInfo.CONFIG_SCREEN_SIZE;
+
 import static com.android.launcher3.AbstractFloatingView.TYPE_ALL;
 import static com.android.launcher3.AbstractFloatingView.TYPE_REBIND_SAFE;
 import static com.android.launcher3.AbstractFloatingView.TYPE_SNACKBAR;