# DiffTORI: Differentiable Trajectory Optimization for Deep Reinforcement and Imitation Learning

arXiv: https://arxiv.org/abs/2402.05421 · HTML: https://mlsanity.com/paper/2402.05421/

## Key facts

- Cited by 6 papers in the map (data through 14 Oct 2026).
- Placed in Robot manipulation imitation learning, one of 13 directions in Offline RL and imitation learning from data.
- Builds on 45 of its 53 references that are in the map.
- First posted on arXiv 8 Feb 2024.

## Placed in

- Offline RL and imitation learning from data / [Robot manipulation imitation learning](https://mlsanity.com/offlinerl/robot-manipulation-imitation-learning/)

## Builds on (most cited first, 30 of 45)

- [Representation Learning with Contrastive Predictive Coding](https://mlsanity.com/paper/1807.03748/) (2018)
- [PointNet: Deep Learning on Point Sets for 3D Classification and Segmentation](https://mlsanity.com/paper/1612.00593/) (2016)
- [Soft Actor-Critic: Off-Policy Maximum Entropy Deep Reinforcement Learning with a Stochastic Actor](https://mlsanity.com/paper/1801.01290/) (2018)
- [Continuous control with deep reinforcement learning](https://mlsanity.com/paper/1509.02971/) (2015)
- [Diffusion Policy: Visuomotor Policy Learning via Action Diffusion](https://mlsanity.com/paper/2303.04137/) (2023)
- [Learning Fine-Grained Bimanual Manipulation with Low-Cost Hardware](https://mlsanity.com/paper/2304.13705/) (2023)
- [Learning Latent Dynamics for Planning from Pixels](https://mlsanity.com/paper/1811.04551/) (2018)
- [Mastering Diverse Domains through World Models](https://mlsanity.com/paper/2301.04104/) (2023)
- [Meta-World: A Benchmark and Evaluation for Multi-Task and Meta Reinforcement Learning](https://mlsanity.com/paper/1910.10897/) (2019)
- [Mastering Atari with Discrete World Models](https://mlsanity.com/paper/2010.02193/) (2020)
- [DeepMind Control Suite](https://mlsanity.com/paper/1801.00690/) (2018)
- [What Matters in Learning from Offline Human Demonstrations for Robot Manipulation](https://mlsanity.com/paper/2108.03298/) (2021)
- [OptNet: Differentiable Optimization as a Layer in Neural Networks](https://mlsanity.com/paper/1703.00443/) (2017)
- [Diffusion Policies as an Expressive Policy Class for Offline Reinforcement Learning](https://mlsanity.com/paper/2208.06193/) (2022)
- [Differentiable Convex Optimization Layers](https://mlsanity.com/paper/1910.12430/) (2019)
- [Implicit Behavioral Cloning](https://mlsanity.com/paper/2109.00137/) (2021)
- [Temporal Difference Learning for Model Predictive Control](https://mlsanity.com/paper/2203.04955/) (2022)
- [Mastering Visual Continuous Control: Improved Data-Augmented Reinforcement Learning](https://mlsanity.com/paper/2107.09645/) (2021)
- [ManiSkill2: A Unified Benchmark for Generalizable Manipulation Skills](https://mlsanity.com/paper/2302.04659/) (2023)
- [ManiSkill: Generalizable Manipulation Skill Benchmark with Large-Scale Demonstrations](https://mlsanity.com/paper/2107.14483/) (2021)
- [Differentiable MPC for End-to-end Planning and Control](https://mlsanity.com/paper/1810.13400/) (2018)
- [Model Predictive Path Integral Control using Covariance Variable Importance Sampling](https://mlsanity.com/paper/1509.01149/) (2015)
- [Objective Mismatch in Model-based Reinforcement Learning](https://mlsanity.com/paper/2002.04523/) (2020)
- [Deep Declarative Networks: A New Hope](https://mlsanity.com/paper/1909.04866/) (2019)
- [Differentiable Integrated Motion Prediction and Planning with Learnable Cost Function for Autonomous Driving](https://mlsanity.com/paper/2207.10422/) (2022)
- [Deep Reinforcement Learning in Parameterized Action Space](https://mlsanity.com/paper/1511.04143/) (2015)
- [Theseus: A Library for Differentiable Nonlinear Optimization](https://mlsanity.com/paper/2207.09442/) (2022)
- [ToolFlowNet: Robotic Manipulation with Tools via Predicting Tool Flow from Point Clouds](https://mlsanity.com/paper/2211.09006/) (2022)
- [Pontryagin Differentiable Programming: An End-to-End Learning and Control Framework](https://mlsanity.com/paper/1912.12970/) (2019)
- [Mismatched No More: Joint Model-Policy Optimization for Model-Based RL](https://mlsanity.com/paper/2110.02758/) (2021)

## Built on by (most cited first, 6 of 6)

- [Differentiable Model Predictive Control on the GPU](https://mlsanity.com/paper/2510.06179/) (2025)
- [Bi-Level Reinforcement Learning Pathway for Sim-to-Real Optimality](https://mlsanity.com/paper/2510.17709/) (2025)
- [Model-free Reinforcement Learning for Model-based Control: Towards Safe, Interpretable and Sample-efficient Agents](https://mlsanity.com/paper/2507.13491/) (2025)
- [Quantifying Generalisation in Imitation Learning](https://mlsanity.com/paper/2509.24784/) (2025)
- [CORL: Reinforcement Learning of MILP Policies Solved via Branch and Bound](https://mlsanity.com/paper/2512.11169/) (2025)
- [Vela: Scaling Vision-Language-Action Models with Adaptive Action Curve Parametrization](https://mlsanity.com/paper/2610.05230/) (2026)

## Abstract (arXiv)

This paper introduces DiffTORI, which utilizes Differentiable Trajectory Optimization as the policy representation to generate actions for deep Reinforcement and Imitation learning. Trajectory optimization is a powerful and widely used algorithm in control, parameterized by a cost and a dynamics function. The key to our approach is to leverage the recent progress in differentiable trajectory optimization, which enables computing the gradients of the loss with respect to the parameters of trajectory optimization. As a result, the cost and dynamics functions of trajectory optimization can be learned end-to-end. DiffTORI addresses the ``objective mismatch'' issue of prior model-based RL algorithms, as the dynamics model in DiffTORI is learned to directly maximize task performance by differentiating the policy gradient loss through the trajectory optimization process. We further benchmark DiffTORI for imitation learning on standard robotic manipulation task suites with high-dimensional sensory observations and compare our method to feed-forward policy classes as well as Energy-Based Models (EBM) and Diffusion. Across 15 model-based RL tasks and 35 imitation learning tasks with high-dimensional image and point cloud inputs, DiffTORI outperforms prior state-of-the-art methods in both domains. Our code is available at https://github.com/wkwan7/DiffTORI.