Branch-Prox: Joint Arrival-Time and Action Optimization with World Models

Goal-conditioned planning with world models often optimizes action sequences by matching the predicted terminal representation to a goal representation at a fixed horizon. However, this objective does not explicitly encourage reaching the goal sooner. We introduce Branch-Prox, a trajectory optimization method that jointly optimizes arrival time and action sequences through a differentiable world model. Each candidate arrival time defines a branch whose objective combines elapsed time with a penalty for exceeding a goal-matching tolerance in latent space. An earlier arrival time may appear unfavorable under the current action sequence yet become preferable after action refinement. Rather than selecting an arrival time based on the current trajectory alone, Branch-Prox solves a convex proximal subproblem for each candidate and jointly selects the arrival time and action update using the optimized local model values. Each box-constrained subproblem reduces to a one-dimensional dual search. Under standard smoothness assumptions on the goal-matching costs and exact branch solves, we establish sufficient decrease of the planning objective with backtracking line search. Branch-Prox integrates into model predictive control without world-model fine-tuning or auxiliary model training. Across three LeWorldModel case studies with 64 matched episodes and three planner seeds, Branch-Prox achieves the lowest mean failure-capped steps among common-horizon reference methods. On Reacher, the mean decreases from the strongest reference baseline's 34.40 steps to 13.57 with 100% native success. On DINO-WM Wall, Branch-Prox reaches all 50 goals with 19.30 mean capped steps, compared with 37.44 for iCEM.