diff --git a/quickstep/src/com/android/quickstep/util/SplitAnimationController.kt b/quickstep/src/com/android/quickstep/util/SplitAnimationController.kt
index 8601350b8f..7dc7b9ec31 100644
--- a/quickstep/src/com/android/quickstep/util/SplitAnimationController.kt
+++ b/quickstep/src/com/android/quickstep/util/SplitAnimationController.kt
@@ -36,6 +36,7 @@ import android.view.RemoteAnimationTarget
 import android.view.SurfaceControl
 import android.view.SurfaceControl.Transaction
 import android.view.View
+import android.view.WindowManager.TRANSIT_CHANGE
 import android.view.WindowManager.TRANSIT_OPEN
 import android.view.WindowManager.TRANSIT_TO_FRONT
 import android.window.TransitionInfo
@@ -65,6 +66,7 @@ import com.android.launcher3.util.MultiPropertyFactory.MULTI_PROPERTY_VALUE
 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
@@ -730,7 +732,7 @@ class SplitAnimationController(val splitSelectStateController: SplitSelectStateC
         val launchAnimation = AnimatorSet()
 
         val splitRoots: Pair<Change, List<Change>>? =
-            SplitScreenUtils.extractTopParentAndChildren(transitionInfo)
+            extractTopParentAndChildren(transitionInfo)
         check(splitRoots != null) { "Could not find split roots" }
 
         // Will point to change (0) in diagram above
@@ -981,35 +983,19 @@ class SplitAnimationController(val splitSelectStateController: SplitSelectStateC
         progressUpdater.setDuration(QuickstepTransitionManager.APP_LAUNCH_DURATION)
         progressUpdater.interpolator = Interpolators.EMPHASIZED
 
-        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
+        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 {
+                (TransitionUtil.isOpeningMode(it.mode) || it.mode == TRANSIT_CHANGE)
+                        && it.taskInfo != null
+                        && it.taskInfo!!.parentTaskId in stageRootTaskIds
             }
-        }
-
-        // 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" }
+            .toList()
 
         // Starting position is a 34% size tile centered in the middle of the screen.
         // Ending position is the full device screen.
@@ -1037,13 +1023,38 @@ class SplitAnimationController(val splitSelectStateController: SplitSelectStateC
             t.apply()
         }
 
+
         // When animation ends,  run finishCallback
         progressUpdater.addListener(
             object : AnimatorListenerAdapter() {
                 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()
+                }
             }
+
         )
 
         launchAnimation.play(progressUpdater)
diff --git a/quickstep/src/com/android/quickstep/util/SplitScreenUtils.kt b/quickstep/src/com/android/quickstep/util/SplitScreenUtils.kt
index 4005c5a813..1919f68d4a 100644
--- a/quickstep/src/com/android/quickstep/util/SplitScreenUtils.kt
+++ b/quickstep/src/com/android/quickstep/util/SplitScreenUtils.kt
@@ -17,14 +17,13 @@
 package com.android.quickstep.util
 
 import android.util.Log
-import android.view.WindowManager.TRANSIT_OPEN
-import android.view.WindowManager.TRANSIT_TO_FRONT
+import android.view.WindowManager.TRANSIT_CHANGE
 import android.window.TransitionInfo
 import android.window.TransitionInfo.Change
 import android.window.TransitionInfo.FLAG_FIRST_CUSTOM
 import com.android.launcher3.util.SplitConfigurationOptions
+import com.android.wm.shell.shared.TransitionUtil
 import com.android.wm.shell.shared.split.SplitBounds
-import java.lang.IllegalStateException
 
 class SplitScreenUtils {
     companion object {
@@ -54,14 +53,7 @@ class SplitScreenUtils {
         ): Pair<Change, List<Change>>? {
             val parentToChildren = mutableMapOf<Change, MutableList<Change>>()
             val hasParent = mutableSetOf<Change>()
-            // filter out anything that isn't opening and the divider
-            val taskChanges: List<Change> =
-                transitionInfo.changes
-                    .filter { change ->
-                        (change.mode == TRANSIT_OPEN || change.mode == TRANSIT_TO_FRONT) &&
-                            change.flags < FLAG_FIRST_CUSTOM
-                    }
-                    .toList()
+            val taskChanges: List<Change> = getNonClosingChanges(transitionInfo)
 
             // 1. Build Parent-Child Relationships
             for (change in taskChanges) {
@@ -90,5 +82,16 @@ class SplitScreenUtils {
                 null
             }
         }
+
+
+        /** @return includes only opening + [TRANSIT_CHANGE] changes and the divider */
+        private fun getNonClosingChanges(transitionInfo: TransitionInfo): List<Change> {
+            return transitionInfo.changes
+                .filter { change ->
+                    (TransitionUtil.isOpeningMode(change.mode) || change.mode == TRANSIT_CHANGE)
+                            && change.flags < FLAG_FIRST_CUSTOM
+                }
+                .toList()
+        }
     }
 }