Multiple Shooting for Gradient-Based Planning with World Models

Planning with world models is typically framed as a trajectory optimization problem where the goal is to find an action sequence that drives a system from an initial to a target state. Since world models are trained purely for prediction, the task-specific objective is specified only at inference, making the planner a key determinant of success. Standard approaches formulate this as a single shooting problem, where actions are the only optimization variables and the objective is evaluated by unrolling a learned latent dynamics model. However, optimizing through these rollouts compounds model errors, exposes the optimizer to many local minima, and is inherently sequential. Multiple shooting formulations sidestep these issues by treating states as additional optimization variables but traditionally rely on second-order methods that scale poorly to high-dimensional latent spaces. We propose PALM, a first-order multiple shooting planner for learned latent spaces. PALM treats latent states as explicit optimization variables and enforces dynamics consistency through an augmented Lagrangian. By relaxing dynamics constraints during optimization, PALM can move through infeasible regions and escape local minima where single shooting methods typically get stuck. We evaluate PALM on navigation (Wall, PointMaze) and manipulation (PushT) tasks. On Wall and PointMaze, PALM outperforms gradient-based baselines and degrades more slowly as the goal gets further: on the longest Wall crossing tasks, it succeeds on $82\%$ of tasks, against at most $31\%$ for other gradient-based planners. CEM achieves higher final success, but PALM reaches competitive success in less wall-clock time. On PushT, all planners reach similar success rates.