Research map / Reinforcement learning / Policy gradient and actor-critic
DiffTORI: Differentiable Trajectory Optimization for Deep Reinforcement and Imitation Learning
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes DiffTORI, which uses differentiable trajectory optimization (Theseus) as the policy representation, learning cost and dynamics end-to-end for model-based RL (on TD-MPC) and imitation learning on image and point-cloud benchmarks.
Problem
Policy representations for deep RL and imitation learning with high-dimensional observations; model-based RL suffers from objective mismatch where surrogate losses do not optimize task return.
Approach
Use trajectory optimization with neural dynamics and cost functions as the policy, and backpropagate a policy-gradient loss (RL) or imitation loss through the differentiable optimizer (Theseus) so the learned dynamics and cost are optimized directly for task performance; a CVAE decoder gives multimodal actions.
Key results
- On MetaWorld 22 tasks, DiffTORI+DP3 averages 56.5% success vs 48.8% for DP3 and 49.8% for Residual+DP3.
- On Robomimic/Push-T, DiffTORI+Diffusion Policy reaches an average failure rate of 0.07 vs 0.10 for Diffusion Policy.
- On ManiSkill (9 tasks), DiffTORI+BC averages 0.42 success vs 0.34 for BC; on DMC, the paper reports outperforming TD-MPC on 14 of 15 tasks.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| add differentiable trajectory optimization and policy-gradient learning to TD-MPC | TD-MPC on DeepMind Control Suite image tasks · DeepMind Control Suite, 15 tasks | return | Helped |
| add policy-gradient loss and backpropagate it through trajectory optimization (ablation) | DiffTORI · DeepMind Control Suite, 15 tasks | return | Helped |
| replace MPPI with trajectory optimization without adding policy-gradient loss (ablation) | TD-MPC · DeepMind Control Suite, 15 tasks | return | No clear effect |
| add DiffTORI action refinement initialized from DP3 | DP3 · MetaWorld 22 tasks | success rate: 48.8 -> 56.5 | Helped |
| add DiffTORI action refinement initialized from BC | BC · ManiSkill 9 tasks | success rate: 0.34 -> 0.42 | Helped |
| add DiffTORI action refinement initialized from Diffusion | Diffusion · Robomimic tasks average | failure rate: 0.10 -> 0.07 | Helped |
| remove the reward prediction loss (ablation) | DiffTORI model-based RL · DeepMind Control Suite | return | Hurt |
| remove the dynamics prediction loss (ablation) | DiffTORI model-based RL · DeepMind Control Suite | return | Hurt |
| remove initialization from the learned policy (ablation) | DiffTORI model-based RL · DeepMind Control Suite | return | Hurt |
| replace base-policy initialization with zero initialization (ablation) | DiffTORI · Robomimic and ManiSkill | failure rate on Robomimic; success rate on ManiSkill | Hurt |
| replace base-policy initialization with random initialization (ablation) | DiffTORI · Robomimic and ManiSkill | failure rate on Robomimic; success rate on ManiSkill | Hurt |
| increase the imitation-learning planning horizon from H = 1 to H = 3 (ablation) | DiffTORI · Robomimic tasks | failure rate: 0.08 -> 0.08 (Square); 0.04 -> 0.06 (Transport); 0.08 -> 0.08 (ToolHang); 0.09 -> 0.12 (Push-T) | No clear effect |
Surprising findings
- Zero or random action initialization for DiffTORI degrades performance to near or below vanilla BC on ManiSkill, so a good base-policy initialization is needed.
- Residual learning on top of DP3 sometimes lowers success rates relative to the base policy.
Limitations
- Training wall-clock time is substantially higher than TD-MPC due to differentiating through trajectory optimization.
- Imitation performance depends on a good action initialization from a base policy; zero/random init degrades results.
- Learned objective is not guaranteed to be a reward function and connections to inverse RL are left unexplored.
Experimental setup
Datasets: DeepMind Control Suite Robomimic ManiSkill1 ManiSkill2 MetaWorld IBC Push-T
Benchmarks: DeepMind Control Suite (15 visual tasks) MetaWorld (22 tasks) Robomimic (Square, Transport, ToolHang) Push-T ManiSkill (9 tasks)
Metrics: normalized return success rate failure rate cost
Baselines: SAC DrQ-v2 TD-MPC DP3 Residual + DP3 IBC BC-RNN Diffusion Policy BC BC + residual
Builds on: TD-MPC Theseus Diffusion Policy DP3 Amos et al. differentiable MPC
Closest papers by meaning
- DiffStitch: Boosting Offline Reinforcement Learning with Diffusion-based Trajectory Stitching (ICML 2024 · similarity 0.89)
- Diffusion Policy Policy Optimization (ICLR 2025 · similarity 0.88)
- Action-Constrained Imitation Learning (ICML 2025 · similarity 0.88)
- TMPO: Trajectory Matching Policy Optimization for Diverse and Efficient Diffusion Alignment (NeurIPS 2026 · similarity 0.87)
- Model-based Diffusion for Trajectory Optimization (NeurIPS 2024 · similarity 0.87)
- Implicit Drifting Policy: One-Step Action Generation via Conditional Expert Geometry (NeurIPS 2026 · similarity 0.87)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.