DiffTORI: Differentiable Trajectory Optimization for Deep Reinforcement and Imitation Learning
Authors: Weikang Wan
Ziyu Wang
Yufei Wang
Zackory Erickson
David Held
Venue: NeurIPS 2024 spotlight (2024)
Keywords: imitation learning
model-based reinforcement learning
differentiable trajectory optimization
Primary area: deep_learning_architectures
TL;DR: This paper introduces DiffTORI, which uses Differentiable Trajectory Optimization as the policy representation to generate actions for deep reinforcement and imitation learning, and outperforms prior state-of-the-art methods in both domains.

Abstract: This paper introduces DiffTORI, which utilizes $\textbf{Diff}$erentiable $\textbf{T}$rajectory $\textbf{O}$ptimization as the policy representation to generate actions for deep $\textbf{R}$einforcement and $\textbf{I}$mitation 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 feedforward 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.

DiffTORI: Differentiable Trajectory Optimization for Deep Reinforcement and Imitation Learning
Weikang Wan't, Ziyu Wang",, Yufei Wang?;
Zackory Erickson", David Held
Computer Science and Engineering Department, University of California San Diego Institute for Interdisciplinary Information Sciences, Tsinghua University Robotics Institute, Carnegie Mellon University
#2wan@ucsd.c edu, ziyu-wan21 @mails. tsinghua. edu. en yufeiw20andrew. cmu, cdu, zackoryšcmu, edu, dheld@andrew cmu.edu edu
Abstract
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 dynamics function. The key to approach is leverage the recent progress in differen- tiable trajectory optimization, which enables computing the gradients of the loss with respect the parameters of trajectory optimization. As 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 RI. algorithms, the dynamics model DiR' is learned to directly maximize task performance by differentiating the policy gradient loss through the trajectory optimization pro- cess. 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-Basse Models (EBM) and Diffusion. Across 15 model- based RL tasks and imitation learning tasks with high-dimg -ith-mmnnniinal image and point cloud inputs, DiffTORI outperforms prior state-of-te the-hhhr methods in both domains.
Introduction
Recent works have shown that the representation of a policy can have a substantial impact on Prior works have explored the use of feed-forward neural networks [4) energy-based based models 2 or diffusion [1.5] as the policy representation. In this paper, as the policy representation to generate actions for deep reinforcement learning (RL) and imitation learning (IL) with high- dimensional sensory observations (images/point clouds).
the learning performance [
we propose use differentiable trajectory optimization
[ 6
Trajectory optimization an effective and widely used algorithm in control, defined with a cost function and dynamics function. It can be viewed as a policy 63, where the parameters of the policy specify the cost function and the dynamics function. Given the learned cost and dynamics functions as well as the input state (c.g.. images, point clouds, robot joint states), the policy then computes the actions by solving the trajectory optimization problem. Trajectory optimization can also be made to be differentiable, which allows back pack-propggating through the trajectory optimization [1] 12 [3]. In prior work, differentiable trajectory optimization has been
process B
10 6
Equul contribution. This work was performed when Weikang Wan and Ziyu Wang were interning CMU.
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

applied to system identification 6 inverse optimal control [6], imitation learning 3 and control/planning for robotics problems with low-dimensional states 3 8 [5]
6 8 14
In this paper, we propose to combine differentiable trajectory optimization with deep model-based RI. algorithms. Because we use differentiable trajectory optimization generate actions [T]] are able to compute the policy gradient loss on the generated actions to learn the dynamics and cost functions to optimize the reward. This approach addresses the "objective mismatch" issue [16 of current model-based RL algorithms, i.e. models that achieve better training performance (e.g., lower MSE) in learning a dynamics model are not necessarily better for control. Our method addresses this issue, as the latent dynamics and reward models are both optimized to maximize the task performance by back-pree pooo-pppggggiing the policy gradient loss through the trajectory optimization process. We show that our method outperforms prior state-of-the- ofrt model-based RL algorithms on 15 tasks from the DeepMind Control Suite 18 with high-dimensional image inputs.
We further benchmark our method for imitation learning on standard robotic manipulation task suites with high-dimensional sensory observations and compare our method to reed-forward policy classes as well as Energy-Based Models (EBM) [2] and Diffusion (1. and term our method DifTTORI (Differentiable Trajectory Optimization for Reinforcement and Imitation Learning). We observe that our training procedure using differentiable trajectory optimization leads to better performance compared the EBM approach used in prior work, which can suffer from training instability due to the requirement of sampling high-quality negative examples []] We also outperform diffusion- based approaches [ due our procedure of learning a cost function that we optimize at test time. We show DifT achieves state-of-the-arr art performance across 35 different tasks: 5 tasks from Robomimic 19 with image inputs, tasks from Maniskill1 (00 and Maniskil12 21 with point
cloud inputs, and 22 tasks from Meta World [22 with point cloud inputs.
in employing differentiable trajectory optimization as a policy representation. Compared to these prior work, we are the first to show how differentiable trajectory optimization can be combined with deep model based RL algorithms, training dynamics, reward, Q function, and the policy end-to-end using task loss. In contrast, prior work either focuses on imitation learning [| 8 assumes known dynamics and reward structures and learns only a few parameters [3) or first learns the dynamics model with the dynamics prediction loss (instead of the task loss), and then uses the fixed learned dynamics for control [ We are also the first to show that the policy class represented by differentiable trajectory optimization can scale up to high-dimensional sensory observations like images and point clouds, achieving state-of tht-hhhrrrr art performances in standard RL and imitation learning benchmarks. In contrast, prior works [3 & 6 only test their methods in customized tasks with ground-truth low-level states, and do not report performance on standard benchmarks with more complex tasks and digeeniimenii observations.
Our work closely related to prior work 3 8
In summary, the contributions of our paper are as following:
We introduce DifTTORI, which uses differentiable trajectory optimization as the policy representa- tion for deep reinforcement learning and imitation learning.
We conduct extensive experiments to compare DiffTORI against prior state-of-ther -ftt----t-- methods on 15 tasks for model-based RI. and 35 tasks for imitation learning in standard benchmarks with high-dimensional sensory observations, and show that DiffTORT achieves superior performances.
We perform analysis and ablations of DiffTORI to provide insights into its performance gains.
2 Related Works
Differentiable optimization and implicit policy representation: Our work follows the line of research on differentiable optimization, which embeds optimization problems as a layer in neural networks for end-to-end learning. Early works focus on differentiating through convex optimization problems (23 24 Recent works extend the range of optimization problems that can be made
differentiable IT
The mostly related prior work to ours are Amos et al. B and Jin et al. O which first proposed to treat trajectory optimization as an implicit policy and demonstrated its effectiveness the setting of behavior cloning, system identification, and control for robotics problems with low-dimensional states. Another closely related recent work is Romero et al. [15) where they embed differentiable quadratic program with learnable cost matrices and known dynamics into the last layer of the actor in PPO, with applications for quadcopter flying. Ours differ
6
9 10

---

from this work as we learn non- non-linear costs parameterized by a full neural network, and we also learn the dynamics instead of assuming is known. We also show our method work with high-dimensional sensory inputs such as images and point clouds. Cheng et al. 22 26 proposes to learn the parameters a a PID controller by unrolling the controller and system dynamics into a computation graph and optimizing the controller parameters via gradient descent with respect to the task loss, assuming known dynamics. DiffTORI does not assume any prior knowledge on the dynamics or policy class; Instead of representing the policy trajectory optimization with the learned dynamics, reward and Q functions represented as neural networks. Sacks et al. [27] proposes to learn the update rule in MPPI, represented as neural network, using reinforcement learning, with known dynamics and cost functions. Instead of learning the update rule, we learn the dynamics, reward, Q function used in trajectory optimization to generate the actions. We perform differentiable trajectory optimization instead of RI. to optimize the parameters of these functions. Differentiable optimization has also been applied in other robotics domains such as autonomous driving [14] (28, 29). navigation 7 30), motion planning [31 [2], and state
predefined controller, our policy is represented as performing
estimation [32] We are the first to show how differentiable trajectory optimization can be combined with deep model-based RL.
Model-based reinforcement learning: Compared to model-free RL, model-based RL usually has higher sample efficiency as it is solving a simpler supervised learning problem when learning the dynamics model. Recently, researchers have identified a fundamental problem for model- based RL, known as "objective mismatch" [16]. Recent works have proposed a single objective which is a lower bound on the true return of the policy, for joint model and policy learning in model-based RL [77 Our approach also addresses the objective mismatch problem. In contrast to these two prior work which only optimizes lower bound on the true return, our approach directly optimizes the task reward. Further, these approaches are only demonstrated using low-dimensional state- state-based observations whereas our approach is able handle high-dimensional image or point cloud observations. In contrast to these works, we use Theseus [T0] to analytically compute the gradient of the true objective for updating the model. Another related work, Nikishin al. 34 proposes to learn a dynamics and reward model in model-based RL, and derive an implicit policy as the softmax policy associated with the optimal Q function under the learned dynamics and reward, learned by back proppagaatgnnn the RL loss via implicit function theorem. In contrast, we derive the implicit policy as the optimal solution from performing trajectory optimization with the learned dynamics, reward and Q function.
Policy architecture for deep imitation learning: Imitation learning can be formulated as the supervised regression task of learning map observations actions from demonstrations. Some recent work explores different policy architectures (e.g., explicit policy, implicit policy 12]. diffusion policy (T) and different action representations (e.g.. mixtures of Gaussian (35] 191, spatial action
maps [36]. action flow 4. or parameterized action spaces (77 to achieve more accurate learning model the multimodal distributions of demonstrations, and to capture sequential correlation. Our method outperforms explicit or diffusion policy approaches due to our procedure of learning a cost function that we optimize at test time. In comparison with the implicit policy, which also employs test-time optimization with a learned obective, we use a different and more stable training procedure via differentiable trajectory optimization.
from demonstrations,
3 Background
3.1 Differentiable Trajectory Optimization
In robotics and control, trajectory optimization solves the following type of problems:
T-1
min
Σ c(8,, at) + C(sr)
(1)
8c+1 d( (1t az)
where c(st, at) and C(sr) are the cost functions, and 8t+1 d(su, az) is the dynamics function. In this paper, we consider the case where the cost function and the dynamics functions are neural networks parameterized by 0: ca(st, Co(S)). and do(st,a))
Let ao())) (0) be the optimal solution to the trajectory optimization problem, which is function of the model parameters 0. Differentiable trajectory optimization is a class of method that enables
3

---

TD-MPC
DiffTORI (Ours)
(magpy, reward)
Envirorment
(mrgge, reword]
Envirommnnt
Enooder
Policy Gradiont Loss
Differemiable
Inferenee
Training
Overview of DifTTORI for model-based RL. In contrast to prior work in model-based RL [38] that uses non-differentiable MPPI (left), we utilize differentiable trajectory optimization to generate actions (right). DiffTORI computes the policy gradient loss on the generated actions and back-propagates through the optimization process, optimize the encoder as well other latent space models (latent reward predictor and latent dynamics function) to maximize task performance.
Figure 1:
computation of the gradient of the actions with respect to the model parameters (o.(0) Specifically, in this paper we use Theseus [10), which is an efficient application-agnossee differentiable nonlinear least squares optimization. Theseus works well with high-ddmensional states, e.g., images or point clouds, along with using neural networks as the cost and dynamics functions.
open source library for
3.2 Model-Based RL preliminaries
We use the standard MDP formulation: A. T, 7) where S is the state space, A is the action space, (s, a) is the reward function, T(- a) is the transition dynamics function, and y E [0, 1) is the is the discount factor. The goal is to learn policy π to maximize the expected return: E, Ex 12 7 R(se, d.)]. In this paper we work on problems where the state space S are high- dimensional sensory observations, e.g., images or point clouds. Model-based RL algorithms first learn a dynamics model, and then use it for learning policy. When applied to model- RL, our method builds upon TD-MPC (38), recently proposed model-based RL algorithm which review briefly here. We choose MD-MPC for its simplicity and sate-of-the-art performance. However, our method is compatible with any model-based RL algorithm that learns dynamics model and reward function. TD-MPC consists of the following components: first, an encoder hg. which encodes the high-dimensional sensory observations, e.g., images, into low-dimensional state =, ho In the latent space, a latent dynamics model de is also learned: 2++1 = do (2t, az). A latent reward predictor Rg is learned which predicts the task reward r: f Rg [2,, (2,, az). Finally, a value predictor Qu learns to predict the Q value: Q = Q2 ac). Note that we use to denote all learnable parameters including the encoder, the latent dynamics model, the reward predictor, and the Q value predictor. These models are trained jointly using the following objective:
(TD-MP((0:T) Σ - (TD)MPC(0:TT).
(2)
where T ~ B is a trajectory (st, $1+1) sampled from a replay buffer B, X E R4 is constant that weights near-term predictions higher, and the single-step loss is:
Lusanc (0; v.) =cs| (R.,,, - rull + Qu(z., (r. + 7Q,- (2. (1)
rewand
calldd (2., a.) N1 (8L+ 11
where 8- are parameters of target networks that are periodically updated using the parameters of the learning networks. As shown in (3), the parameters 0 is optimized with set of surrogate losses (reward prediction, value prediction, and latent consistency), rather than directly optimizing the task performance, known as the objective mismatch issue (16). At test time, model predictive path integral (MPPI) (3 is used for planning actions that maximize the predicted rewards and Q functions in the latent space. policy Te is further learned in the latent space using the latent Q-value function, which is used to generate action samples in the MPPI process.

---

4 Method
4.1
Overview
The core idea of DiffTORI is to use trajectory optimization as the policy Te, where 0 parameterizes the dynamics and cost functions. Given a state s, DiffTORI generates the actions a()) by solving the trajectory optimization problem in ① with so = s. To optimize the policy parameters 6, we use differentiable trajectory optimization to compute the gradients of the loss [(a(f)) with respect the (C(a(())
policy parameters:
where the exact form of the loss depends on the problem setting.
An overview of applying DifiTORI to model-based RL is shown in Figure | Existing model- based RL algorithms such as TD-MPC suffer from the objective mismatch issue: the latent dynamics and reward (cost) functions are learned to optimize a set of surrogate losses (as in (3), instead of optimizing the task performance directly. DifTTORI addresses this issue: by computing the policy gradient loss on the optimized actions from trajectory optimization and differentiating through the trajectory optimization process, the dynamics and cost functions are optimized directly to maximize the task performance. We describe DiffTORI for model-based RL in Section 4.2
We also apply DifTORRI to imitation learning: an overview is shown in Figure In contrast to explicit policies that generate actions at test-time by forward passes of the policy network, DifTTORI generates the actions via test-time trajectory optimization with a learned cost function. This is in the same spirit of implicit behaviour cloning [2] which learns an energy function and optimizes with respect to it to generate actions test-time. However, we observe that our training procedure using differentiable trajectory optimization leads to better performance compared to the EBM approach used in prior work, which can suffer from training instability due to the requirement of sampling high-quality negative examples [ We describe DiffTORI for imitation learning in detail in Section 4.3
4.2 Differentiable trajectory optimization applied to model-based RL
We build DifiTORI top of TD-MPC for model-based RL. Similar to TD-MPC, DifTTORI consists of an encoder he. a latent dynamics model de, a reward predictor Rg, and a Q-value predictor Qs (see Sec. 3.2p We use 0 to denote all learnable parameters to be optimized in DiffTORI. As shown in Figure I the key to DiffTORI is to change the non-differentiable MPPI planning algorithm in TD-MPC to differentiable trajectory optimization, and include the policy gradient loss the generated actions to optimize the model parameters directly for task performance.
Formally, given a state 8t, we use the encoder he encode it to the latent state 21, and then construct the following trajectory optimization problem in the latent space:
a (0)
arg max
Σ
Ho(o ar) + , " "Q(Ez,
(4)
C H 1st
s.L. Z041 do(2),a))
where H is the planning horizon. In this paper we leverage Theseus 10
solve differentiable
way. Since Theseus only supports solving non-linear least-square optimization problems without remove the dynamics constraints in the above optimization problem by manually rolling out the dynamics into the objective function. For example, with a planning horizon of H 2, we the above optimization problem into the following one:
constraints,
(0)
arg max Ro a.) + Re(de (zi, at), a.41) + Qa (ds (z1, at). ar41). a.12)
(5)
We set the values of H following the schedule as in TD-MPC, and we use the Levenberg- Marquardd algorithm in Theseus to solve the optimization problem. Following TD-MPC, we also learn a policy Ty in the latent space using the learned Q-value predictor Qa. and the output from the policy is used as the action initialization for solving
Let (0) be the solution of the above trajectory optimization problem, obtained using Theseus as described above. DiffTORI learned with the following objective, which jointly optimizes the encoder, latent dynamics model, latent reward model, and the Q-value predictor:
BL (DSTTTOR(
Σ X (CTD-MPo((: coCra(0:Ti)
(6)
ist
(Pa(0:T) = (o(s,,
5

---

argenin (o,
(o)
argminn & (o.o)
K
(a(o.a)
Pnjjechery ncttMnaaT.oo
Explicit Policy
Implicit Policy
Diffusion Policy
DiffTORI (Ours)
Figure 2: Overview of method Imitation Learning. DifTTORI (right) learns cost function via differentiable trajectory optimization and performs test-time optimization with it, which is different from prior work (left) that uses explicit policy diffusion without test-time optimization. Although implicit policy shares the same spirit as DifTTORI, we observe that the training procedure of DiffTORI using differentiable trajectory optimization leads to better performance compared the EBM approach used in prior work (2) which can suffer from training instability.
where Qo is the Q function learned via Bellman updates 40 which is used to compute the detemi- nistic policy gradient [41], and co is the weight for this loss term. Qo is learned in the original state space instead of the latent space provide accurate policy gradients. The key idea here is that can backpropagate through the policy gradient loss Lrg. which backpropagates through a()) and then through the differentiable trajectory optimization procedure of Equation to update 0.
4.3 Differentiable Trajectory Optimization applied to imitation learning
We also use DiffTORI for model-based imitation learning. A comparison of DiffTORI to other types of policy classes used in prior work is shown in Figure 2 In this approach, DiffTORI consists of an encoder he and a latent dynamics function dg, as before. However, in the setting of imitation learning, do not assume access to reward function R(s, a). Instead, generate actions by solving the following trajectory optimization problem:
(0)
arg max
fo(z), Eo'r)o(z at)
(7)
s.t. 2141 do az),
in which fo(=), a4) is a function the latent state and actions a that will optimize using the imitation learning loss, as described below. Similarly, We use to denote all learnable parameters to be optimized in DiftTORI, including the parameters of the encoder hg. the latent dynamics model dg. and the function So in the imitation learning setting.
In imitation learning, we assume access an expert dataset D = {(ss af)) of state-action pairs (s,, of) In the most basic form, the loss L for DifiTORI can be the mean square error between the the expert actions of and the actions a()) returned from solving 7
(ec Σ (a(() afll
The key idea here is that we can backpropagate through the imitation loss Lac, which backpropagates through (0) and then through the differentiable trajectory optimization procedure of Equation to update 0. This enables us to learn the function fo(2,, ar) used in the optimization Equation directly by optimizing the imitation loss (B). Because this loss is optimized through the trajectory optimization procedure (Equation we will learn function fo(z,, ax) such that optimizing Equation returns actions that match the expert actions.
Multimodal DifTORI:
The loss in Equation 8 will not be able to capture multi-modal action distributions in the expert demonstrations, To address this, we use Conditional Variational Auto- Encoder (CVAE) [42] as the policy architecture, which has the ability to capture multi-modal action distribution 43 The CVAE encoder encodes the state 8; and the expert action of into a
latent state vector zj. The key idea in our approach is that the decoder in CVAE takes the form of trajectory optimization algorithm, given by Equation 7 takes as input the sampled latent from the Gaussian Prior, and the state 8; and uses differentiable trajectory optimization to decode the

---

Average
Cartpnle- Swwnoup- STAINA
cheetah-nin
cup-catch
800
800
2.6
800
500
540
400
400
M
0.00
8.0
1.000
singer-spin
1,000
finger turn casy
1200
Fingen hard
1.000
hopper stand
820
800
800
500
230
200
200
200
8.0
0.2
3.4
3.6
5.6
1.0
8.0
0.2
044
0.6
3.8
1.0
8.0
0.2
0.4
0.6
5.6
8.0
0.2
0.4
0.6
0.8
1.
condduuum inqup
quadruped.oom
quad
roachereasy
800
600
600
600
500
4.0
400
400
400
9.0
8.0
8.0
8.
1.000
mwachen hard
1.000
Walbor-rum
1000
walken stand
1,000
waltorwall
8.0
800
800
800
6.00
000
600
200
200
200
8.0
02
241
3.6
O.B
1.0
0.7 D0B
2.0
3.0 8.0
0.2
6.4
D.6
5.B.
3.0
0.7
6.8
0.6 6.8
Figure Performance of DifTTOII in comparison prior state-of-the-art model-based and model-free RL algorithms, on 15 tasks from DeepMind control suite. DifTTORI achieves the best performance when averaged across all tasks. Results are averaged with seeds, and the shaded regions represent the 95% confidence interval.
action (0). Because this trajectory optimization is differentiable, we can backpropagate through learn the parameters for the encoder, dynamics dg, and the function used in Equation 7 See Appendix D for further details.
Action refinement: We note that DiffTORI provides a natural way to perform action refinement on top base policy. Given an action from any base policy, we can use this action as the initialization of the action variables for solving the trajectory optimization problem; the trajectory optimizer will iteratively refine this action initialization with respect to the optimization objective of Equation 7 In experiments, find DifTTORI always outperforms the base policies when using their actions as the initialization and other ways of performing action refinement, such as residual learning.
5 Experiments
5.1 Model-based Reinforcement Learning
We conduct experiments on 15 DeepMind Control suite tasks, which involve simulated locomotion and manipulation tasks, such making cheetah run swinging ball into cup. All tasks use image observations and the control policy does not have direct access to the underlying states.
commonly used off-policy model-free RL algorithm. DrQ-v2 [45], state-of-the-art model-free RL algorithm for image observations that adds data augmentation which DiffTORI builds on. All training details such as hyper-pant parameters, and pysudo-code can be found in Appendix B All experiments use NVIDIA 2080 Ti GPUs.
We compare to the following baselines: SAC 44
top of SAC TD-MPC 38 state-of-the- -f-----rtttt model-based RL algorithm,

---

Figure 3 shows the learning curves for all methods on all tasks. The top-left subplot shows the normahzed performance averaged across all 15 tasks, which is computed as the achieved return divided by the max return from any algorithm. As shown, DifTTORI (red curve) outperforms all compared baselines by a noticeable margin. On 14 out of the 15 tasks (except Quadruped-war walkk, DifTTORI achieves the highest performance among compared algorithms. We especially note that the performance of DiftTORI is much higher than TD-MPC, which DiffTORI builds on, showing the benefit of adding the policy gradient loss and directly differentiatingg through it to optimize the learned latent spaces, Although DifTTORI achieves higher sample efficiency, one limitation of DI is that it requires more wall-clock time for training, due to the need for solving and differentiating through the trajectory optimization process. We show detailed results on computational efficiency (return wall-clock time) of DiffTORI in Appendix A.1.2 We also perform ablation studies to examine how each loss term in (6 contributes to the final performance of DiffTORI in Figure 6 in Appendix A.1.3
Table I: Success rates (f) of DIRT DP3 and Residual + DP3 on 22 Meta World tasks. DifTTORI consistenly achieves higher or on-par success rates on all 22 tasks.
World (Medium)
Soccer
Push Wall
Peg Insert Side
Picking
Basketball
Box Close
Coffee Pull
DP3
3.
42.0+22
17.3
13.0+11
88.0-5.
22.0
65.023. 15.3
Residual DP3
33.0422
43.3
59,3 64.7 4-2.
13. 13.3
87.3131
67. 3,50
DITTORI (Ours) DP3
40.7
50.0:120
22.0
88.0
73.3
70.7
Metaworld (Medium)
Metaworld (Hard) Pick out of Hole
Coffee Push
Hammer
Sweep
Sweep Into
Assemble
Hand Insert
Pick Place
DP3 Residual + DP3
53.0
32.313.
74.713.
30.31112
68.711.
18.3121
55.0148 61.322 63.3-6.1
56.7125
19.3:42
31.3:1a
71.3au
45.7
64.7as1
19.7422
52.0120
DiITORI (Ours)+ DP3
60.7=50
37.3+8.1
90.7s8.1
45.3
74.0a40
24.7+1.2
61.3
Melaworld (lard)
Metaworid (Very
Hard)
Push
Push Back
Shelf Place
Disassemble
Stick Pull
Stick Push
Pick Place Wall
Average
DP3 Residual + DP3
21 Te
55.324.
27.7a29
31.0ess
33.0ass
94.3-2.1
38.3
48.8
20 20.0+s.5
35.32
35.3+42
35.3+1.
56 06.00
91.342.8
44.7=4.
49.8
DIITTORI (Ours) DP3
30.0=3.
64,7=8.1
42,0a8.0
40.7
59,3
94.0
44.7=2s
56.5
Table 2: Failure rates (4) of all methods on the Robomimic tasks. DiffTORI achieves the lowest failure rates on all tasks with diffusion policy as the base policy.
Residual +BC-RNN
DHITORI (Ours)
Residual + Diffussion
DiFTORI (Ours)
IBC
BC-RNN
BC-RNN
Diffusion
Diffusion
+ Diffussio
Square
0.961 8o
0.1823..0
0.16110.01
0.10220.2
0.32 0.12:0.05
0.12:0.02
0.08:20.01
Transport
1.001 oo
0.2810..0 03
0.261 0n
0.1720.02
0.0720.0s
0.92 oa
0.08
0.04 01
ToolHang
1.00# co
0.33+o0.04 04
0.28
0.18:1..0
0.10=0000
0.94=0.00 01
0.10=0.00
0.08-1..0
Push-T
0.89+0.01
0.30+0.02
0.28
0,25++0.2
0.09+0.00
0.92+9.01
0.09
0.09+0.01
Average
0.96
0.27
0.25
0.18
0.10
0.78
0.10
0,07
5.2 Imitation Learning
Below we show results of DiffTORI on 3 commonly used imitaiton learning benchmarks: Meta- World (22), RoboMimic [19), ManiSkill [20). and the comparison to state-of-the-art methods on these three benchmarks. We also compare to one closely related prior work 3 on one of their customized tasks in Appendix A.3
5.2.1 Meta World
MetaWorld [22] is a large-scale benchmark that includes 100 robotic manipulation tasks, and has been recently used for evaluating different imitation learning algorithms [46). The policy observation point clouds of the scene, and the action is the 3d translation of the robot end-effector. We test on 22 tasks with different levels of difficulties: Medium, Hard, and Very Hard (See Table for all the tasks). 10 demonstrations are used for all tasks [46). We compare DiffTORI with the following baselines: DP3 [46), 3D version of diffusion policy that achieves state-of-ln of-the-art performances on this benchmark, outperforming other algorithms such as the original diffusion policy 0 with 2d image inputs. Residual + DP3: Since DiffTORI refines the actions from a base pre-trained DP3 policy, we additionally compare to this baseline that also leverages the actions from a base pre-trained DP3 policy. Specifically, we leamn a residual policy top of the base pre- trained policy, which takes
8

---

Table 3: On Maniskill tasks, DiffTORI consistently achieves higher success rates (1) on all tasks.
OpenCabinet
OpenCabinet
PickClibe
Hong
Fieewate
Pour
Drawer
Dour
PuchChuir
MoveRackar
Average
BC
8.19
0.7259.04
0.76
0.25 0.3350.01
0.4740.03
0.354
0:12=0.0:
0:1040.0:
0.34
BC ressdual
0.214
0.75++..2
0.75+0.02
0.27+0.03
0.32+n.01
0.49+0.02
0.38+0.05
0.15-0.02
0.10+0.0:
DIITORL - BC
0.32 .2
0.220.00
0.83=0.05 us
0.29 0.17=0.02
153
0.45+0.02
0.20 oz
0.15
0.42
as imput the action from the base policy, and outputs a delta action that is added the base action. This is the most standard and simple way of doing residual learning. All training details such as hyper-parameters and pseudo-code can be found in Appendix B
Table presents the task success rates, averaged over 50 evaluation episodes, of all compared algorithms. As shown, DifTTORI consistently achieves higher (or on par) success rates than the other 2 compared baselines. The improvement in success rates is larger on tasks where the original DP3 policy struggles, e.g., a 15% improvement on the task of Shelf Place and Sweep Into; and as expected, when the base DP3 policy is already doing well on the task, there is not much room of improvement left for DifTTORI, e.g., on Basketball and Stick Push. The simple way of learning a residual policy on top of the DP3 policy does not always improve the performance of the base policy, and even leads to lower success rates. This demonstrates that DiffTORI is more effective way to leverage a pre-trained policy. On average, the success rates of DifiTORI is 7.7% higher than that of DP3, a substantial improvement with only 10 demonstrations.
5.2.2 Robomimic
Robomimic [9] is another commonly used benchmark designed to study imitation learning for robot manipulation. The benchmark encompasses a total of 5 tasks with two types of demonstrations: collected from proficient humans (PH) or mixture of proficient and profffiiiciet humans. We use the PH demonstrations, and evaluate on three of the most challenging tasks: Square, Transport, and ToolHang. We use image-based observations and the default velocity controller for all the tasks. In addition to Robomimic, we compare to another task, Push-T from the diffusion policy [I task set, to demonstrate that we can learn multimodal cost functions by using the CVAE training loss.
We compare to the following baselines: IBC (2 conditioned on both action and observation using the InfoNCE loss (47). BC-RNN [9:] A variant of BC that uses a Recurrent Neural Network (RNN) as the policy network to encode a history of observations. This is the best- best-performing baseline in the original Robomimic [99 paper. Residual + BC-RNN: We use a pretrained BC-RNN as the base policy, and learn a residual policy on top of it. The residual policy takes as input the action from the base policy, and outputs delta action which added to the base action. Diffusion Policy [I A policy that uses the diffusion model as the policy class. It refines noise into actions via learned gradient field. IBC + Diffusion: A version of IBC that uses the action from a pre-trained Diffusion Policy as the action initialization in the test-time optimization process. Residual Diffusion: Similar to Residual BC-RNN, but using a pre-trained Diffusion Policy as the base policy. For DiffTORI, we compare two variants of it: DifiTORI+ + BC-RNN and DiffTORI+ Diffusion Policy, which uses pre- pre-trained BC-RNN or a pre-trained diffusion policy as the base policy to generate the initialization action for solving the trajectory optimization problem. In Appendix A.2. we also present results of DifTTORI with zero initialization or random initialization, instead of initializing the action from a base policy.
An implicit policy that learns an energy function
The results are shown in Table We find that DifTTRI-Diifusion Policy achieves the lowest failure rates consistently across all tasks. Even though Diffusson Policy has almost saturated on these tasks with very low failure rates, DiffTORI can still further reduces it. Furthermore, irrespective of the base policy used - whether BC-RNN or Diffusson Policy - DiffTORI always brings noticcable improvement in the performance over the base policy. While learning a residual policy does lead to improvements upon the base policy, DifiTORI shows a significantly greater performance boost. In addition, by comparing DiLTTORiDiiffuiin Policy with IBC+Diffusion Policy, we find that using the same action initialization for IBC is considerably less effective than using the same action initialization in DiffTORI. In many tasks, even when the base Diffusion Policy already exhibits low failure rates, IBC+Diffusion Policy still results in poor performances, indicating the training objective used in IBC actually deteriorates the base actions.
We also show the benefit of using a CVAE architecture for DiffTORI, which enables DifFTORT to capture multimodal action distributions. With different latent samples from CVAE, we get

---

2D Reward Landscape
Reward Landscape
Goal Object
Objective Functonn
d2
1
009
Ginil
Oinit
0g
130
Action
Action
Figure 4: By using a CVAE, DifITORI can learn multimodal objectives functions via sampling different latent vectors from CVAE (right). By performing trajectory optimization with these two different objective functions, DifiTORI can generate multimodal actions (left).
different objective functions fa(z, a) and dynamics functions dz(z, a). allowing DiffTORI to generate the multimodal objective function learned by
different actions from the same state. DifITORI (right), and the resulting multimodal actions (left). The left subplot shows that when starting from the same action initialization asnit. with two different latent samples, DiffTORI optimizes Gjnit into two different actions, åj and 02 that move in distinct directions. The trajectory optimization procedure that iteratively updates the action is represented by dashed lines transitioning from faint to solid. From these two actions, two distinct trajectories are subsequently generated to push the T-shape object towards its goal. The middle and right subplots show the objective function landscapes for the different samples, as well as the initial action Oinits and the final optimized action di and d2- We note the two landscapes are distinct from each other with different optimal solutions, showing that DifTTORI can generate multimodal objective functions and thus capture multimodal action distributions. We note that the learned objective function f is not necessarily a "reward" function as those learned via inverse RL [48]. It is just learned "objective function"" such that optimizing it with trajectory optimization would yield actions that minimize the imitation learning loss with respect to the expert actions in the demonstration. We leave exploring the connections with inverse RL for future work.
5.2.3 ManiSkill
ManiSkill 20 21] is a benchmark for learning generalizable robotic manipulation skills with 2D & 3D visual input. It includes series of rigid body tasks and soft body tasks. We choose 9 tasks (4 soft body tasks and 5 rigid body tasks) from ManiSkill1 (20] and ManiSkil12 [21] and use 3D point cloud input for all the tasks. We use the end-effector frame as the observation frame 44) and use the PD controller with the end-effector delta pose as the action.
We build our method on top of the strongest imitation learning baseline in ManiSkil12 released by the authors, which is a Behavior Cloning (BC) policy with PointNet [50] as the encoder. Again, we also compare to BC+residual, which learns residual policy that takes input the action from the BC policy and outputs a delta correction. The results are shown in Table 3 As shown, DiffTORI + BC consistently achieves higher success rates than both baselines on all tasks, demonstrating the strong effectiveness of using diffcrentiable trajectory optimization as the policy class.
Conclusion and Discussion
We introduce DifTTORI that uses differentiable trajectory optimization generate the policy actions for deep reinforcement learning and imitation learning. The key is to utilize the recent progress in differentiable trajectory optimization compute the gradients of the loss with respect to the parameters of the cost and dynamics function of trajectory optimization, and learn them end-to-end. When applied to model-based reinforcement learning. DiffTORI addresses the "objective mismatch" issue of prior methods. We also test DiffTORI for imitation learning standard robotic manipulation task suites with high- nigh-dimensional sensory observations and compare it feed-forward policy classes as well as Energy-Based Models (EBM) and Diffusion. Across 15 model-based RI tasks and 35 imitation learning tasks with high-dimensional image and point cloud inputs, DiffTORI outperforms prior stte--of-the-ar methods.
10

---

References
[1] Cheng Chi, Siyuan Feng, Yilun Du, Zhenjia Xu, Eric Cousineau, Benjamin Burchfiel, and Shuran Song. Diffusson policy: Visuomotor policy learning via action diffusion. arXiv preprint arXiv:2303.04133 04.33...... 2023.
[2] Pete Florence, Corey Lynch, Andy Zeng, Oscar A Ramirez, Ayzaan Wahid, Laura Downs, Adrian Wong. Johnny Lee, Igor Mordatch, and Jonathan Tompson. Implicit behavioral cloning. In Conference on Robot Learning, pages 168-168, PMLR, 2022.
[3] Brandon Amos, Ivan Jimenez, Jacob Sacks, Byron Boots, and J Zico Kolter. Differentiable mpe for end-to-end planning and control. Advances in neural information processing systems, 31, 2018.
[4] Daniel Seita, Yufei Wang, Sarthak J Shetty, Edward Yao Li, Zackory Erickson, and David Held. Toolflownet: Robotic manipulation with tools via predicting tool flow from point clouds. In Conference on Robot Learning, pages 1038-1049. PMLR, 2023.
[5] Zhendong Wang, Jonathan J Hunt, and Mingyuan Zhou. Diffusion policies as an expressive policy class for offline reinforcement learning. arXiv preprint arXiv:2288.06193, 2022.
[6] Wanxin Jin, Zhaoran Wang, Zhuoran Yang, and Shaoshuai Mou. Pontryagin differentiable programming: An end-to-end learning and control framework. Advances in Neural Information Processing Systems, 33:7979-7992,, 2020.
[7] Xuesu Xiao, Tingnan Zhang, Krzysztof Choromanski, Edward Lee, Anthony Francis, Jake Varley, Stephen Tu, Sumeet Singh, Peng Xu, Fei Xia, et al. Learning model predictive controllers with real-time attention for real-world navigation. arXiv preprint arXiv:2209. 10780, 2022.
[8] Ming Xu, Timothy Molloy, and Stephen Gould. Revisiting implicit differentiation for learning problems in optimal control. arXiv preprint arXiv:2310. 14468, 2023.
[9] Wanxin Jin, Shaoshuai Mou, and George Pappas. Safe pontryagin differentiable programming. Advances in Neural Information Processing Systems, 34:16034-16050, 2021.
[10] Luis Pineda, Taosha Fan, Maurizio Monge, Shobha Venkataraman, Paloma Sodhi, Ricky TQ Chen, Joseph Ortiz, Daniel DeTone, Austin Wang, Stuart Anderson, et al. Theseus: A library for differentiablee nonlinear optimization. Advances in Neural Information Processing Systems, 35:3801-381 2022.
[11] Stephen Gould, Richard Hartley and Dylan Campbell. Deep declarative networks. IEEE Transactions on Pattern Analysis and Machine Intelligence, 44(8):3988-400 2021.
[12] Benoit Landry, Zachary Manchester, and Marco Pavone. A differentiable augmented lagrangian method for bilevel nonlinear optimization. arXiv preprint
2019.
[13] Stephen Tu, Roy Frostig, and et al. trajax, 2021.
[14] Jatan Shrestha, Simon Idoko, Basant Sharma, and Arun Kumar Singh. End-to-end learning of behavioural inputs for autonomous driving dense traffic. In 2023 IEEE/RS. International Conference on Intelligent Robots and Systems (IROS), pages 10020-10027. IEFF, 2023.
[15] Angel Romero, Yunlong Song, and Davide Scaramuzza. Actor-critic model predictive control. arXiv preprini arXiv:2306.00985. 2023.
[16] Nathan Lambert, Brandon Amos, Omry Yadan, and Roberto Calandra. Objective mismatch in bodel-based reinforcement learning. arXiv preprint arXiv:2002.0455 2020.
[17) Benjamin Eysenbach, Alexander Khazatsky, Sergey Levine, and Russ R Salakhutdinov. Mis- matched no more: Joint model-policy optimization for model-based Advances in Neural Information Processing Systems, 15:23230-22222 2022.
[18] Yuval Tassa, Yotam Doron, Alistair Muldal, Tom Erez, Yazhe Li, Diego de Las Casas, David Budden, Abbas Abdolmaleki, Josh Merel, Andrew Lefrancq, et al. Deepmind control suite. 2018.
arXiv preprint

---

[19] Ajay Mandlekar, Danfei Xu, Josiah Wong, Soroush Nasiriany, Chen Wang, Rohun Kulkarni, Li Fei-Fei, Silvio Savarese, Yuke Zhu, and Roberto Martín-Martin What matters in learning from offline human demonstrations for robot manipulation. arXiv preprint 20xx:2208.0039:0 2021.
[20] Tongzhou Mu, Zhan Ling, Fanbo Xiang, Derek Yang, Xuanlin Li, Stone Tao, Zhiao Huang, Zhiwei Jia, and Hao Su. Maniskill: Generalizable manipulation skill benchmark with large-scale demonstrations. arXiv preprint arXiv:2107. 14483, 2021.
[21] Jiayuan Gu, Fanbo Xiang, Xuanlin Li, Zhan Ling, Xiqiang Liu, Tongzhou Mu, Yihe Tang, Stone Tao, Xinyue Wei, Yunchao Yao, et al. Maniskill2: unified benchmark for generalizable manipulation skills. arXiv preprint arXiv:2302.0440
2023.
[22] Tianhe Yu, Deirdre Quillen, Zhanpeng He, Ryan Julian, Karol Hausman, Chelsea Finn, and Sergey Levine. Meta-world: A benchmark and evaluation for multi-task and meta reinforcement learning. In Conference on robot learning, pages (094-1100. PMLR, 2020.
[23] Brandon Amos and J Zico Kolter. Optnet: Differentiable optimization as a layer in neural networks. In International Conference Machine Learning, pages 136-145. PMLR, 2017.
[24] Akshay Agrawal, Brandon Amos, Shane Barratt, Stephen Boyd, Steven Diamond, and J Zico Kolter. Differentiable convex optimization layers. Advances in neural information processing systems, 32, 2019.
[25] Sheng Cheng, Minkyung Kim, Lin Song, Chengyu Yang, Yiquan Jin, Shenlong Wang, and Naira Hovakimyan. Difftune: Auto-tuning through auto-differentiation. IEEE Transactions Robotics, 2024.
[26] Sheng Cheng, Lin Song, Minkyung Kim, Shenlong Wang, and Naira Hovakimyan. Difftune O:: Hyperparameter--iee auto-tuning using auto-differentiation. In Learning for Dynamics and Control Conference, pages 170-183, PMLR, 2023.
[27] Jacob Sacks, Rwik Rana, Kevin Huang, Alex Spitzer, Guanya Shi, and Byron Boots. Deep model predictive optimization. In 2024 IEEE International Conference on Robotics and Automation (ICRA), pages 16945-16953 IEEE, 2024.
[28] Zhiyu Huang, Haochen Liu, Jingda Wu, and Chen Lv. Differentiable integrated motion prediction and planning with learnable cost function for autonomous driving. IEEE transactions on neural networks and learning systems, 2023.
[29] Christopher Dichl, Tobias Klosek, Martin Krueger, Nils Murzyn, Timo Osterburg, and Torsten Bertram. Energy-based potential games for joint motion forecasting and control. In Conference on Robot Learning. pages 31 1112-3141 3111. PMLR, 2023.
[30] Christopher Diehl, Tobias Klosek, Martin Krüger. Nils Murzyn, and Torsten Bertram. On connection between differential games, optimal control, and energy- based models for multi- agent interactions. arXiv preprint ar3iv: 16539, 2023.
[31] Mohak Bhardwaj, Byron Boots, and Mustafa Mukadam. Differentiable gaussian process motion planning. In 2020 FEEE international conference on robotics and automation (ICRA), pages 10598-10666 IEEE, 2020.
[32] Brent Yi, Michelle A Lee, Alina Kloss, Roberto Martín-Martín and Jeannette Bohg. Differ- entiable factor graph optimization for learning smoothers. In 2021 JEFE/RSI International Conference on Intelligent Robots and Systems (IROS), pages 1339-1345. IEEE, 2021.
[33] Raj Ghugare, Homanga Bharadhwaj, Benjamin Eysenbach, Sergey Levine, and Ruslan Salakhut- dinov. Simplifying model-based rl: leaming representations, latent-space models, and policies with one objective. arXiv preprint mrXiv:2209.084666 2022.
[34] Evgenii Nikishin, Romina Abachi, Rishabh Agarwal, and Pierre-Luc Bacon. Control-oriented model-based reinforcement learning with implicit differentiation. In Proceedings of the AAAI Conference Artificial Intelligence, volume 36, pages 7886-7894, 2022.
12

---

[35] Christopher M Bishop. Mixture density networks. 1994.
[36] Jimmy Wu, Xingyuan Sun, Andy Zeng, Shuran Song, Johnny Lee, Szymon Rusinkiewicz. and Thomas Funkhouser. Spatial action maps for mobile manipulation. arXiv preprint arX1v:2004.09141, 2020.
[37] Matthew Hausknecht and Peter Stone. Deep reinforcement learning in parameterized action space. arXiv preprint arXiv: 1014:1351,04143, 2015.
[38] Nicklas Hansen, Xiaolong Wang, and Hao Su. Temporal difference learning for model predictive control. arXiv preprint arXiv:2203.0499 04955, 2022.
[39] Grady Williams, Paul Drews, Brian Goldfain, James M Rehg. and Evangelos Theodorou. Aggressive driving with model predictive path integral control. In 2016 IEEE International Conference on Robotics and Automation (ICRA), pages 1433-1440. IEFF, 2016
[40]
Christopher JCH Watkins and Peter Dayan. Q-learning. Machine learning, 8:279-292, 1992.
[41] Timothy P Lillicrap. Jonathan Hunt, Alexander Pritzel, Nicolas Heess, Tom Erez, Yuval Tassa, David Silver, and Daan Wierstra. Continuous control with deep reinforcement learning. arXiv preprint arXiv:150.00290 2015.
[42] Kibyuk Sohn, Honglak Lee, and Xinchen Yan. Learning structured output representation using deep conditional generative models. Advances in neural information processing systems, 28, 2015.
[43] Tony 7. Zhao, Vikash Kumar, Sergey Levine, and Chelsea Finn. Learning fine-grained bimanual manipulation with low-cost hardware. arXiv preprint arXiv:2304. 13705, 2023.
[44] Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with stochastic actor, 2018.
[45] Denis Yarats, Rob Fergus, Alessandro Lazaric, and Lerrel Pinto. Mastering visual continuous control: Improved data-augmented reinforcement learning. arXiv preprint arXiv:2107.00645, 2021.
[46] Yanjie Ze, Gu Zhang, Kangning Zhang, Chenyuan Hu, Muhan Wang, and Huazhe Xu. 3d diffusion policy: Generalizable visuomotor policy learning via simple 3d representations, 2024.
[47] Aaron van den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive predictive coding. arXiv preprint vrXiv:1807.03777 2018.
[48] Andrew Y Ng, Stuart Russell, et al. Algorithms for inverse reinforcement learning. In Icml, volume 1, page 2, 2000.
[49] Minghua Liu, Xuanlin Li, Zhan Ling, Yangyan Li, and Hao Su. Frame mining: free lunch for learning robotic manipulation from 3d point clouds. arXiv preprint arXiv:22 2022.
[50] Charles R Qi, Hao Su, Kaichun Mo, and Leonidas Guibas. Pointnet: Deep learning on point the IEEE conference on computer
sets for 3d classification and segmentation. In Proceedings the vision and pattern recognition, pages 652-660, 2017.
[51] Grady Williams, Andrew Aldrich, and Evangelos Theodorou. Model predictive path integral control using covariance variable importance sampling. arXiv preprint arXiv:1509.0/ 149, 2015.
13

---

A Additional results
A.1 Model-based Reinforcement Learning
A.1.1 DiTTORI without policy gradient loss
In model- based reinforcement learning, the key distinctions between DiffTORI and TD-MPC [8 are: TD-MPC employs the Model Predictive Path Integral (MPPI (5T] in the planning stage, whereas we utilize trajectory optimization. 2) In addition to the original loss used in TD-MPC, we use an additional policy gradient loss and back-prop paoo-ppro through the differentiable trajectory optimization process to update the model parameters. Figure 5 shows that the improvement of DiffTORI over TD-MPC comes from the addition of the policy gradient loss, instead of purely changing MPPI to trajectory optimization. To be more specific, compare TD-MPC with DiftTORI (w/o backward), a variant DiffTORI that removes the policy gradient loss for updating the model parameters. The results indicate that TD-MPC and the DifiTORI (w/o backward) variant perform comparably, suggesting that using MPPI or trajectory optimization at test-time for action generation have similar performances. With the inclusion of the policy gradient loss, DifiTRllsiniffcaanly outperforms both TD-MPC and the DiffTORI (w/o backward) variant, demonstrating the efficacy of adding the policy gradient loss in DifTTORI.
Avuoe
Carpalle
Chentah
Cap
3023
3.0
9.0
Spin
Finger hard
Stand
800
100
200
200
200
0.5
3.2
Pordulam Swingua
Quadiured Run
Ouadhuped W
Readierr Easy
820
630
620
400
2.0
200
200
200
0.5
6.6
3.5
3.2
0.4
6/
8.6
9.2
0.2
5.6
Dil
3.0
0.2
1.0
Boacher Iland
Via Run
Waller Stand
Wallee Walk
1000
030
600
400
200
200
8.
0.2
5.0
03
64
25
3.0
TD-MPC
DIFITORI (w/o backward)
DIYTORII (Ours)
Figure 5: Performance of DiffTORI, in comparison to TD-MPC and DiffTORI (w/o backward) on 15 tasks from DeepMind control suite.
A.1.2 Computational efficiency DifTOORI
In addition comparing the sample efficiency of DifITORI to prior methods, we also compare the computational efficiency of DiffTORI versus TD-MPC on some of the environments. This is shown in Figure where the y-axis is the return, and the x-axis the wall-clock time (tested on NVIDIA RTX 2080 Ti GPU) used to train DiffTORI and TD-MPC for 1M environment steps. As shown, takes more wall-clock time for DiffTORI to finish the training. In terms of computational efficiency, the results are anvironment-dependen DiffTORI achieves better computational efficiency on reacher-hard and cup-catch. pendum-ssingee TD-MPC converges to a sub-optimal value in the early training stage and DifITORI outperforms it within 24 hours of training time. DifITORI
14

---

has similar computational efficiency slightly worse computational efficiency on cheetah-run and walker-stand compared TD-MPC. The gap is larger on hopper-stand. The major reason for DiffTORI to take longer time for training that solving and hack-propagating through the trajectory optimization problem in 4 is slower than doing MPPI as used in TD-MPC. As reference, to infer the action at one time step, it takes 0.052 second to use Theseus to solve and differentiate through the trajectory optimization problem in 4 and 0.0092 second for using MPPI in TD-MPC. However, we also want to note that the community is actively developing better and faster algorithms/software libraries for differentiable trajectory optimization, which could improve the computation efficiency of DiffTORI For example, in all our experiments, we used the default CPU-based solver in Theseus. Theseus also provides more advanced solver named BaSpaCho, which is a batched sparse Cholesky solver with GPU support. When we switch from the default CPU-based solver to BaSpaCho, the time cost of solving the trajectory optimization problem in (4) is reduced by 22% from 0.052 second 0.041 second. With better libraries/algen gorithhms in the future for differentiable trajectory optimization, we believe the computational efficiency of DifTTORI would further improve.
seieewswwnguppassasss reacher- easy, and finger-spin, and
Cheetah
Finger Easy
Easy
Epttade retum
Envtronment stops (x1e6)
w/o dynamic loss
w/o action inttalizatiion
w/o reward lass
DIfTOR (Ours]
Figure 6: Ablation study DifTTORI to examine the contribution of each loss terms towards the final performance, subset of 4 tasks. We find the reward prediction loss, action initialization, and dynamics prediction loss are all essential for DifiTORI to achieve good performance.
A.1.3 Ablation study on the loss terms
We also perform ablation studies to examine how each loss term in (6) contributes to the final performance of DifTTORI, shown in Figure We find that removing the reward prediction loss causes DifTTORI to completely fail. Removing the dynamics loss, or not using the action initialization from the learned policy Tę for solving the trajectory optimization, both lead to a decrease in the performance. These shows the necessity of using all the loss terms in DifTTORI for learning good latent space achieve strong performance.
A.2 Imitation Learning
A.2.1 DifFTORI with zero and random action initialization
We also present results DiffTORI with zero initialization or random initialization, where instead of initializing the action from base policy, the action is initialized to be O, randomly sampled from ((,01) 1), on RoboMimic and Maniskill.
The results on RoboMimic is shown in Table 4 zero or andomly-initialicn actions, possibly due to the convergence to bad local minima during nonlinear trajectory optimization without a good action initialization. We note this would not be drawback of applying DifTORRI in practice for imitation learning: one could always first learn base policy using any behavior cloning algorithm, and then use DiffTORI to further refine the actions.
We notice a drop in performance of DiffTORI with
The results Maniskill is shown in Table Again, if with DiffTORI, the performance drops to be similar to or slightly worse than vanilla BC. Therefore, think a good practice of using DifT for imitation learning would be always try provide with a good action initialization, e.g., by first training a BC policy and use its action as the initialization in DiffTORI.
zero random action initialization
A.2.2 Results of positional controller on RoboMimic
Note that for the three tasks in Table from Robomimic, we use the default velocity controller from Robomimic. We note the use of the velocity controller leads to small decline in the performance of
15

---

1000
1000
cap.catch
penculum swing.p
800
600
Return
400
200
1.0
2,0
3.0
4.0
0.0
2.0
8.0
Traising time (days)
cartpo e-swing. poserse
1000
reacher easy
onger
Ratur
600
Retum
Naturn
600
200
200
8.0
0.5
Triining time (days)
1.0
15
2.0
8.
2.0
Thiining Idayss)
4.0
0.0
0.0
Thaining (days) iepperstand
1400
cheetah-run
2000
welloe-siaree
1000
800
800
800
600
600
Nature
600
400
400
.00
200
200
200
6.0
2.0
4.0
6.0
8.0
0.5
2.5
5.0
7.5
10.0
12.5
8.0
2.0
6.0
Triiiin time (day5)
Traniing time (drys)
Training time (days)
TD-MPC
DiffTORI (Ours)
Figure 7: Return vs wall-clock time of DifTTORI and TD-MPC on some of the RL environments. The x-axis is the training time in days (24 hours), and the y-axis is the return. Both methods are trained for IM environments steps. The training takes a long time (a few days on some environments) because the policy observation is high- nigh-dimensional images.
Table 4: Failure rates (1) of DiffTORI and all other mehtods on the Robomimic tasks. DIRITORI achieves the best performances on all tasks when using diffusion policy as the base policy. zero or random initialization are used in DIRT the performance drops, possibly due to the convergence bad local minima during nonlinear trajectory optimization without good action initialization.
IBC
BC-RNN
Residal +BC-RNK
DITTORI (Ours) BC-RNN
DirTusion
IBC
Diifusion
+ Diffission
Residual
DiRTORR (Ous)
Diffusion
DiITORI (Oun))
zero inil.
DUTTORI (Ours)
+ tandom inil.
Square
0.96--0.00
0.18:20.00
0.16:10.01
0.104:0:22
0.124-0.03
0.32::0005
0.12:80.02
0.08 0.00
0.16:1:0022
0.20:1:0:00
Transpott
10000.00
0:282003
0.26
0.77:022
0,07:0044
0.2220003
0.08::0.01
0.0400001
0.58:20.01
0.44::004
ToolHanng
1000-0000
0.33:00.4
022820005
0,18-:000
0,101:0000
0944-001
0.10-00.00
0,08 0.00
1.00:-0.00
1:0++00000
P
0.891-0.01
0.30=00:2
0.28:80.02
0.254:0.02
0,091:0:00
0.92-100001
0.38:1-0044
the Diffusion Policy compared to its performance in the original paper where a positional controller used. The Push-T task still uses the default position controller in the diffusion policy paper. Below we evaluate the performance of DiffTORI and Diffusion Policy with the positional controller.
In the original Diffusion Policy [ paper, it was observed that the of positional controllers yielded superior results for Diffusson Policy compared to the default velocity controller on Robomimic [19] tasks. We evaluate Diffusion Policy, which is the strongest baseline, and DifTTORI the most difficult three tasks with ph (proficient-human demonstration) and mh demonstrations using positional controller. The results with the positional controller are presented in Table 6 Diffusion Policy already achieves nearly the maximal possible performance on most tasks with the positional controller. DifiTORI, however, is able achieve similar even higher performances on most of these tasks.
human demonstration)
A.2.3 Ablation on planning horizon H
Additionally, we do ablation experiments on the planning horizon H for imitation learning, with the results presented in Table We observe that simply increasing the planning horizon H in imitation learning does not necessarily enhance performance. As the horizon increases from H 1 to H 3.
16

---

Table 5: Success rates (1) of all methods on the Maniskill benchmark. DifTTORI consistently outperforms both baselines on all tasks with action initialization from the BC policy. If zero or random initialization are used in DIfT the performance drops, possibly due to the convergence to bad local minima during nonlinear trajectory optimization without good action initialization.
OpenCabinet Drawer
OpenCahinet
PickCabe
Hong
Excavate
Pour
PushChair
MoveBRcket
0,194-003
0.74-004
0.2440022
0.25++0.02
0.134-0.01
0.474-0.03
0.354-0.4
0.124-0.01
0.10-0001
BC + resicual
07511002
0.531002
0,2700033
0.12:0..1
0.91002
0.3610.3
0.15
0.10::0011
DIFT BL
0,32-0002
4024000
0-8540033
0.29
0,17 +0,0
0.53+-0002
0,45
0.30--0.02
0.15 0.10
DIETORI (Ouns
0,3010003
07762003
072200022
0:25:0001
0.0420000
0.501:004
0.3420004
0.04:0001
0.06
+ aroo init
DIrTORI (Oirs)
0.18:0022
0:6:::00
066:00011
0.19:00:4
0,04::0.00
0.391:0.04
0.50::0.02
0.00::0000
0.05
Table 6: Failure rates (L) of DiffTORI and Diffusion Policy using Positional Controllers on
Robomimic Tasks.
Square (ph)
Square (mh)
Transport (ph)
Transport (mh)
ToolHang (ph)
Diffusion
0.02:10.01 0.02++0.01
0.03+-0.02 0.04++0.02
0.00 0.00
0.12++0.02 0.09++0.01
0.05+0.02 0.04
DiffTORI Diffusion
the performance remains nearly the same; however, when H increase to 5, we observe slight decline in the performance.
Table 7: Failure rates (J) of different planning horizon H for DiffTORI on RoboMimic tasks.
Square (ph)
Transport (ph)
ToolHang (ph)
Push-T
H 1
0.08
0.04:10.01
0.08
0.09
H 3
0.08
0.06++0.02
0.08
0.12+10.02
I 5
0.09
0.06:1:001
0.10++0.00
0.12+00.01
A.3 Comparison to prior work Amos et al. 3
because we target different experiments. These related works all conduct experiments on customized tasks with ground-truth low-level states. In contrast, we test our method on standard RL and robotic imitation learning benchmarks, with high- dimensional sensory observations like images and point clouds. As these prior works have not been demonstrated on high-dimensional observations or more complex tasks, we originally compared to more recent state-of- of-e--of-the-art methods on these benchmarks, e.g.. 3D Diffusion Policy (4)..
In our main paper, we did not compare to 3
We have now included a comparison with Amos et al. in one of their tasks (pendulum swing-up with ground-truth low-level states) under imitation learning settings. Unlike Amos et al. (3), who assumes known dynamics and reward structures and only learns 10 parameters, our method uses neural networks represent both dynamics and reward functions without such assumptions. The metric is the cost of the learned policy. As in Amos et al., we test in two settings, pendulum without damping and with damping. Following Amos al., their method does not model the damping effect in the assumed dynamics, so the ground-truth dynamics model is not realizable in the damping case. We also compared to an additional baseline in Amos al., which uses a LSTM to predict the expert action. The results in Table 8 show our method performs slightly worse in the no damping case but noticeably better in the damping case. This because Amos al. assumes correct dynamics in the no damping case and learns only 10 unknown parameters, whereas the assumed dynamics structure is incorrect in the damping case; we use fully-connected neural networks to represent the dynamics function, avoiding such assumptions. It is generally difficult to know the exact correct dynamics function structure, especially for tasks with complex dynamics (e g., with contacts) and high-dimensional observations (images and point clouds).
17

---

Table 8: Cost of different algorithms on the Pendulum swingup tasks from Amos al. As in Amos al., we test in two settings, pendulum without damping and with damping. Lower cost means the better performance. DiffTORI performs slightly worse the no damping case but noticeably better in the damping case.
Expert Policy
Amos
LSTM policy
(inTOR[(ours)
Pendulum w/o damping Pendulum with dampling
13.126
13.576 + 0.012
15,962 + 0.164
14.603 + 0.190
10.132
14.874 + 0.600
12.098 + 0.031
10.644 + 0.029
B Implementation Details
In this section, we describe the implementation details of DiftTORI for the model-based RL ex- periments. For the imitation learning part, the code structure very similar to this model-based RL implementation. For more detailed information, please refer to the code we will release upon acceptance of the paper. We implement DiffTORI on top the open-source implementation of TD-MPC (38) from the authors. Below we show the pscudo-code of the training function in DORL
def train():
Trasning code u.e
for step in range(total_steps) obs env reset()
Differentiable trajectory optimization and update model
action, info agent.plan_theeeus. caheseus_update(on Env step obs, zeward, done, # collect data in buffer and update model (TD-MPC loss)
env env. stwppactiooonn.... (cep(action.cp()) .ummy ())
replay_ buffer (oba, action, reward, done) agent. update (replayybrell
Then, we demonstrate how the policy gradient loss is computed by differentiablee trajectory optimiza- tion in DifITORI with PyTorch-like pscudocode:
def *ie
Differentiable trajectory optimization and update model using policy gradient loss.
h, R, Q, d: model components.
c0: loss coefficients. u.e
import theseus as th
Encode first observation self mobe)..(ob)))
Get initialization action from pi
init_actions = self model. (z)
Theseus variable actions th. Vector (tensormactions, name="actions")
obs th. Variable (obs, name="obs")
Cost Function
Objective
cost fost_function th. tb..AuttoifCCCothutttiicc [obs], [action] vallu_cost__0 objective th. Ob)eetieveee add(costtfunction)
Trajectory optomization optimizer
casseus_opting th. TheseusLayer (th_optimizer)
Theseus layer forward
18

---

theseus_inputs = ("actions"; init_actions, "obs": obs}
inlated__iputs, info theseus_ wttp:aaoooiiio..odaddaaaaadaeaaa wtw.arrddceeessuuuiiprtr
apdated_actions updated_inputs actions
Update model using policy gradient losss
a_loss - torch. wwt.mmmiiieseeomoooooo model. Q_s(obs, (pdated_actionss)))))1+++ a a_loss. backward
aorim_asstepl
-For model-based reinforcement learning, We provide the network details for the added networks we used upon TD-MPC, which are the twin Q networks Q₂ learned in the original state space for computing the deterministic policy gradient.
(Qsl): Sequential
(0): Sinear(in_5) (ea_ureuss out_features =256)
(1): HLU(alpha=1.0)
(2): Linear(in_fetturet
out_features==))
(3): Linear (in_fetinnttutrrss out_features=511
(4): LayerNorm((512.)) elementwise__fffneeTTuu
(5): T)
(6): Linear (inear(ii_featurr (inattuen_, out_features=512)
(7): ELU(alpha=1.0)
(8): Linear
out_features=!))
(Q s2): Sequential
(0): Linear(in_featurrssSS out out_feature==255
(1): HLU(alpha=1.0)
(2): Lin_f
out_features=2))
(3): Linear(in_eeatures
out_features==11]
(4): DayerNorm((512.)) elementwise_dffinc=TTue
(5): Ta
(6): Linefarin_eattures
out_features=512]
(7): ELU(alpha=1.00
(8): Linear(in (in__catures =512, out out_fcatures=!)
For Imitation Learning. The default network details are follows. Note that for Robomimic [19 and Push-T tasks, we use the RNN-encoder from Robomimic; for ManiSkill 20 21] tasks, we use the PointNet encoder from ManiSkill2 [2]].
(ho): Sequential
(0): Linear(in (in_etturrrss
out_features=2355
(1): ELU((apha=1.00
(2):
out_fcatures=255 =256)
(3): ELU(alpha=1.00
(4): Lincar(nn out_fcatures=2s))
(ha): Identity
(hl) ): Sequential
(0): (infettin_fatturee--Z5 out_features-22
(1): ELU(alpha=1.0)
(2): Linear(in_ectturs=225 out_features=256)
(3): FLU(alpha=1.00)
(4): Linear(in out_features=128)
(R): Sequentiall
(0): Lineartin_fcctuures-::::44 out out_features=512)
(1): ELU(alpha=1.0)
(2): Linear(i__eattures5511 out out_features=512)
(3): ELU(alpha=1.0))
(4): Linea out out_features=!))
(d): Sequential
(0): (inaarninnfattuures--eaa out_features=512)
(1): #LU((Ipha=1.0)
(2): (inefr(in_eattures-55 out_features=:12)
19

---

(3): ELU(alpha=1.0)
(4): Linear iin_fttuures=seeeeeiiiiiirrseess =512, out oul_features=22464)
Hyperparameters used for DifTTORI for both model-based RL and imitation learning are shown in Tab 9 In model-basee RL., the parameters with TD-MPC 38 whenever possible.
Table Hyperparameters used in DifTTORI.
Hyperparameter
Value
Model-based RL
Max planning iterations
100 (50)
Planning step size
le-4 (5e-3)
Discount factor
0.99
Action loss coefficient (c0)
optimizer
Adam(3 = 0.9, 32 = 0.999)
Gradient Norm
10
Planning horizon schedule
1 5 (25k steps)
Batch size
256
Latent dimension
50
Sampling technique
PER(a 0.6, /3 0.4)
Learning rate
1e-3
Imitation Learning
Max planning iterations
100
Planning step size
le-4
Planning horizon schedule
Latent dimension
Posterior Gaussian dimension
64
KL coefficien
Learning rate
3e-4
Learning rate (MetaWorld)
3e-3
GMM Num Modes
5
RNN Seq Len
16
RNN Hidden Dim
1000
Point Cloud Sampled Points (ManiSkill)
1200
Point Cloud Sampled Points (MelaWorld)
512
C Environment Details
For model-based reinforcement learning evaluation, we use 15 visual continuous control tasks from Deepmind Control Suite (DMC). For imitation learning, we use tasks (detailed information can be found in Table from Robomimic [19], IBC (2), ManiSkillp 200, and ManiSkil12 21
We visualize the keyframes of the imitation learning tasks in Fig 8 and Fig
D More implementation details on using CVAE for imitation learning
We provide more details on how we instantiate DI with CVAE in imitation learning, in which the goal is to reconstruct the expert actions conditioned on the state. The CVAE encoder is composed of three networks: the first network is a state encoder ha that encodes the state into latent feature vector z = hg (s;), which is the conditional information in our case. The second an action encoder hg that encodes the expert action into latent feature vector - (a()) ). The last is fusing encoder (2", 2²) that takes input the concatenation of the state and action latent features, and outputs the mean f and variance o² of the posterior Gaussian distribution M o²). During training, the final latent state for state 8; used (7) is the concatenation of sampled vector ž from the posterior Gaussian distribution (- o2), and the latent state feature vector z: [2, 2"), NV-uu,, o2).
20

---

Table 10: Imitation Learning Tasks Summary.
Tank
Type
Object
Deseription
Square
Ralveimmic
Img
Rigid
200
Pick a squane not and place on mdd Transfer a hammer from concainer a bis
Transport
Robomimicc
Img
14
Rigid
200
Toollfang
Robomimic
Rigid
200
assemble frame consisting of base und book
Push-T
Rigid
T-shaped object spocified position Open a specific
OperCabinetDrawer
Pount
Rigid
3000bb
cabinet
OpenCahinelDoor
ManiSkillI
Point Cloud
Rigid
300
Open a specific doce of the cubinct
PhstChairr
ManiSkilt
Pount Cloud
Rigid
Padi the swivel chair to the sarget position
MoveBacket
ManiSkillt
Point Cloud
22
Rigid
300bby
Move a bucket and lift it cato pladfom
PickCabe
ManiSkil12
Point Cloud
Rigid
1000
Pick up cuhe and move a goal position
MuniSkili2
Cloud
Sont
Fill
bucket the turget beaker
Hang
ManiSkiliz
Cloud
Hang a
target cod
Facavate
ManiSkill2
Cloud
sevoun clay tanget heiggt
Pour
ManiSkil22
Point Cloud
Soct
200
Pou: liguid from bottle into boskee
Soccer
MetaWorkd (metium)
Point Cloud
Rigid
10
Kick soccer into the goal
Push Wall
MenWWork (medium)
Print Cloud
Rigid
10
Rypess wall and push puck goal
MataWorkd (medium)
Cloud
peg sideways
picking
MetaWorkd (metium)
Rigid
Grasp pack
place another
Badketholl
MenWWrkd (metium)
Print Cloud
Rigi
10
Dunk bockethall into the backet
Box close
MotaWordd (metium)
Point Cloud
Rigid
10
Grasp the cover and close the box with
Coffee pull
MetaWorkd (methum)
Point Cloud
Righd
10
Pull mug from coffee machine
Cottie push
MetaWorkd (medium)
Point Cloud
Rigid
10
Push mugg under coffee machine
Hammer
MetuWarkd (medium)
Point Cloud
Rigid
10
Hammer
wall
Sweep
MetaWorkd (methum)
Rigid
Sweep
Sweep into
MotaWorkd (nedium)
Pout Cloud
Rigid
Sweep pock
Assemble
MotWWordd (hard)
Point Cloud
Rigid
10
Pick up zut and place onso a peg Insert the gripper into a hele
Hand ansert
MonWWork (lan))
Point Cloud
Rigid
10
Pick out of hole
MotoWorkd (lar))
Posnt Cloud
Rigid
10
Pick puek from a hole
Phck place
MetuWarid
Point Cloud
10
and place goal
Push
MenaWorkl
the
Push back
MotaWorid
Posit Cloud
Rigid
puck to goal
Shelf place
MenWWordd (very hard)
Point Cloud
Rigid
10
pick and place puck onto a shelf pick a nut out the a pag
Disssemble
Men-Workd (very hard)
Print Cloud
Rigid
10
Stick pull
MetaWorid (very hand)
Point Cloud
Rigid
10
Grasp stick and pull box with the stick Grasp slick and push bux using the stick puck.hypess
Stick push
MetaWorkd (very hard)
Point Cloud
Rigid
10
place wall
MenaWorkd (very
Rigid
plase puck
The latent state will then be used as input for the decoder, which consists of the reward function Re- and the dynamics function dg. Trajectory optimization is performed with the reward and dynamics function in the decoder to solve 7 to generate the reconstructed action a" (0; st). The loss for training the CVAE the evidence lower bound (ELBO) on the demonstration data:
CLL (1) (0) = Σ |aa(0 s,) - a KL(N((H.,
where (L(PIIQ) denotes the KL. divergence between distributions P and Q. At test time, only the decoder of the CVAE is used for generating the actions. Given state s, the latent state is the concatenation of the encoded latent state feature =", and a sampled vector 2 from the prior distribution (O 1).
21

---

Square
Trensport
Tool Hang
Push T
OpenCabinet Door
OpenCabinng Draver
Push Chair
Move Bucket
Pick Cubbe
FiI
Hang
Excavate
Pour
Figure 8: Visualization of the tasks for imitation learning in RoboMimic and ManiSkill.
22

---

Assembly
insert
Peg side
hole
Bin picking
Pick place
Basketball
Box close
Push
back
Coffee pull
Sheff place
Coffee push
Disassemble
Hammer
Slick pull
Sweep
Stick push
Sweep into
Pick place
wall
Figure Visualization of the tasks for imitation learning Metaworld.
23

---

NeurIPS Paper Checklist
1. Claims
Question:
Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer:[ Yes].
Justification: The main claims in the abstract and introduction are all supported by our experimental results.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation as long as is clear that these goals attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification: We discussed that one limitation of DiffTORI is that it requires longer time to train Section 5.1, with detailed results on its computational efficiency in Figure 5 in Appendix A.1.I.
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification. asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect the scope of the claims made, e.g., if the approach was only tested on a few datasets or with few runs. In general, empirical results often depend implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full of assumptions and complete (and correct) proof?
24

---

Answer:
[NA].
Justification:
This is not a theoretical paper and we did not include any theoretical results.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Does the paper fully disclose all information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: We included all the training details, such as the hyper-parameters, training settings, environment details, and pseudo-code of our experiments in the main paper and appendix B & C, in order for others to reproduce our experiments.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, other means that appropriate to the research performed.
also be provided via detailed
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reprodueibility may be tricky in some cases, in which authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers have some path reproducing or verifying the results.
5. Open access to data and code
25

---

Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Question:
Answer: [No]
Justification:
We still cleaning code and will release it soon.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.cr cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https guwdes/ddeeeee CooeeSumeeeeeeeeeeeeee ssi onyllicy) for more details.
/ //nips.
The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data,
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental Setting/Details
Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Question:
Answer: [Yes]
Justification: We have included all the training and test details in the main paper and Appendix B.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment Statistical Significance
Does the report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification: For every experiment introduced, run with multiple random seeds and reported both mean and std (or confidence interval) averaged over these multiple random seeds.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by emor bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
26

---

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
If error bars reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: As mentioned in the main paper, used NVIDIA RTX 2080Ti for the experiments, We suggest running our experiments with GPU of at least 12G VRAM.
Guidelines:
The answer NA means that the paper does include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental rums as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics htttss///eeerrpp..cc/p/b/////// ctwpeeeeeeiicciceeeeeeeenee
Answer: [Yes]
Justification:
We are strictly following NeurIPS Code of Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurlPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. there a special consid- eration due to laws regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: There is socictal impact of this work. Our paper proposes policy representation for reinforcement learning and imitation learning, and the authors do not feel there would be any direct social impact of the works performed.
Guidelines:
27

---

The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No. they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g.. gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
The paper poses no such risks.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification:
We cited all the creators or original owners of assets mentioned in the paper.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset used and, possible, include a URL
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
28

---

If assets are released, the license, copyright information, and terms of use in the thcode. has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
package should be provided. For popular datasets, paperswi
latasets.paperssi
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information not available online, the authors are encouraged to reach out to the asset's s creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: The paper does not release new assets.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/moode as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: This paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country institution) were obtained?
Answer: [NA]
Justification: This paper does involve crowdsoureing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
29

---

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
30


Reviews:
[{"review_id": "nvd16FGY96", "reviewer": "Reviewer_kMzm", "summary": "The paper presents a method that uses Differentiable Trajectory Optimization as a policy representation. The proposed method extends the work of Temporal difference learning for model predictive control  (TD-MPC) by incorporating a policy-gradient loss for which analytical backpropagation is possible thanks to the differentiable properties of the used trajectory optimization scheme.\n\nThe paper evaluates the effectiveness of trajectory optimization as a policy representation for both model-based RL and Imitation Learning, benchmarking the approach in the DeepMind control tasks, the MetaWorld benchmark, the Robomimic benchmark, and the ManiSkill benchmark. Furthermore, the work presents comparisons against other policy representations like traditional feed-forward policies, Energy-based methods (EBMs), and Diffusion Policies.", "questions": "- Is there any specific reason why the Levenberg-Marquardt solver was used? Using a second-order solver might lead to fewer iterations to get to a solution, but the time per iteration might also be higher. A simpler solver (SGD) could do the job and be faster, improving the overall training time and inference speed.\n\n- Are there any insights on why increasing the time horizon for the trajectory-optimization policy leads to the slight decrease in performance reported in Appendix A2.3? Typically, a longer time horizon leads to trajectories of better quality. Does the trajectory optimization problem reach convergence for the longer time horizons that were tested?", "limitations": "Yes. Limitations of the method are mentioned and I agree with the author's claims that there are no direct potential societal impacts.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper is well-written, the work is interesting and it presents a simple yet effective extension of TD-MPC.\n\n- The method enables the use of trajectory optimization using high-dimensional observations and it is also able to deal with multimodality in the solution space.\n\n- The method is evaluated thoroughly in different setups (Model-based RL, IL) and against other classes of policy representations, outperforming most baselines in terms of sample efficiency and final reward.", "weaknesses": "[Minor]\n- The main weakness of the approach is the high computational cost of solving trajectory optimization at test time. The appendix reports 0.052 seconds to infer the action for one timestep (20 Hz), despite the fact that the time horizon is considerably short (1 to 3 steps). Such inference speed might hinder the deployment in certain real-world scenarios.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We want to extend our heartfelt gratitude for taking the time to review our paper. Thank you for all valuable comments and suggestions on improving the quality of the paper. Below we respond to each of your comments in detail.\n\n**Q: [minor] The main weakness of the approach is the high computational cost of solving trajectory optimization at test time. The appendix reports 0.052 seconds to infer the action for one timestep (20 Hz), despite the fact that the time horizon is considerably short (1 to 3 steps). Such inference speed might hinder the deployment in certain real-world scenarios.**\n\nWe thank the reviewer for bringing up this issue. We have updated section 5 of the paper to include a more detailed discussion on the inference speed of our method for real-world applications: \n\n“Our current inference speed is 0.052 seconds, which is a control frequency of 20 Hz. We note that such inference speed is comparable or higher to other deep robot learning algorithms that take high-dimensional image or point cloud observations as inputs, e.g., Diffusion Policy [1] reports a control frequency of 10 Hz, and PerAct [2] reports a control frequency of 2.23 Hz. As these methods are shown to be able to be deployed with real-world robots for many tasks, we believe our method’s inference speed of 20 Hz would work well for most real-world robot tasks as well. ”\n\nWe would also like to clarify that since we also learn a value function that predicts the future accumulated rewards and use it during the planning process, the policy is able to reason more than 3 steps into the future at inference time. \n\n**Q: Is there any specific reason why the Levenberg-Marquardt solver was used? Using a second-order solver might lead to fewer iterations to get to a solution, but the time per iteration might also be higher. A simpler solver (SGD) could do the job and be faster, improving the overall training time and inference speed.**\n\nWe thank the reviewer for this insightful question. We use the Theseus [3] library for differentiating through the trajectory optimization algorithms, and the solvers supported by Theseus include second-order solvers like Gauss-Newton, Levenberg–Marquardt, and Dogleg, and linear solvers such as CHOLMOD. We choose Levenberg-Marquardt as we find it to perform better in early experiments. We didn’t use SGD since it is not supported by the Theseus library, possibly due to its slower convergence rate when dealing with highly non-linear problems. Although implementing SGD to solve the trajectory optimization problem itself is not difficult, robustly differentiating through it may not be trivial and is out of the scope of the current paper. We look forward to trying more solvers, such as SGD, in future work. \n\n**Q: Are there any insights on why increasing the time horizon for the trajectory-optimization policy leads to the slight decrease in performance reported in Appendix A2.3? Typically, a longer time horizon leads to trajectories of better quality. Does the trajectory optimization problem reach convergence for the longer time horizons that were tested?**\n\nWe thank the reviewer for this interesting question. We do notice a slight performance drop with longer prediction horizons. The reason could be as follows: since there will always be errors in the learned dynamics function and the reward/value function, there is a tradeoff between the errors in the dynamics model and the errors in the reward/value function when choosing the prediction horizon. With a longer prediction horizon, compounding errors in the learned dynamics model will dominate, whereas with a shorter prediction horizon we expect errors in the learned value function to be more prevalent [4]. The optimal value for this parameter is highly likely to be application-dependent, but as shown in Appendix A2.3, our method can demonstrate robustness to different prediction horizons. When solving the trajectory optimization problem, we run the Levenberg-Marquardt solver for 100 iterations and it has reached convergence within 100 iterations. \n\n[1] Cheng Chi, Siyuan Feng, Yilun Du, Zhenjia Xu, Eric Cousineau, Benjamin Burchfiel, Shuran Song, “Diffusion Policy: Visuomotor Policy Learning via Action Diffusion”, RSS 2023  \n[2] Mohit Shridhar, Lucas Manuelli, Dieter Fox, “Perceiver-Actor: A Multi-Task Transformer for Robotic Manipulation”, CoRL 2022  \n[3] Luis Pineda, Taosha Fan, Maurizio Monge, Shobha Venkataraman, Paloma Sodhi, Ricky T. Q. Chen, Joseph Ortiz, Daniel DeTone, Austin Wang, Stuart Anderson, Jing Dong, Brandon Amos, Mustafa Mukadam, “Theseus: A Library for Differentiable Nonlinear Optimization”, NeurIPS 2022  \n[4] Harshit Sikchi, Wenxuan Zhou, David Held, “Learning Off-Policy with Online Planning”, CoRL 2021", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "GB94LsHTFZ", "reviewer": "Reviewer_fycy", "summary": "The paper introduces DiffTOP, a novel policy class for reinforcement learning (RL) and imitation learning (IL) that utilizes differentiable trajectory optimization to generate policy actions. DiffTOP leverages recent advancements in differentiable trajectory optimization, allowing end-to-end learning of cost and dynamics functions through gradient computation. The approach addresses the \"objective mismatch\" problem in model-based RL by optimizing dynamics and reward models to directly maximize task performance. For imitation learning, DiffTOP optimizes actions with a learned cost function at test time, outperforming previous methods.\n\nThe authors benchmark DiffTOP on 15 model-based RL tasks and 35 imitation learning tasks with high-dimensional inputs like images and point clouds. The results demonstrate that DiffTOP surpasses prior state-of-the-art methods in both domains. The paper includes analysis and ablation studies to provide insights into DiffTOP's learning procedure and performance gains.", "questions": "- For ManiSkill tasks, have the authors tried more advanced BC baselines (e.g., Diffusion Policy)? It seems the performance of baselines is not very strong.", "limitations": "The authors have discussed a few limitations in the paper.", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 4, "strengths": "This paper has several strengths:\n\n- The paper presents a robust and technically rigorous study with extensive experiments and ablation studies across a wide range of challenging environments. The proposed method DiffTOP achieves superior results in both RL and IL tasks with high-dimensional sensory observations.\n- The approach effectively tackles the important \"objective mismatch\" problem inherent in model-based reinforcement learning.\n- The ability to compute policy gradients directly with respect to the parameters describing the observation and transition model is a significant advancement, eliminating the need for sample-based estimates.\n- DiffTOP alleviates the model mismatch problem by learning the model concurrently with optimization, similar to TD-MPC.", "weaknesses": "While the paper has many strengths, there are also some potential weaknesses or areas that could be improved:\n\n- The paper attempts to condense a large amount of information into limited space, which may affect readability and clarity, particularly for readers less familiar with TD-MPC.\n- The trajectory optimization solver used (Theseus) does not support constraint optimization, requiring manual unrolling of dynamics instead of presenting it as a constraint to the optimizer.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We want to extend our heartfelt gratitude for taking the time to review our paper. Thank you for all valuable comments and suggestions on improving the quality of the paper. Below we respond to each of your comments in detail.\n\n**Q: The paper attempts to condense a large amount of information into limited space, which may affect readability and clarity, particularly for readers less familiar with TD-MPC.**\n\nWe thank the reviewer for bringing this to our attention. Indeed, the model-based RL part of DiffTOP builds upon TD-MPC, thus requiring some knowledge of TD-MPC for understanding the paper. Following the reviewer’s suggestion, we have updated section 3.2 with more explanations of TD-MPC to add more buildup for introducing our algorithm, and for a better understanding of the paper. We have also updated our paper to include more background information and details on TD-MPC in the appendix for readers less familiar with it.\n\n**Q: The trajectory optimization solver used (Theseus) does not support constraint optimization, requiring manual unrolling of dynamics instead of presenting it as a constraint to the optimizer.**\n\nWe thank the reviewer for this feedback. Indeed, Theseus does not support constrained optimization at the time of our paper submission, and thus we have to unroll the dynamics when solving the trajectory optimization problem. The original Theseus paper has the following discussion about constrained optimization in their limitation section: “The nonlinear solvers we currently support apply constraints in a soft manner (i.e., using weighted costs). Hard constraints can be handled with methods like augmented Lagrangian or sequential quadratic programs [99, 100], and differentiating through them are active research topics”. Based on this, it seems that the support of constrained optimization might be added in the future. We leave integrating DiffTOP with constrained optimization as important future work. \n\n**Q: For ManiSkill tasks, have the authors tried more advanced BC baselines (e.g., Diffusion Policy)? It seems the performance of baselines is not very strong.**\n\nWe thank the reviewer for this question. ManiSkill tasks all use point cloud as the policy inputs; the original Diffusion Policy paper only tested their method with image inputs, and has not been implemented and tested to work with point cloud inputs. The baseline we compared to in the ManiSkill tasks was the best method that was introduced by the ManiSkill authors, and we have further tuned this baseline method to make the results stronger than presented in their original paper. The improvement of DiffTOP over the ManiSkill baselines in Table 2 provides evidence to the effectiveness of our proposed method. \n\nWe also note that we have compared to more advanced BC baselines such as Diffusion Policy and 3D Diffusion Policy (DP3) in other benchmarks where these baseline methods are originally tested, i.e., MetaWorld and RoboMimic, and DiffTOP outperforms them in both benchmarkings,  showing the effectiveness of our proposed method.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "x4FZyHqfLc", "reviewer": "Reviewer_epML", "summary": "The paper introduces DiffTOP, an approach leveraging differentiable trajectory optimization as the policy representation to enhance performance in deep reinforcement learning (RL) and imitation learning (IL). By utilizing the advancements in differentiable trajectory optimization, DiffTOP addresses the \"objective mismatch\" issue prevalent in prior model-based RL algorithms by optimizing the dynamics model to directly maximize task performance. The method is benchmarked across various robotic manipulation tasks with high-dimensional sensory inputs and demonstrates superior performance over state-of-the-art methods in both model-based RL and IL domains.", "questions": "see the weaknesses section.", "limitations": "yes", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "1. The paper is well-written, structured, and intuitive, making it easy to understand.\n2. The experiments are well-executed with various benchmarks. The comparison with baselines is thorough, including state-of-the-art methods like DP3.\n3. The paper provides experiments in both imitation learning and reinforcement learning settings, which is commendable.\n4. DiffTOP is combined with various existing approaches to demonstrate the generalization of the method.", "weaknesses": "1. The work is not well-contextualized. While the paper makes a general claim about a policy class using differentiable trajectory optimization, using such technique is not new in the field, as the authors point out in the related work section. Despite the authors mention some differences in implementation (e.g., whether the dynamics model is learned), the paper presents this differentiable trajectory optimization as a new contribution. It would be beneficial for the introduction and diagrams to highlight the specific differences between the proposed work and other differentiable trajectory optimization approaches, rather than focusing on the differences among EBM Diffusion policy and TD-MPC. Additionally, calling the method Differentiable Trajectory Optimization might be inaccurate since it encompasses a broad range of works and could be misleading, failing to capture the differences.\n2. While the execution is commendable, the design choices could be better justified, and more analysis would strengthen the paper. Particularly, as mentioned in the related work section, there are closely related works like [2][19][42], but there is no direct comparison with them in the experiments. More analysis on the design choices, such as using TD-MPC for model-based RL and learning dynamic functions, would also be appreciated.\n3. The idea is not particularly novel. The combination of TD-MPC and differentiable trajectory optimization is rather straightforward.\n4. It would also be beneficial to have more validation on real robots.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We want to extend our heartfelt gratitude for taking the time to review our paper. Thank you for all valuable suggestions and comments on improving the quality of the paper. Below we respond to each of your comments in detail.\n\n**Q:  The work is not well-contextualized … It would be beneficial for the introduction and diagrams to highlight the specific differences between the proposed work and other differentiable trajectory optimization approaches …  calling the method Differentiable Trajectory Optimization might be inaccurate since it encompasses a broad range of works and could be misleading, failing to capture the differences.**\n\nThank you for this suggestion. We have updated the introduction and diagrams of our paper to highlight the differences between our method and other differentiable trajectory optimization approaches: “1) We are the first to show how differentiable trajectory optimization can be combined with deep model based RL algorithms, training dynamics, reward, Q function, and the policy end-to-end using task loss. In contrast, prior work focuses on imitation learning [2, 42], assumes known dynamics and reward structures and learns only a few parameters [2], or first learns the dynamics model with the dynamics prediction loss (instead of the task loss), and then uses the fixed learned dynamics for control [19]. 2) We are the first to show that the policy class by differentiable trajectory optimization can scale up to high-dimensional sensory observations like images and point clouds, achieving state-of-the-art performances in standard RL and imitation learning benchmarks. In contrast, prior works [2, 19, 42] only test their methods in customized tasks with ground-truth low-level states, such as CartPole or Pendulum, and do not report performance on standard benchmarks with more complex tasks and high-dimensional observations. ”\n\nThank you for the suggestion on the name of the method, and we will revise it to better reflect our unique contributions.\n\n**Q: ... there is no direct comparison with closely related works [2][19][42] in the experiments.**\n\nThank you for the feedback. We did not compare to [2][19][42] because we target different experiments. These related works all conduct experiments on customized tasks with ground-truth low-level states. In contrast, we test our method on standard RL and robotic imitation learning benchmarks, with high-dimensional sensory observations like images and point clouds. As these prior works have not been demonstrated on high-dimensional observations or more complex tasks, we originally compared to more recent state-of-the-art methods on these benchmarks, e.g., 3D Diffusion Policy [1]. \n\nWe have now included a comparison with Amos et al. in one of their tasks (pendulum swing-up with ground-truth low-level states) under imitation learning settings. Unlike Amos et al., who assumes known dynamics and reward structures and only learns 10 parameters, our method uses neural networks to represent both dynamics and reward functions without such assumptions. The metric is the cost of the learned policy. As in Amos et al., we test in two settings, pendulum without damping and with damping. Following Amos et al., their method does not model the damping effect in the assumed dynamics, so the ground-truth dynamics model is not realizable in the damping case. We also compared to an additional baseline in Amos et al., which uses a LSTM to predict the expert action. \n\nThe results (Table 1 in the PDF uploaded via global rebuttal) show our method performs slightly worse in the no damping case but noticeably better in the damping case. This is because Amos et al. assumes correct dynamics in the no damping case and learns only 10 unknown parameters, whereas the assumed dynamics structure is incorrect in the damping case; we use fully-connected neural networks to represent the dynamics function, avoiding such assumptions. It is generally difficult to know the exact correct dynamics function structure, especially for tasks with complex dynamics (e.g., with contacts) and high-dimensional observations (images and point clouds).\n\n**Q: More analysis on the design choices, such as using TD-MPC for model-based RL and learning dynamic functions, would also be appreciated.**\n\nThank you for the suggestion. We have updated section 5 of our paper to include more discussion and experiments on these design choices:\n“We choose TD-MPC for its simplicity and state-of-the-art performance. However, our method is compatible with any model-based RL algorithm that learns a dynamics model and a reward function. To show this, we have added experiments implementing our method on top of Dreamer-V3 [3], another state-of-the-art image-based model-based RL algorithm, and tested it on 4 Deepmind Control Suite tasks. The results (Fig. 1 of the PDF uploaded in the global rebuttal) show that integrating our method with Dreamer-V3 improves performance in 3 out of 4 tasks, indicating that our method can enhance other model-based RL algorithms as well. In our experiments, we have also combined our method with 3D Diffusion Policy for imitation learning, and it achieves significant improvements.\n\nWe choose to learn the dynamics function as we work directly with high-dimensional sensory inputs like images and point clouds, where manually specifying the analytic dynamics function is challenging. This is in contrast to learning from ground-truth low-level states where the dynamics model can be derived using physical laws. Therefore, we need to learn a dynamics model in the latent space, similar to prior work such as TD-MPC [4]. ”\n\n**Q: It would also be beneficial to have more validation on real robots.**\n\nThank you for the suggestion. While real-world validation is beyond the scope of this paper, which focuses on standard benchmarks in simulation, we believe that validation of our method on real robots would be valuable for the robotics community. we leave this as important future work.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ukU5mRRKTb", "reviewer": "Reviewer_XFZZ", "summary": "This paper proposes DiffTOP as a model-based approach to reinforcement learning and behavior cloning. DiffTOP learns a cost function and a dynamics model through differentiable trajectory optimization, and then uses the learned model during inference for online optimization. For model-based RL tasks, DiffTOP is built on a recent work of TD-MPC and resolves the objective mismatch issue that TD-MPC failed to address. For behavior cloning, the DiffTOP can be used with a different loss function tailed to BC, with an additional consideration of learning a multimodal policy which is achieved by the use of CVAE. The proposed DiffTOP framework is evaluated on 15 model-based RL tasks and 35 imitation learning tasks with high-dimensional observation spaces.", "questions": "* Regarding the use of CVAE for capturing multimodal action distributions for behavioral cloning, it seems that the distribution representation in the latent space is still unimodal Gaussian. Although it should work just fine, I wonder if the authors have considered to leverage a discrete categorical latent variable for explicitly modeling multi-modality even in the latent space to possibly improve the performance (see e.g., [6])\n\n[6] Ivanovic, B. and Pavone, M., 2019. The trajectron: Probabilistic multi-agent trajectory modeling with dynamic spatiotemporal graphs. In Proceedings of the IEEE/CVF international conference on computer vision (pp. 2375-2384).", "limitations": "[EDIT] The concerns were resolved throughout the rebuttal. In particular, my initial assessment of the paper regarding the ignorance of TD-MPC2 was inaccurate.\n\n* As mentioned above, the contribution of this paper is quite limited especially in light of the recent development of TD-MPC2.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "* The authors demonstrate broad applicability of the proposed DiffTOP framework in both reinforcement learning and behavior cloning settings, with varying problem complexity ranging from simple continuous control (such as cartpole) to complex object manipulation tasks (such as robomimic or maniskill).", "weaknesses": "[EDIT] The concerns were resolved throughout the rebuttal. In particular, my initial assessment of the paper regarding the ignorance of TD-MPC2 was inaccurate.\n\nDespite the thorough results presented in the paper, I have to say that it is not well contextualized relative to the existing literature.\n\n* First and foremost, the idea of leveraging model-based differentiable optimization for reinforcement learning and learning-based control is not new, and the field is rapidly expanding with new methods emerging in multiple domains, including control theory, robotics, and machine learning. For instance, Nikishin et al. [1] proposes an approach to differentiable optimization of trajectories (for an infinite horizon) by leveraging implicit function theorem and deep neural networks. Cheng et al. [2][3] perform differentiable trajectory optimization to learn a feedback policy, although they assume a known model-based function class for the dynamics and the policy. Sacks et al. [4] propose to optimize the inner-loop of an MPC optimization algorithm, though their method is based on MPPI. Given the abundance of similar ideas in the literature, the authors are encouraged to spend time to conduct a more thorough literature review. \n\n* Most critically, the authors of TD-MPC recently published a new paper titled TD-MPC2 [5] at ICLR 2024, in which the improvement from TD-MPC includes the mitigation of objective mismatch. Since the motivation is very similar to that of this paper and the authors of this paper completely disregard the new work, the contribution of this paper is highly questionable (even though the algorithmic details are still different between TD-MPC2 and DiffTOP). In particular, the following motivational statement made in Section 4 is no longer true: “Existing model-based RL algorithms such as TD-MPC suffer from the objective mismatch issue … DifftTOP addresses this issue …” \n\n[1] Nikishin, Evgenii, Romina Abachi, Rishabh Agarwal, and Pierre-Luc Bacon. \"Control-oriented model-based reinforcement learning with implicit differentiation.\" In Proceedings of the AAAI Conference on Artificial Intelligence, vol. 36, no. 7, pp. 7886-7894. 2022.\n\n[2] Cheng, Sheng, Minkyung Kim, Lin Song, Chengyu Yang, Yiquan Jin, Shenlong Wang, and Naira Hovakimyan. \"Difftune: Auto-tuning through auto-differentiation.\" arXiv preprint arXiv:2209.10021 (2022).\n\n[3] Cheng, Sheng, Lin Song, Minkyung Kim, Shenlong Wang, and Naira Hovakimyan. \"DiffTune $^+ $: Hyperparameter-Free Auto-Tuning using Auto-Differentiation.\" In Learning for Dynamics and Control Conference, pp. 170-183. PMLR, 2023.\n\n[4] Sacks, Jacob, Rwik Rana, Kevin Huang, Alex Spitzer, Guanya Shi, and Byron Boots. \"Deep model predictive optimization.\" arXiv preprint arXiv:2310.04590 (2023).\n\n[5] Hansen, Nicklas, Hao Su, and Xiaolong Wang. \"TD-MPC2: Scalable, Robust World Models for Continuous Control.\" In The Twelfth International Conference on Learning Representations (2024).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We want to extend our heartfelt gratitude for taking the time to review our paper. Thank you for all valuable comments and suggestions on improving the quality of the paper. Below we respond to each of your comments in detail.\n\n**Q: ... Given the abundance of similar ideas in the literature, the authors are encouraged to spend time to conduct a more thorough literature review.** \n\nWe thank the reviewer for bringing up these related works. We agree that the field of using differentiable optimization for end-to-end learning is rapidly expanding, and we have already discussed and cited many relevant works in the “Differentiable optimization” paragraph in the related work section. We are sorry for the omission of these 4 related works and we have updated section 2 of our paper to include a thorough discussion of them, detailed as follows:\n\nNikishin et al. [1] proposes to learn a dynamics and reward model in model-based RL, and they derive an implicit policy as the softmax policy associated with the optimal Q function under the learned dynamics and reward. The dynamics and reward model are learned by back-propagating the RL loss through this implicit policy via implicit function theorem. In contrast, we derive the implicit policy as the optimal solution from performing trajectory optimization with the learned dynamics, reward and Q function. Nikishin et al. only tested their methods in simple tasks with ground-truth low-level states (e.g., CartPole with only 2d action space and 4d state space), while we show our methods work with much more complex tasks in DeepMind Control Suite with high-dimensional image observations, and it outperforms several prior state-of-the-art methods.\n\nCheng et al. [2, 3] proposes to tune the parameters of a controller by unrolling the controller and system dynamics into a computation graph and optimizing the controller parameters via gradient descent with respect to the task loss, with applications in tracking trajectories for drones. The method assumes a known dynamics and policy class (e.g., a PID controller), and only optimizes some parameters of the controller. Our method does not assume any prior knowledge on the dynamics or policy class, instead, the dynamics and policy are all neural networks which we learn end-to-end using the task loss. Instead of representing the policy as a predefined controller and learning its parameters, our policy is represented as performing trajectory optimization with the learned dynamics, reward and Q functions.\n\nSacks et al. [4] proposes to learn the update rule in MPPI (the mean and variance used to sample the actions), represented as a neural network, using reinforcement learning, with known dynamics and cost functions. To apply RL for learning the update rule, they design an auxiliary MDP and treat the MPPI process as part of the dynamics of this auxiliary MDP. Instead of learning the update rule, we learn the dynamics, reward, Q function used in trajectory optimization to generate the actions. We perform differentiable trajectory optimization instead of RL to optimize the parameters of these functions.\n\nWe thank the reviewer again for introducing us to these 4 related works. As discussed above, we believe our proposed method is quite different from these 4 related works and our contributions remain valid.  \n\n**Q: The authors of TD-MPC recently published a new paper titled TD-MPC2 [5] at ICLR 2024, in which the improvement from TD-MPC includes the mitigation of objective mismatch. Since the motivation is very similar to that of this paper and the authors of this paper completely disregard the new work, the contribution of this paper is highly questionable.**\n\nWe thank the reviewer for raising this issue. We respectfully disagree with this assessment. TD-MPC2 [5] does not include any improvements over TD-MPC [6] that mitigate the objective mismatch issue. TD-MPC2 inherits the same training objective from TD-MPC, as evidenced by Equation (3) in TD-MPC2 and Equations (8), (9), and (10) in TD-MPC. This means that TD-MPC2 still suffers from the objective mismatch issue, where the dynamics model is optimized to predict future states in latent space (via the latent state consistency loss), which is not necessarily aligned with the goal of achieving high task return when using the dynamics model for planning. The primary improvements in TD-MPC2 over TD-MPC are related to exploring different architectural variations and new algorithmic design choices for more stable training in multi-task settings, rather than addressing the objective mismatch issue. Therefore, we believe our contribution and motivation remains valid. We have updated section 2 of our paper to include a discussion of TD-MPC2. \n\n**Q: I wonder if the authors have considered to leverage a discrete categorical latent variable for explicitly modeling multi-modality even in the latent space to possibly improve the performance**\n\nThank you for this great suggestion. We chose CVAE due to its simplicity, and as shown in our experiments, it achieves good performance and is able to capture the multi-modality well (Fig. 4). We believe that a more advanced technique, such as the discrete categorical latent variable as suggested, would further improve the performance. We look forward to exploring this in future work. \n\n\n[1] Nikishin et al., Control-oriented model-based reinforcement learning with implicit differentiation, AAAI Conference on Artificial Intelligence. 2022.  \n[2] Cheng et al., Difftune: Auto-tuning through auto-differentiation, IEEE Transactions on Robotics, 2024  \n[3] Cheng et al., DiffTune$^\\dagger$: Hyperparameter-Free Auto-Tuning using Auto-Differentiation, Learning for Dynamics and Control Conference, 2023.  \n[4] Sacks et al., Deep model predictive optimization, arXiv preprint, 2023.  \n[5] Hansen et al., TD-MPC2: Scalable, Robust World Models for Continuous Control, ICLR 2024  \n[6] Hansen et al., Temporal Difference Learning for Model Predictive Control, ICML 2022", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
Dear reviewers,

We want to extend our heartfelt gratitude for taking the time to review our paper. Thank you for all valuable comments and suggestions on improving the quality of the paper. We respond to each of your comments in detail in the individual rebuttal.

Following the reviewer's suggestions, the attached PDF contains two additional experiments:
- Comparison to closely related work Amos et al. (Reviewer epML)
- Combining DiffTOP with Dreamer-V3 as an ablation study on the underlying model-based RL algorithm (Reviewer epML)

[1] Amos, B., et al., Differentiable mpc for end-to-end planning and control. NeurIPS, 2018

Best,  
Authors

Meta-review: The paper presents a method using Differentiable Trajectory Optimization as a policy representation, extending TD-MPC with a policy-gradient loss. It addresses the "objective mismatch" issue by optimizing dynamics models to maximize task performance.

Overall, it's a technically strong paper, demonstrating good performance across diverse tasks. The reviewers had concerns about contextualizing the contribution within existing research, computational efficiency as well as novelty.  However the rebuttal does a good job in clarifying these. 

All reviewers lean accept, and the Metra-reviewer finds concurs.  
The authors are advised to review the final comments, and update the manuscript accordingly.
Decision: Accept (spotlight)