PaperlayerView exact revision on arXiv

Source-linked AI summary

Temporal Difference Learning for Model Predictive Control

Nicklas Hansen, Xiaolong Wang, Hao Su

arXiv:2203.04955v2cs.LGcs.RO

TL;DR

Model-based RL must balance planning and sample efficiency against expensive long-horizon optimization and difficult model learning. TD-MPC addresses this with jointly learned task-oriented latent dynamics and terminal value functions for short-horizon planning and long-term return estimation. It achieves superior sample efficiency and asymptotic performance across state- and image-based continuous-control tasks.

  • Problem

    Model-based RL offers planning and sample-efficiency advantages, but long-horizon planning is costly and accurate environment models and value functions are challenging to obtain.

  • Method

    TD-MPC jointly learns a task-oriented latent dynamics model and terminal value function with temporal-difference learning, using model rollouts for local trajectory optimization and value estimates for long-term returns.

  • Results

    TD-MPC achieves superior sample efficiency and asymptotic performance over prior model-based and model-free methods across DMControl and Meta-World continuous-control tasks.

  • Takeaways & Limitations

    Task-oriented latent modeling lets TD-MPC support diverse continuous-control settings, including sparse rewards, arbitrary input modalities, and image-based tasks.

Abstract

from arXiv · show

Data-driven model predictive control has two key advantages over model-free methods: a potential for improved sample efficiency through model learning, and better performance as computational budget for planning increases. However, it is both costly to plan over long horizons and challenging to obtain an accurate model of the environment. In this work, we combine the strengths of model-free and model-based methods. We use a learned task-oriented latent dynamics model for local trajectory optimization over a short horizon, and use a learned terminal value function to estimate long-term return, both of which are learned jointly by temporal difference learning. Our method, TD-MPC, achieves superior sample efficiency and asymptotic performance over prior work on both state and image-based continuous control tasks from DMControl and Meta-World. Code and video results are available at https://nicklashansen.github.io/td-mpc.

1. Introduction

Model-based RL offers planning and potential sample-efficiency benefits, but long-horizon planning is costly and accurate models and value functions are difficult to learn. TD-MPC combines short-horizon latent-model planning with jointly learned long-term value estimates and performs strongly across continuous-control benchmarks.

  • Model-based methods can plan ahead and improve sample efficiency, but long-horizon planning is expensive and model bias can propagate into learned policies.
  • MPC limits planning to a finite horizon, while a terminal value function estimates discounted return beyond that horizon.
  • TD-MPC jointly learns a task-oriented latent dynamics model and terminal value function through temporal-difference learning.
  • TD-MPC uses model rollouts for short-term reward estimates and value estimates to guide long-term objectives such as running direction.
  • The model-learning design uses reward-focused representations, multi-step gradient propagation, and latent-space temporal consistency instead of predicting all observations.
  • TD-MPC achieves superior sample efficiency and asymptotic performance across DMControl and Meta-World, solving Humanoid and Dog tasks in as little as 1M environment steps.

2. Preliminaries

The preliminaries formulate infinite-horizon continuous-control MDPs and contrast fitted Q-iteration with finite-horizon model predictive control. MPC obtains local action sequences by optimizing model-generated trajectories and executing only the first action.

  • Problem formulation: The MDP formulation uses continuous state and action spaces, transition dynamics, rewards, discount factor γ, and an initial-state distribution.
  • Fitted Q-iteration: Fitted Q-iteration estimates the optimal state-action value function by repeatedly fitting Qθ to targets using a replay buffer and slowly updated target parameters.
  • Model Predictive Control: Actor-critic methods use a policy network to approximate globally optimal actions from learned value estimates.
  • Model Predictive Control: MPC estimates actions over a finite horizon H, obtains a local trajectory-optimization solution, and executes only the first action at each step.
  • Model Predictive Control: Unlike fitted Q-iteration, standard MPC typically uses γ = 1 and optimizes sampled model trajectories without directly predicting long-term rewards.

3. TD-Learning for Model Predictive Control

TD-MPC performs receding-horizon planning with MPPI over trajectories generated by learned latent dynamics, rewards, and terminal values. It combines sampled optimization with policy guidance, exploration schedules, and warm starts.

  • TD-MPC framework: TD-MPC combines MPPI planning, learned latent dynamics and reward models, a terminal state-action value function, and a policy network.
  • MPPI planning: The planner samples model-generated trajectories, estimates their returns, and updates a time-dependent Gaussian action distribution using importance-weighted high-return trajectories.
  • Inference: After a fixed number of optimization iterations, TD-MPC samples an action sequence from the final return-normalized distribution and executes the first action.
  • Inference: Trajectory returns combine discounted predicted rewards over H steps with a terminal value estimate at the rollout endpoint.
  • Inference: TD-MPC reuses a shifted previous distribution to warm-start optimization while retaining large initial variance to avoid local minima.
  • Exploration by planning: Exploration and planning schedules control optimization: action-noise decay is regulated, and the planning horizon grows from 1 to H while the model is inaccurate.
  • Policy-guided trajectory optimization: Policy trajectories are sampled alongside optimized trajectories, providing guidance while the learned policy remains less effective than planning on complex tasks.

4. Task-Oriented Latent Dynamics Model

TOLD is a task-oriented latent dynamics model jointly learned with a terminal value function through temporal-difference learning. It supports short-horizon trajectory optimization by predicting reward-relevant latent dynamics, rewards, values, and actions while using latent consistency rather than observation prediction.

  • Model overview: TOLD jointly learns latent dynamics and a terminal value function using temporal-difference learning.The model is used with TD-MPC for trajectory optimization and long-term return estimation.
  • Model components: TOLD comprises representation, latent dynamics, reward, value, and policy components that recurrently predict task-relevant quantities.The representation encodes observations; subsequent components predict future latent states, rewards, Q-values, and approximately value-maximizing actions.
  • Training objective: Training minimizes a temporally weighted objective combining reward prediction, value prediction, and latent state consistency losses.Near-term predictions receive higher weight through λ, while the trajectory is sampled from a replay buffer.
  • Training objective: Recurrent predictions remain entirely in latent space after only the first observation is encoded, allowing gradients from all three losses to back-propagate through time.This differs from approaches that learn state or video prediction separately from policy or value learning.
  • Value learning: The learned policy approximates Q-function maximization for efficient value learning because directly computing the planned maximum is costly.For complex tasks, the paper reports that planning remains superior to the learned policy, while the policy is sufficiently expressive for value learning.
  • Latent state consistency: Latent state consistency matches predicted future latents to target-network encodings of corresponding observations, avoiding direct observation prediction.The design reduces the need to model task-irrelevant details such as shading and is agnostic to observation modality.

5. Experiments

TD-MPC is evaluated across diverse continuous-control benchmarks, modalities, and computational settings against model-free and model-based baselines. It generally matches or outperforms competing methods while maintaining favorable sample and wall-time efficiency.

  • Tasks and benchmarks: TD-MPC is evaluated on 92 diverse tasks spanning DMControl and Meta-World, including sparse rewards, image observations, multimodal inputs, and multitask settings.
  • Baselines: The experiments compare TD-MPC with SAC, LOOP, simulator-based MPC, pixel-based model-free methods, latent-model methods, and discrete-action planning approaches.
  • Image-based control: On image-based tasks, TD-MPC uses shared hyperparameters across tasks and remains competitive with carefully tuned state-of-the-art methods.
  • Comparison to other methods: TD-MPC outperforms or matches baselines on most tasks, with larger gains on complex Humanoid, Dog, and Bin Picking tasks.
  • Training wall-time: TD-MPC solves Walker Walk 16× faster than LOOP while using 3.3× less compute per 500k environment steps.

6. Related Work

Related work spans model-free, model-based, and hybrid reinforcement-learning methods. TD-MPC combines elements explored previously but distinguishes itself through jointly learning a model and terminal value function with TD-learning for continuous control.

  • Temporal Difference Learning: Model-free TD-learning methods such as DDPG and SAC learn policies and value functions without learning an environment model.
  • Model-based RL: Model-based reinforcement learning uses learned models either for planning or for generating data to train model-free algorithms.
  • Model-based RL: EfficientZero learns latent dynamics from rewards, but its MCTS action selection is incompatible with continuous action spaces.
  • TD-MPC: TD-MPC is presented as the first framework to jointly learn a model and terminal value function through TD-learning in continuous control.

7. Conclusions and Future Directions

The paper concludes that TD-MPC performs strongly across diverse control settings and supports transfer between related tasks. It identifies model use, exploration, and architecture as directions for further improvement.

  • Conclusions: TD-MPC outperforms model-based and model-free methods on diverse continuous-control tasks and matches state-of-the-art image-based RL with trivial modifications.
  • Future Directions: Future improvements include using the learned model in new ways, better exploration strategies, and architectural innovations.
  • Model Generalization: Freezing the representation nearly preserves finetuning performance, whereas freezing both representation and dynamics substantially slows convergence.
  • Model Generalization: The transfer results suggest that the representation captures information transferable across related tasks, while the dynamics predictor encodes more task-specific behavior.
  • Computational Budget: For complex tasks, more planning generally improves performance, while inference cost can be reduced by 50% without performance loss by reducing planning iterations.

D. Latent Dynamics Objective

The section studies alternatives to TD-MPC’s latent dynamics objective and examines planning-driven exploration under different task difficulties. It also situates the comparison within the paper’s broader component-level analysis.

  • Latent Dynamics Objective: The latent dynamics objective is ablated by replacing the proposed consistency loss with contrastive, reconstruction, or no-regularization alternatives.
  • Comparison to Prior Work: The paper presents TD-MPC as the first complete continuous-control MPC framework using a model learned by TD-learning.
  • Computational Budget: The computational-budget analysis finds larger planning-policy gaps and greater budget sensitivity on high-dimensional or complex-dynamics tasks.
  • Exploration: Planning uncertainty decreases during training and converges as Humanoid Stand, Walk, and Run become solved.
  • Exploration: Exploration decreases more slowly on harder tasks, which the authors conjecture reflects greater variance in reward and value estimates.

F. Implementation Details

The implementation uses task-specific neural components, recurrent latent-model training, and extensive baseline and ablation configurations. It also evaluates alternative latent dynamics objectives and documents MPC comparison settings.

  • Network architectures: State-based experiments use a one-hidden-layer 256-unit MLP representation, while image-based experiments use a four-layer CNN with 32 filters per layer.Other components use two-layer 512-unit MLPs.
  • Training procedure: TOLD training recurrently predicts rewards, values, and latent states, weighting losses by temporal discount rho and updating a slow-moving target network.The update combines reward, value, and consistency losses before backpropagation.
  • Objective evaluation: Figure 10 compares TD-MPC objectives on 15 state-based DMControl tasks, including no regularization, state reconstruction, contrastive loss, and latent state consistency.Results are means over five runs with 95% shaded intervals.
  • Baselines: TD-MPC is compared with vanilla simulator-based MPC using a planning horizon of 10 and four iterations, versus six iterations for TD-MPC.The simulator baseline has no terminal value function.
  • Ablations: The no-latent ablation replaces the learned representation with the identity, modeling dynamics directly in state space while retaining reward- and value-gradient backpropagation.Its consistency loss therefore becomes a state prediction loss.
  • Ablations: The no-consistency ablation sets c3 to 0, training the TOLD model only with reward and value prediction losses.This makes the ablation more similar to MuZero.

H. Inference Time

TD-MPC inference time increases approximately linearly with planning computational budget, while the budget can be reduced for latency-sensitive applications. The method also supports multimodal locomotion and complex manipulation evaluations.

  • Planning budget: Inference time has an approximately linear relationship with computational budget as planning horizon or iteration count changes.The measurement varies one planning hyperparameter while holding the other at its default.
  • Latency: TD-MPC’s default training settings require approximately 20ms per decision step, equivalent to 50Hz, on an RTX3090 GPU.The paper identifies this as fast enough for many real-time robotics applications assuming an onboard GPU.
  • Latency: Reducing Quadruped Run’s planning horizon from 5 to 1 causes no significant performance reduction and lowers inference time to approximately 12ms per step.The computational budget can therefore be adjusted when inference time is critical.
  • Policy comparison: The jointly learned model-free policy is nearly 6× faster at inference than planning, although its performance is lower.This compares the policy’s inference speed with TD-MPC planning.
  • Meta-World: On Meta-World, SAC is competitive on most tasks, whereas TD-MPC is more sample efficient on complex manipulation tasks such as Bin Picking, Box Close, and Hammer.The evaluation uses goal-conditioned Meta-World v2 tasks and focuses on sample efficiency.
  • Multimodal locomotion: TD-MPC fuses proprioceptive data and egocentric camera input in Quadruped Corridor and Quadruped Obstacles locomotion tasks.The tasks require visual perception and continuous walking-direction adjustment.

K. Additional Metrics

The additional analysis reports aggregate TD-MPC and baseline performance across 15 state-based DMControl tasks using robust summary metrics and bootstrap confidence intervals.

  • Aggregate evaluation: Aggregate evaluation reports median, interquantile mean (IQM), and mean returns for SAC, LOOP, and TD-MPC after 500k environment steps.The metrics summarize episode returns from trained agents.
  • Uncertainty estimates: The reported confidence intervals are 95% intervals estimated with percentile bootstrap and stratified sampling.This follows the recommendation associated with the rliable toolkit.

L. Task Visualizations

The visualizations show successful TD-MPC trajectories across selected DMControl and Meta-World tasks, while Meta-World figures also compare success rates with SAC.

  • Trajectory visualizations: TD-MPC solves seven selected DMControl and Meta-World tasks in fewer than 1M environment steps.Figure 15 displays key frames rather than complete trajectories.
  • Meta-World visualizations: Figure 14 visualizes success rates for 24 of Meta-World’s 50 goal-conditioned tasks, with TD-MPC and SAC compared across tasks.The displayed tasks are sorted alphabetically and use means over five runs with 95% confidence intervals.
  • Task coverage: The selected trajectories include high-dimensional locomotion, sparse-reward exploration, and three-dimensional pick-and-place tasks.Examples include Dog Walk, Humanoid Walk, Finger Turn Hard, and Bin Picking.
Loading 2203.04955v2…