​

TDMPC

​

Overview

​

Key Features

• Model-Based RL: Learns a latent dynamics model of the environment
• MPC Planning: Uses Cross-Entropy Method (CEM) for trajectory optimization
• Hybrid Learning: Combines world model learning with value-based RL (TD learning)
• Efficient Inference: Leverages learned models for fast planning
• Single-Image Support: Works with single camera observations and proprioceptive state

​

Architecture

​

1. Observation Encoder

• Image Encoder: Convolutional network for processing visual observations
• State Encoder: MLP for processing proprioceptive state
• Latent Dimension: Typically 50-100 dimensional embedding

​

2. Latent Dynamics Model

z_{t+1} = f(z_t, a_t)

​

3. Reward Model

r_t = g(z_t, a_t)

​

4. Value Functions

• Q-Function Ensemble: Multiple Q-networks for uncertainty estimation
• V-Function: State value function trained with expectile regression

​

5. Policy Network (π)

a_t = π(z_t)

​

Training

​

Loss Components

1. Reward Loss: Predicts immediate rewards accurately
2. Value Loss: TD learning for Q and V functions
3. Consistency Loss: Ensures latent dynamics consistency
4. Policy Loss: Advantage-weighted regression for the policy

​

Training Command

lerobot-train \
    --dataset.repo_id=your_dataset \
    --policy.type=tdmpc \
    --output_dir=./outputs/tdmpc_training \
    --job_name=tdmpc_training \
    --policy.device=cuda \
    --batch_size=256 \
    --steps=100000

​

Key Training Parameters

​

Inference

​

Model Predictive Control

1. Initialize: Sample action sequences from a Gaussian distribution
2. Rollout: Use learned world model to simulate trajectories
3. Evaluate: Compute trajectory values using Q-functions
4. Update: Re-fit Gaussian to elite trajectories
5. Iterate: Repeat for several CEM iterations
6. Execute: Return first action(s) from best trajectory

​

CEM Parameters

​

Policy-Only Mode

​

Configuration

​

Basic Configuration

from lerobot.policies.tdmpc import TDMPCConfig

config = TDMPCConfig(
    # Input/Output
    n_obs_steps=1,
    n_action_repeats=2,
    horizon=5,
    n_action_steps=1,
    
    # Architecture
    image_encoder_hidden_dim=32,
    state_encoder_hidden_dim=256,
    latent_dim=50,
    q_ensemble_size=5,
    mlp_dim=512,
    
    # RL
    discount=0.9,
    
    # MPC
    use_mpc=True,
    cem_iterations=6,
    n_gaussian_samples=512,
    n_pi_samples=51,
    
    # Training
    reward_coeff=0.5,
    value_coeff=0.1,
    consistency_coeff=20.0,
    pi_coeff=0.5,
)

​

Normalization

normalization_mapping = {
    "VISUAL": "IDENTITY",  # Images in [0, 255]
    "STATE": "IDENTITY",   # Raw state values
    "ENV": "IDENTITY",     # Environment state
    "ACTION": "MIN_MAX",   # Actions normalized to [-1, 1]
}

​

Training Data Augmentation

• max_random_shift_ratio (default: 0.0476): Maximum random shift as proportion of image size
• Applied to square images only
• Improves robustness to small visual variations

​

Target Networks

• Target Model: Exponential moving average (EMA) of the main model
• target_model_momentum (default: 0.995): EMA coefficient for target updates

# Target network update
θ_target ← α * θ_target + (1 - α) * θ

​

Action Execution

​

Action Repeats

• n_action_repeats (default: 2): Number of times to repeat each action
• Reduces planning frequency and improves stability
• Common in model-based RL for real-world robotics

​

Action Steps from Plan

n_action_steps = 2  # Execute first 2 actions from plan
n_action_repeats = 1  # No repeats
use_mpc = True  # Required

​

Use Cases

• Manipulation tasks with visual and proprioceptive observations
• Simulated environments where model learning is effective
• Offline RL scenarios with pre-collected datasets
• Sample-efficient learning when online interaction is expensive

​

Limitations

• Currently supports only single camera observations
• Requires square images for random shift augmentation
• Single observation step (no observation history)
• Model learning can be challenging in complex/stochastic environments

​

Performance Tips

1. Latent Dimension: Increase for complex tasks (50-100)
2. Horizon: Longer horizons allow better planning but increase computation
3. Ensemble Size: More Q-functions improve uncertainty estimation
4. CEM Iterations: More iterations improve plan quality
5. Consistency Weight: Higher values enforce stronger model consistency

​

Example: PushT

lerobot-train \
    --dataset.repo_id=lerobot/pusht \
    --policy.type=tdmpc \
    --output_dir=./outputs/tdmpc_pusht \
    --policy.horizon=5 \
    --policy.use_mpc=true \
    --policy.latent_dim=50 \
    --batch_size=256 \
    --steps=100000

​

Comparison with Other Policies

​

Implementation Notes

​

FOWM Extensions

• Improved offline-to-online finetuning
• Better initialization strategies
• Enhanced uncertainty estimation

​

Code Structure

lerobot/policies/tdmpc/
├── configuration_tdmpc.py  # Configuration class
├── modeling_tdmpc.py       # Policy implementation
└── processor_tdmpc.py      # Data preprocessing

​

Citation

@inproceedings{Hansen2022tdmpc,
  title={Temporal Difference Learning for Model Predictive Control},
  author={Nicklas Hansen and Xiaolong Wang and Hao Su},
  booktitle={ICML},
  year={2022}
}

​

Related Papers

• TD-MPC: Temporal Difference Learning for Model Predictive Control
• FOWM: Finetuning Offline World Models in the Real World

​

See Also

• HIL-SERL Guide - RL training guide
• Policy Concepts - Understanding different policy types