diff --git a/quickstep/src/com/android/quickstep/util/SplitAnimationController.kt b/quickstep/src/com/android/quickstep/util/SplitAnimationController.kt
index 96a5733..8601350 100644
--- a/quickstep/src/com/android/quickstep/util/SplitAnimationController.kt
+++ b/quickstep/src/com/android/quickstep/util/SplitAnimationController.kt
@@ -65,7 +65,6 @@
 import com.android.launcher3.util.SplitConfigurationOptions.SplitSelectSource
 import com.android.launcher3.views.BaseDragLayer
 import com.android.quickstep.TaskViewUtils
-import com.android.quickstep.util.SplitScreenUtils.Companion.extractTopParentAndChildren
 import com.android.quickstep.views.FloatingAppPairView
 import com.android.quickstep.views.FloatingTaskView
 import com.android.quickstep.views.GroupedTaskView
@@ -982,14 +981,35 @@
         progressUpdater.setDuration(QuickstepTransitionManager.APP_LAUNCH_DURATION)
         progressUpdater.interpolator = Interpolators.EMPHASIZED
 
-        val splitTree: Pair<Change, List<Change>>? = extractTopParentAndChildren(transitionInfo)
-        check(splitTree != null) { "Could not find a split root candidate" }
-        val rootCandidate = splitTree.first
-        val stageRootTaskIds: Set<Int> = splitTree.second.map { it.taskInfo!!.taskId }.toSet()
-        val leafTasks: List<Change> =
-            transitionInfo.changes
-                .filter { it.taskInfo != null && it.taskInfo!!.parentTaskId in stageRootTaskIds }
-                .toList()
+        var rootCandidate: Change? = null
+
+        for (change in transitionInfo.changes) {
+            val taskInfo: RunningTaskInfo = change.taskInfo ?: continue
+
+            // TODO (b/316490565): Replace this logic when SplitBounds is available to
+            //  startAnimation() and we can know the precise taskIds of launching tasks.
+            if (
+                taskInfo.windowingMode == windowingMode &&
+                    (change.mode == TRANSIT_OPEN || change.mode == TRANSIT_TO_FRONT)
+            ) {
+                // Found one!
+                rootCandidate = change
+                break
+            }
+        }
+
+        // If we could not find a proper root candidate, something went wrong.
+        check(rootCandidate != null) { "Could not find a split root candidate" }
+
+        // Recurse up the tree until parent is null, then we've found our root.
+        var parentToken: WindowContainerToken? = rootCandidate.parent
+        while (parentToken != null) {
+            rootCandidate = transitionInfo.getChange(parentToken) ?: break
+            parentToken = rootCandidate.parent
+        }
+
+        // Make sure nothing weird happened, like getChange() returning null.
+        check(rootCandidate != null) { "Failed to find a root leash" }
 
         // Starting position is a 34% size tile centered in the middle of the screen.
         // Ending position is the full device screen.
@@ -1023,42 +1043,6 @@
                 override fun onAnimationEnd(animation: Animator) {
                     finishCallback.run()
                 }
-
-                override fun onAnimationStart(animation: Animator) {
-                    // Reset leaf and stage root tasks, animation can begin from freeform windows
-                    for (leaf in leafTasks) {
-                        val endAbsBounds = leaf.endAbsBounds
-
-                        t.setAlpha(leaf.leash, 1f)
-                        t.setCrop(
-                            leaf.leash,
-                            0f,
-                            0f,
-                            endAbsBounds.width().toFloat(),
-                            endAbsBounds.height().toFloat(),
-                        )
-                        t.setPosition(leaf.leash, 0f, 0f)
-                    }
-
-                    for (stageRoot in splitTree.second) {
-                        val endAbsBounds = stageRoot.endAbsBounds
-
-                        t.setAlpha(stageRoot.leash, 1f)
-                        t.setCrop(
-                            stageRoot.leash,
-                            0f,
-                            0f,
-                            endAbsBounds.width().toFloat(),
-                            endAbsBounds.height().toFloat(),
-                        )
-                        t.setPosition(
-                            stageRoot.leash,
-                            endAbsBounds.left.toFloat(),
-                            endAbsBounds.top.toFloat(),
-                        )
-                    }
-                    t.apply()
-                }
             }
         )
 
diff --git a/quickstep/src/com/android/quickstep/util/SplitScreenUtils.kt b/quickstep/src/com/android/quickstep/util/SplitScreenUtils.kt
index 7787e30..4005c5a 100644
--- a/quickstep/src/com/android/quickstep/util/SplitScreenUtils.kt
+++ b/quickstep/src/com/android/quickstep/util/SplitScreenUtils.kt
@@ -46,8 +46,8 @@
          * Given a TransitionInfo, generates the tree structure for those changes and extracts out
          * the top most root and it's two immediate children. Changes can be provided in any order.
          *
-         * @return null if no root is found, otherwise a [Pair] where first -> top most split root,
-         *         second -> [List] of 2, leftTop/bottomRight stage roots
+         * @return a [Pair] where first -> top most split root, second -> [List] of 2,
+         *   leftTop/bottomRight stage roots
          */
         fun extractTopParentAndChildren(
             transitionInfo: TransitionInfo