Universal Planning Networks

1. Introduction

• Raw-pixel goal rewards are uninformative and overconstrained, while prior visual representations were often learned with objectives only indirectly connected to control utility.The paper identifies effective representation learning for planning and control as an open problem.Source: §1 · pp. 1–2
• UPNs embed differentiable planning in a goal-directed policy and train its latent encoder and forward dynamics end-to-end against an imitation objective.The method directly optimizes for representations in which gradient-based planning succeeds on the goal-directed task.Source: §1 · p. 2
• UPN representations support gradient-based visual imitation and provide a latent distance metric for specifying goals and rewards for reinforcement learning on new target states.The paper reports more effective learning with image targets using these distance-based rewards.Source: §1 · p. 2
• UPNs learn visual goal-directed policies with less data than traditional imitation learners and transfer task-related semantics to more challenging agents and tasks.The reported findings also indicate that planning can improve with more test-time updates, including in lower-data settings.Source: §1 · p. 2

2. Universal Planning Networks

• UPNs directly optimize latent representations and forward dynamics for gradient-based trajectory optimization that produces task-desired actions, rather than relying on surrogate unsupervised objectives.This approach is motivated by modeling errors, irrelevant details, and task-misaligned plans in learned forward models.Source: §2 · p. 2

2.1. Learning to Plan

• Given initial and goal observations, the GDP uses an encoder and latent forward model to iteratively optimize an action sequence toward the encoded goal.Each planning update unrolls latent dynamics and uses gradients with respect to the planned actions to reduce terminal-state distance to the goal.Source: §2.1 · pp. 2–3, Figure 2 · p. 3

2.2. Imitation as the Outer Objective

• An outer imitation loss supervises the full planning computation graph, training representations that make iterative planning effective for the demonstrated task.The learned encoder thereby provides an optimizable latent distance metric that can support reward functions for reinforcement learning.Source: §2.1 · p. 3, §2.2 · p. 4

2.3. Reinforcement Learning with a UPN Latent Space

• The UPN encoder’s latent distance supplies rewards for image-goal reinforcement learning, while the agent receives its own joint state and the goal’s latent feature as inputs.The described setup uses perceptual rewards rather than extrinsic environment rewards.Source: §2.3 · pp. 4–5, Figure 3 · p. 5

3. Related Work

• Unlike unsupervised approaches using iLQR with image targets, UPNs optimize representations directly for plannability; unlike self-supervised intermediate objectives, they learn representations for new tasks from a single goal image.Pixel reconstruction may encode state variables unhelpful for planning or discard visually inconspicuous ones; Sermanet et al. (2017) instead recover rewards to mimic specific demonstrations.Source: §3 · p. 5
• Prior state-representation work for model-free RL uses external rewards, whereas this work learns representations for metric-based rewards on new tasks; its connection to IRL and reward shaping likewise yields rewards effective on new tasks.IRL from raw pixels is challenging and successful image-based methods cited here rely on human domain knowledge or pre-training.Source: §3 · p. 5
• Unlike value iteration networks limited to discrete planning or differential MPC approaches requiring available cost functions, this work targets continuous motor control from raw pixels when cost functions are unavailable.The authors position their work as the first to connect differentiable planning procedures with learning reusable representations that generalize across complex visuomotor tasks.Source: §3 · p. 5
• Earlier gradient-based planning used known analytic dynamics, while later latent-space planning approaches relied on unsupervised pretraining or extrinsic rewards; this work instead focuses on continuous control without extrinsic rewards.The cited latent-space approaches predict forward or decode a value function; distance to goal in UPN's learned latent space can also be viewed as a value function.Source: §3 · pp. 5–6
• Goal-conditioned imitation policies also accept goal images, but the authors distinguish their approach by reusing its learned representation for planning and reward specification.The paper frames its representation as broadly applicable to performing new tasks from a single goal image.Source: §3 · pp. 5–6

4.1. UPNs Learn Effective Imitation Policies

• UPN supports test-task generalization from fewer demonstrations; AIL nearly matches its performance as the dataset grows.Figure 5 compares success across four generalization conditions and shows UPN is more sample efficient.Source: §4.1 · p. 6, Figure 5 · p. 7

4.2. Analysis of the Gradient Descent Planner

• 38.1% to 64.44% task success: a 20,000-demonstration UPN improves from 40 to 160 test-time planning steps; expert average success is 73.12%.With 160 steps, UPN matches expert performance; the relative benefit of extra steps shrinks with more demonstrations.Source: §4.2 · p. 7, Figure 6 · p. 7

4.3. Latent Space Visualization

• The learned latent distance is obstacle-aware and task-specific: poses near but below the initial position are less desirable than farther-above poses.Figure 7 maps larger latent distances to lighter colors.Source: §4.3 · p. 7, Figure 7 · p. 8

4.4. Transfer to Harder Scenarios

• UPN-feature RL outperforms the expert and UPN-160, while RIL- and AIL-feature RL fail and VAE trails UPN on the Figure 4(d) task.Table 1 reports average success rates for fixed and varying goals.Source: Table 1 · p. 8, §4.4 · p. 8

4.5. Transfer Across Robots

• Across robot-transfer scenarios, UPN-derived RL rewards significantly outperform VAE, AIL, and RIL feature spaces, supporting transfer to tasks with different dynamics and action spaces.The comparisons include point-robot-to-ant and pushing-from-poking transfer.Source: §4.5 · pp. 8–9, Figure 9 · p. 9

4.6. 3D 7-DoF Control from an Non-orthographic View

• UPN rewards from poking transfer to pushing, with image-based UPN-RL approaching the task success of a well-shaped reward.The result suggests the latent space captures end-effector proximity to the object as a prerequisite for moving it to the goal.Source: §4.6 · p. 9

4.8. Using UPN rewards for long horizon tasks

• 8000 steps: point-robot UPN rewards guide ant navigation around obstacles, addressing sparse-reward credit assignment over a long horizon.Figure 10 describes locomotion to a goal by the ant, whose policy is optimized over 8000 time steps.Source: Figure 10 · p. 10, §4.8 · p. 10

4.7. Transfer from Point Robot to a Humanoid: Pushing the limits of generalization

• Humanoid transfer re-renders center-of-mass coordinates as point-robot images; an explicit staying-upright reward and decaying curriculum supplement UPN path rewards.Naive fall termination led the humanoid to balance near side walls instead of optimizing the path-planning reward.Source: §4.7 · p. 10

5. Discussion

• UPNs learn effective visual goal-directed policies, transfer task-related semantics to harder agents and tasks through goal-conditioned rewards, and improve planning with additional test-time updates.Source: §5 · p. 10
• The authors propose studying reinforcement learning or self-supervision for training UPN representations and structuring their metrics as value functions for longer-horizon optimization.Source: §5 · p. 10
• UPN-like architectures might make complex real-robot behaviors practically learnable by leveraging simulated behaviors from simpler robots performing the same tasks.Source: §5 · p. 11

A.1. Input / A.2. Encoder fφ

• Inputs are 84 × 84 × 3 RGB images scaled by 1/255 without standardization; the encoder outputs 128-dimensional vectors and uses swish, which worked better than relu.Source: §A.1 · p. 14, §A.2 · p. 14

A.3. Bias Transformation and Embodiment Information

• For torque control, joint angles and velocities supplement image features because current pixels alone leave velocities partially observed; bias-transform variables help control inner-loop gradients.Source: §A.3 · p. 14

A.4. Action Encoder

• The action encoder converts action sequences into 64-dimensional latent vectors for use by the forward dynamics model.Source: §A.4 · p. 15

A.5. Planning by Gradient Descent

• The planner rolls a latent-space dynamics model forward from the current state using its current estimate of the action plan.Source: §A.5 · p. 15

A.5.1. DYNAMICS IN LATENT SPACE

• One-step latent dynamics take the current state and action as input and predict the next latent state.Source: §A.5.1 · p. 15

A.5.2. PLANNING UPDATES

• Planning updates minimize Huber loss between predicted and goal latents by taking gradients with respect to the action plan and applying gradient descent.Source: §A.5.2 · p. 15

A.5.3. HYPERPARAMETER DETAILS

• Table 2 lists planning-module hyperparameters, including maximum rollout horizon; training uses dynamic horizons and favors shorter ones early.Source: Table 2 · p. 15, §A.5.3 · p. 16
• The inner loop clips gradients to [-25, 25] and uses Huber delta 0.85; step sizes are 0.5 for the first update and 0.25 for the next 39.Source: §A.5.3 · p. 16

A.6. Outer Loop Loss and Training

• The outer loop trains updated plans against demonstrated optimal plans with continuous-action mean-squared error; batch size 128 was crucial, and shorter horizons were oversampled early.Short-horizon Poisson sampling was used for the first 300,000 minibatch updates.Source: §A.6 · p. 16

A.7. Things that did not work

• Weight normalization failed; batch sizes 256 matched 128 while 512 destabilized training, and input batch normalization and convolutional layer normalization hurt performance.Weight normalization was tried to control large inner-loop gradients, latent features, and planning error.Source: §A.7 · pp. 16–17

B. Action-selection with Model Predictive Control (MPC)

• At test time UPN plans toward the goal but executes only the first inferred action before replanning; RIL instead predicts one action without MPC.The trajectory-planning and first-action execution cycle is illustrated in Figure 11.Source: B. Action-selection with Model… · p. 17, Figure 11 · p. 17

C. Demo collection

• Because HER demonstrations can be imperfect and policies configuration-specific, hindsight rendering retains successful rollouts or relabels failed ones with their final position.HER policies are trained independently for each obstacle configuration in the training distribution.Source: C. Demo collection · p. 17

D.1. Reactive Imitation Learner (RIL)

• RIL maps the image observation through a convolutional encoder and fully connected decoder to predict an action.The comparator retains the UPN encoder for a fair comparison.Source: §D.1 · p. 17

D.2. Autoregressive Imitation Learner (AIL)

• AIL encodes the observation and goal into an initial hidden state, then recurrently decodes an action sequence over the planning horizon.Its recurrent cell is initialized with the latent representation produced by the encoder.Source: §D.2 · p. 17

RIL

• RIL fuses the retained UPN convolutional encoding with joint information through two 128-dimensional hidden layers before its action output layer.Source: §D.3 · p. 17

AIL

• AIL concatenates joint encoding with the retained UPN encoder and uses a 128-dimensional vanilla RNN over the planning horizon; more sophisticated cells showed no observed benefit.AIL uses skewed batch sampling to encourage curriculum learning and decodes actions from recurrent hidden states.Source: §D.3 · p. 18

• UPN, RIL, and AIL use the same train, validation, and test splits.Source: §D.3 · p. 18

E. Imitation Learning Benchmarks Experimental Setup

• For point-robot and reacher rollouts, success means the end-effector—or point-robot center of mass—is within 0.05 meters of the goal.The task-specific hyperparameter table covers planning horizon, demonstration count, and action count.Source: E. Imitation Learning Benchmarks… · p. 18, Table 5 · p. 18

F. PPO Parameters

• The PPO hyperparameter tables cover simpler RL tasks, long-horizon ant navigation, humanoid experiments, and environment-specific details.Source: F. PPO Parameters · pp. 18–19
• PPO policies and value functions use three 128-dimensional fully connected layers with tanh nonlinearities; each experiment averages three random seeds.The policy is Gaussian with learned state-independent diagonal covariance.Source: F. PPO Parameters · p. 19

• The evaluation asks whether feature-space distance metrics can serve as rewards for new RL problems with image targets but no demonstrations or extrinsic rewards.Source: §G.1 · p. 19
• The comparison includes feature spaces learned by the supervised AIL and RIL models.Source: §G.1 · p. 19
• A VAE is trained on demonstration-trajectory images with a mean-squared-error decoder objective, and its learned representation is used as the VAE feature space.Source: §G.1 · p. 19

G.2. VAE Architecture Details

• The VAE uses a 128-dimensional Gaussian posterior, deconvolutional decoding, pixels scaled by 1/255, and a sigmoid output constrained to [0, 1].The encoder structure is kept consistent with UPN, RIL, and AIL.Source: §G.2 · p. 19
• More expressive pixel-prediction architectures may improve representation metrics, but are unlikely to change the paper’s conclusion because their latent codes lack agency or planning.The authors leave experimenting with autoregressive decoders for future work.Source: §G.2 · p. 19

G.3. Reacher Morphology Transfer

• For 3-link/4-link reacher transfer, UPN shares the encoder across robots but uses different forward models, training on 20,000 demonstrations per robot.The VAE instead trains on aggregated demonstration snapshots without actions; success means reaching within 0.05 meters of the goal.Source: §G.3 · pp. 19–20

G.4. Point Robot to Ant Transfer

• Point-to-ant transfer uses 40,000 trajectories from 40 randomized point-robot appearances, including non-actuated limbs, with ant success defined within 0.5 meters of the goal.The appearance randomization was intended to encourage representations that focus on the robot’s center.Source: §G.4 · p. 20

G.5. Learning to Push from Poking Trajectories

• The pushing experiment trains UPN on 20,000 successful pokes and tests skill extrapolation on the same robot, with success defined within 0.08 meters of the goal.Object initial positions are randomized; the RL policy is not given the goal image embedding at test time.Source: §G.5 · p. 20

H. Camera viewpoint for each task

• The 2D point-robot view depicts navigation around obstacles.Source: Figure 13 · p. 20

H.2. Ant around obstacles task

• The ant-robot view depicts navigation around obstacles.Source: Figure 14 · p. 20

H.3. Reacher around obstacles task

• The 3-link reacher view depicts navigation around obstacles.Source: Figure 15 · p. 20

H.4. Pusher

• The 7-DoF pusher view depicts an attempt to push an object to a goal.Source: Figure 16 · p. 20

I. Exponential of distance metrics for rewards

• Rewards use the exponential of the negative scaled Huber distance between the current and goal states, a normalization that yielded stable PPO training.This transformation was used across all reinforcement-learning experiments; task-specific forms are listed in Table 10.Source: §I · p. 21, Table 10 · p. 21

J. Open Source Frameworks

• Experiments used TensorFlow for differentiation through inner-loop planning, with tasks and RL methods built on Gym and Baselines and environments simulated in MuJoCo.Source: J. Open Source Frameworks · p. 21