DiffTORI: Differentiable Trajectory Optimization for Deep Reinforcement and Imitation Learning

Authors: weikang  wan
ziyu  wang
yufei  wang
zackory  erickson
david  held
san  diego
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-13 20:42:20

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.

Main paper: 
 introduction: 
 Recent works have shown that the representation of a policy can have a substantial impact on the learning performance [1; 2; 3; 4] . Prior works have explored the use of feed-forward neural networks [4] , energy-based models [2] , or diffusion [1; 5] as the policy representation. In this paper, we propose to use differentiable trajectory optimization [3; 6; 7; 8; 9] as the policy representation to generate actions for deep reinforcement learning (RL) and imitation learning (IL) with highdimensional sensory observations (images/point clouds).Trajectory optimization is an effective and widely used algorithm in control, defined with a cost function and a dynamics function. It can be viewed as a policy [3; 6] , 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 (e.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-propagating through the trajectory optimization process [3; 8; 10; 6; 9; 11; 12; 13] . In prior work, differentiable trajectory optimization has been applied to system identification [3; 6; 9] , inverse optimal control [6] , imitation learning [3; 6; 8; 14; 7] and control/planning for robotics problems with low-dimensional states [3; 6; 8; 15] .In this paper, we propose to combine differentiable trajectory optimization with deep model-based RL algorithms. Because we use differentiable trajectory optimization to generate actions [10] , we 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; 17] 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-propagating the policy gradient loss through the trajectory optimization process. We show that our method outperforms prior state-of-the-art 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 feed-forward policy classes as well as Energy-Based Models (EBM) [2] and Diffusion [1] , and term our method DiffTORI (Differentiable Trajectory Optimization for Reinforcement and Imitation Learning). 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 [1] . We also outperform diffusionbased approaches [1] due to our procedure of learning a cost function that we optimize at test time. We show DiffTORI achieves state-of-the-art performance across 35 different tasks: 5 tasks from Robomimic [19] with image inputs, 9 tasks from Maniskill1 [20] and Maniskill2 [21] with point cloud inputs, and 22 tasks from MetaWorld [22] with point cloud inputs.Our work is closely related to prior work [3; 8; 6] 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 [3; 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 [8] . 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-the-art performances in standard RL and imitation learning benchmarks. In contrast, prior works [3; 8; 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 high-dimensional observations. In summary, the contributions of our paper are as following:• We introduce DiffTORI, which uses differentiable trajectory optimization as the policy representation for deep reinforcement learning and imitation learning.• We conduct extensive experiments to compare DiffTORI against prior state-of-the-art methods on 15 tasks for model-based RL and 35 tasks for imitation learning in standard benchmarks with high-dimensional sensory observations, and show that DiffTORI achieves superior performances.• We perform analysis and ablations of DiffTORI to provide insights into its performance gains. 
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 [11; 12; 6; 8; 9; 10] . The mostly related prior work to ours are Amos et al. [3] and Jin et al. [6] , which first proposed to treat trajectory optimization as an implicit policy and demonstrated its effectiveness in 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 a 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 from this work as we learn non-linear costs parameterized by a full neural network, and we also learn the dynamics instead of assuming it is known. We also show our method work with high-dimensional sensory inputs such as images and point clouds. Cheng et al. [25; 26] proposes to learn the parameters of 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 as a predefined controller, our policy is represented as performing 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 a 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 RL 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; 12] , and state 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 [17; 33] . Our approach also addresses the objective mismatch problem. In contrast to these two prior work which only optimizes a lower bound on the true return, our approach directly optimizes the task reward. Further, these approaches are only demonstrated using low-dimensional state-based observations whereas our approach is able to handle high-dimensional image or point cloud observations. In contrast to these works, we use Theseus [10] to analytically compute the gradient of the true objective for updating the model. Another related work, Nikishin et 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-propagating 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 to map observations to actions from demonstrations. Some recent work explores different policy architectures (e.g., explicit policy, implicit policy [2] , diffusion policy [1] ) and different action representations (e.g., mixtures of Gaussian [35; 19] , spatial action maps [36] , action flow [4] , or parameterized action spaces [37] ) to achieve more accurate learning from demonstrations, to 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.3 Background 
differentiable trajectory optimization: 
 In robotics and control, trajectory optimization solves the following type of problems:EQUATIONwhere c(s t , a t ) and C(s T ) are the cost functions, and s t+1 = d(s t , a t ) is the dynamics function.In this paper, we consider the case where the cost function and the dynamics functions are neural networks parameterized by θ: c θ (s t , a t ), C θ (s T ), and d θ (s t , a t ).Let a 0 (θ), ..., a T (θ) be the optimal solution to the trajectory optimization problem, which is a function of the model parameters θ. Differentiable trajectory optimization is a class of method that enables 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 it through the optimization process, to optimize the encoder as well as other latent space models (latent reward predictor and latent dynamics function) to maximize task performance.computation of the gradient of the actions with respect to the model parameters ∂at(θ) ∂θ . Specifically, in this paper we use Theseus [10] , which is an efficient application-agnostic open source library for differentiable nonlinear least squares optimization. Theseus works well with high-dimensional states, e.g., images or point clouds, along with using neural networks as the cost and dynamics functions.We use the standard MDP formulation: ⟨S, A, R, T , γ⟩ where S is the state space, A is the action space, R(s, a) is the reward function, T (•|s, a) is the transition dynamics function, and γ ∈ [0, 1) is the is the discount factor. The goal is to learn a policy π to maximize the expected return:E st,at∼π [ ∞ t=1 γ t R(s t , a t )].In this paper we work on problems where the state space S are highdimensional sensory observations, e.g., images or point clouds. Model-based RL algorithms first learn a dynamics model, and then use it for learning a policy. When applied to model-based RL, our method builds upon TD-MPC [38] , a recently proposed model-based RL algorithm which we review briefly here. 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. TD-MPC consists of the following components: first, an encoder h θ , which encodes the high-dimensional sensory observations, e.g., images, into a low-dimensional state z t = h θ (s t ). In the latent space, a latent dynamics model d θ is also learned: z t+1 = d θ (z t , a t ). A latent reward predictor R θ is learned which predicts the task reward r: r = R θ (z t , a t ). Finally, a value predictor Q θ learns to predict the Q value: Q = Q θ (z t , a t ). 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:EQUATIONwhere τ ∼ B is a trajectory (s t , a t , r t , s t+1 ) t:t+H sampled from a replay buffer B, λ ∈ R + is a constant that weights near-term predictions higher, and the single-step loss is:LTD-MPC(θ; τi) =c1∥R θ (zi, ai) -ri∥ 2 2 reward + c2∥Q θ (zi, ai) -ri + γQ θ -(zi+1, π θ (zi+1)) ∥ 2 2 value + c3∥d θ (zi, ai) -h θ -(si+1)∥ 2 2latent state consistency (3) where θ -are parameters of target networks that are periodically updated using the parameters of the learning networks. As shown in (3), the parameters θ is optimized with a 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) [39] is used for planning actions that maximize the predicted rewards and Q functions in the latent space. A policy π ψ is further learned in the latent space using the latent Q-value function, which is used to generate action samples in the MPPI process. 
overview: 
 The core idea of DiffTORI is to use trajectory optimization as the policy π θ , where θ parameterizes the dynamics and cost functions. Given a state s, DiffTORI generates the actions a(θ) by solving the trajectory optimization problem in (1) with s 0 = s. To optimize the policy parameters θ, we use differentiable trajectory optimization to compute the gradients of the loss L(a(θ)) with respect to the policy parameters: ∂L(a(θ)), where the exact form of the loss depends on the problem setting.An overview of applying DiffTORI to model-based RL is shown in Figure 1 . 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. DiffTORI 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 DiffTORI to imitation learning; an overview is shown in Figure 2 . In contrast to explicit policies that generate actions at test-time by forward passes of the policy network, DiffTORI 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 at 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 [1] . We describe DiffTORI for imitation learning in detail in Section 4.3.We build DiffTORI on top of TD-MPC for model-based RL. Similar to TD-MPC, DiffTORI consists of an encoder h θ , a latent dynamics model d θ , a reward predictor R θ , and a Q-value predictor Q θ (see Sec. 3.2). We use θ to denote all learnable parameters to be optimized in DiffTORI. As shown in Figure 1 , the key to DiffTORI is to change the non-differentiable MPPI planning algorithm in TD-MPC to a differentiable trajectory optimization, and include the policy gradient loss on the generated actions to optimize the model parameters θ directly for task performance.Formally, given a state s t , we use the encoder h θ to encode it to the latent state z t , and then construct the following trajectory optimization problem in the latent space:EQUATIONwhere H is the planning horizon. In this paper we leverage Theseus [10] to solve (4) in a differentiable way. Since Theseus only supports solving non-linear least-square optimization problems without constraints, we 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 turn the above optimization problem into the following one:EQUATIONWe set the values of H following the schedule as in TD-MPC, and we use the Levenberg-Marquardt algorithm in Theseus to solve the optimization problem. Following TD-MPC, we also learn a policy π ψ in the latent space using the learned Q-value predictor Q θ , and the output from the policy is used as the action initialization for solving (4).Let a(θ) be the solution of the above trajectory optimization problem, obtained using Theseus as described above. DiffTORI is learned with the following objective, which jointly optimizes the encoder, latent dynamics model, latent reward model, and the Q-value predictor: Although implicit policy shares the same spirit as DiffTORI, we observe that the training procedure of DiffTORI using differentiable trajectory optimization leads to better performance compared to the EBM approach used in prior work [2] , which can suffer from training instability.EQUATIONwhere Qϕ is the Q function learned via Bellman updates [40] which is used to compute the deteministic policy gradient [41] , and c 0 is the weight for this loss term. Qϕ is learned in the original state space S instead of the latent space to provide accurate policy gradients. The key idea here is that we can backpropagate through the policy gradient loss L P G , which backpropagates through a(θ) and then through the differentiable trajectory optimization procedure of Equation 4 to update θ.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 h θ and a latent dynamics function d θ , as before. However, in the setting of imitation learning, we do not assume access to a reward function R(s, a). Instead, we generate actions by solving the following trajectory optimization problem:EQUATIONin which f θ (z t , a t ) is a function over the latent state z t and actions a t that we will optimize using the imitation learning loss, as described below. Similarly, We use θ to denote all learnable parameters to be optimized in DiffTORI, including the parameters of the encoder h θ , the latent dynamics model d θ , and the function f θ in the imitation learning setting.In imitation learning, we assume access to an expert datasetD = {(s i , a * i )} N i=1 of state-action pairs (s i , a * i ).In the most basic form, the loss L for DiffTORI can be the mean square error between the the expert actions a * i and the actions a(θ) returned from solving (7) :EQUATIONThe key idea here is that we can backpropagate through the imitation loss L BC , which backpropagates through a(θ) and then through the differentiable trajectory optimization procedure of Equation 7to update θ. This enables us to learn the function f θ (z t , a t ) used in the optimization Equation 7directly by optimizing the imitation loss L BC (θ). Because this loss is optimized through the trajectory optimization procedure (Equation 7), we will learn a function f θ (z t , a t ) such that optimizing Equation 7 returns actions that match the expert actions.Multimodal DiffTORI: The loss in Equation 8will not be able to capture multi-modal action distributions in the expert demonstrations. To address this, we use a Conditional Variational Auto-Encoder (CVAE) [42] as the policy architecture, which has the ability to capture a multi-modal action distribution [43] . The CVAE encoder encodes the state s i and the expert action a * i into a latent state vector z i . The key idea in our approach is that the decoder in CVAE takes the form of a trajectory optimization algorithm, given by Equation 7. It takes as input the sampled latent z from the Gaussian Prior, and the state s i and uses differentiable trajectory optimization to decode the action a(θ). Because this trajectory optimization is differentiable, we can backpropagate through it to learn the parameters θ for the encoder, dynamics d θ , and the function f θ 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 of a 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 our experiments, we find DiffTORI 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 
model-based reinforcement learning: 
 We conduct experiments on 15 DeepMind Control suite tasks, which involve simulated locomotion and manipulation tasks, such as making a cheetah run or swinging a ball into a cup. All tasks use image observations and the control policy does not have direct access to the underlying states.We compare to the following baselines: SAC [44] , a commonly used off-policy model-free RL algorithm. DrQ-v2 [45] , a state-of-the-art model-free RL algorithm for image observations that adds data augmentation on top of SAC . TD-MPC [38] , a state-of-the-art model-based RL algorithm, which DiffTORI builds on. All training details such as hyper-parameters, and pysudo-code can be found in Appendix B. All experiments use NVIDIA 2080 Ti GPUs.Figure 3 shows the learning curves for all methods on all tasks. The top-left subplot shows the normalized performance averaged across all 15 tasks, which is computed as the achieved return divided by the max return from any algorithm. As shown, DiffTORI (red curve) outperforms all compared baselines by a noticeable margin. On 14 out of the 15 tasks (except Quadruped-walk), DiffTORI achieves the highest performance among compared algorithms. We especially note that the performance of DiffTORI is much higher than TD-MPC, which DiffTORI builds on, showing the benefit of adding the policy gradient loss and directly differentiating through it to optimize the learned latent spaces. Although DiffTORI achieves higher sample efficiency, one limitation of DiffTORI 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 vs 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. 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.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 is 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 1 for all the tasks). 10 demonstrations are used for all tasks [46] . We compare DiffTORI with the following baselines: DP3 [46] , a 3D version of diffusion policy that achieves state-of-the-art performances on this benchmark, outperforming other algorithms such as the original diffusion policy [1] 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 learn a residual policy on top of the base pre-trained policy, which takes as input the action from the base policy, and outputs a delta action that is added to 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 1 presents the task success rates, averaged over 50 evaluation episodes, of all compared algorithms. As shown, DiffTORI 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 DiffTORI, 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 a more effective way to leverage a pre-trained policy. On average, the success rates of DiffTORI is 7.7% higher than that of DP3, a substantial improvement with only 10 demonstrations.Robomimic [19] 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 a mixture of proficient and non-proficient 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 [1] 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]: An implicit policy that learns an energy function conditioned on both action and observation using the InfoNCE loss [47] . BC-RNN [19] : 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-performing baseline in the original Robomimic [19] 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 a delta action which is added to the base action. Diffusion Policy [1] : A policy that uses the diffusion model as the policy class. It refines noise into actions via a 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: DiffTORI + BC-RNN and DiffTORI + Diffusion Policy, which uses a 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 DiffTORI with zero initialization or random initialization, instead of initializing the action from a base policy.The results are shown in Table 2 . We find that DiffTORI+Diffusion Policy achieves the lowest failure rates consistently across all tasks. Even though Diffusion 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 Diffusion Policy -DiffTORI always brings noticeable improvement in the performance over the base policy. While learning a residual policy does lead to improvements upon the base policy, DiffTORI shows a significantly greater performance boost.In addition, by comparing DiffTORI+Diffusion 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 DiffTORI to capture multimodal action distributions. With different latent samples from CVAE, we get We note the two landscapes are distinct from each other with different optimal solutions, showing that DiffTORI 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 a 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.Goal ObjectManiSkill [20; 21] is a benchmark for learning generalizable robotic manipulation skills with 2D & 3D visual input. It includes a 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 ManiSkill2 [21] and use 3D point cloud input for all the tasks. We use the end-effector frame as the observation frame [49] 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 ManiSkill2 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 a residual policy that takes as 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 differentiable trajectory optimization as the policy class. 
conclusion and discussion: 
 We introduce DiffTORI that uses differentiable trajectory optimization to generate the policy actions for deep reinforcement learning and imitation learning. The key is to utilize the recent progress in differentiable trajectory optimization to 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 on standard robotic manipulation task suites with high-dimensional sensory observations and compare it to feed-forward policy classes as well as Energy-Based Models (EBM) and Diffusion. Across 15 model-based RL tasks and 35 imitation learning tasks with high-dimensional image and point cloud inputs, DiffTORI outperforms prior state-of-the-art methods. 
broader impacts: 
 Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed? 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)? 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 a 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 to 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 is challenging, and many papers do not require this, but we encourage authors to take this into account and make a 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 is used and, if 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.• If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • 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 is not available online, the authors are encouraged to reach out to the asset's creators. 
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/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc. • The paper should discuss whether and how consent was 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. 
crowdsourcing 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 as 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 in the supplemental material is fine, but if the main contribution 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. 
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 an equivalent approval/review based on the requirements of your country or institution) were obtained?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. • Depending on 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. 
 
 Appendix: In model-based reinforcement learning, the key distinctions between DiffTORI and TD-MPC [38] are: 1) TD-MPC employs the Model Predictive Path Integral (MPPI [51] ) 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-propagate it through the differentiable trajectory optimization process to update the model parameters. In addition to comparing the sample efficiency of DiffTORI to prior methods, we also compare the computational efficiency of DiffTORI versus TD-MPC on some of the environments. This is shown in Figure 7 , where the y-axis is the return, and the x-axis is the wall-clock time (tested on a NVIDIA RTX 2080 Ti GPU) used to train DiffTORI and TD-MPC for 1M environment steps. As shown, it takes more wall-clock time for DiffTORI to finish the training. In terms of computational efficiency, the results are environment-dependent. DiffTORI achieves better computational efficiency on reacher-hard and cup-catch. On pendum-swingup, TD-MPC converges to a sub-optimal value in the early training stage and DiffTORI outperforms it within 24 hours of training time. DiffTORI has similar computational efficiency on cartpole-swingup-sparse, reacher-easy, and finger-spin, and slightly worse computational efficiency on cheetah-run and walker-stand compared to TD-MPC. The gap is larger on hopper-stand. The major reason for DiffTORI to take longer time for training is that solving and back-propagating through the trajectory optimization problem in ( 4) is slower than doing MPPI as used in TD-MPC. As a 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 a 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 to 0.041 second. With better libraries/algorithms in the future for differentiable trajectory optimization, we believe the computational efficiency of DiffTORI would further improve. Figure 6 : Ablation study of DiffTORI to examine the contribution of each loss terms towards the final performance, on a subset of 4 tasks. We find the reward prediction loss, action initialization, and dynamics prediction loss are all essential for DiffTORI to achieve good performance.We also perform ablation studies to examine how each loss term in (6) contributes to the final performance of DiffTORI, as shown in Figure 6 . We find that removing the reward prediction loss causes DiffTORI to completely fail. Removing the dynamics loss, or not using the action initialization from the learned policy π ψ for solving the trajectory optimization, both lead to a decrease in the performance. These shows the necessity of using all the loss terms in DiffTORI for learning a good latent space to achieve strong performance.We also present results of DiffTORI with zero initialization or random initialization, where instead of initializing the action from a base policy, the action is initialized to be 0, or randomly sampled from N (0, 1), on RoboMimic and Maniskill.The results on RoboMimic is shown in Table 4 . We notice a drop in performance of DiffTORI with zero or randomly-initialized 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 a drawback of applying DiffTORI in practice for imitation learning: one could always first learn a base policy using any behavior cloning algorithm, and then use DiffTORI to further refine the actions.The results on Maniskill is shown in Table 5 . Again, if we use zero or random action initialization with DiffTORI, the performance drops to be similar to or slightly worse than vanilla BC. Therefore, we think a good practice of using DiffTORI for imitation learning would be to always try to provide it with a good action initialization, e.g., by first training a BC policy and use its action as the initialization in DiffTORI.Note that for the three tasks in Table 2 from Robomimic, we use the default velocity controller from Robomimic. We note the use of the velocity controller leads to a small decline in the performance of Table 4 : Failure rates (↓) of DiffTORI and all other mehtods on the Robomimic tasks. DiffTORI achieves the best performances on all tasks when using diffusion policy as the base policy. If zero or random initialization are used in DiffTORI, the performance drops, possibly due to the convergence to bad local minima during nonlinear trajectory optimization without a good action initialization. the Diffusion Policy compared to its performance in the original paper where a positional controller is used. The Push-T task still uses the default position controller as in the diffusion policy paper. Below we evaluate the performance of DiffTORI and Diffusion Policy with the positional controller.In the original Diffusion Policy [1] paper, it was observed that the use of positional controllers yielded superior results for Diffusion Policy compared to the default velocity controller on Robomimic [19] tasks. We evaluate Diffusion Policy, which is the strongest baseline, and DiffTORI on the most difficult three tasks with ph (proficient-human demonstration) and mh (multi-human demonstration) 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. DiffTORI, however, is able to achieve similar or even higher performances on most of these tasks.Additionally, we do ablation experiments on the planning horizon H for imitation learning, with the results presented in Table 7 . 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, the performance remains nearly the same; however, when H is increase to 5, we observe a slight decline in the performance. In our main paper, we did not compare to [3; 6; 8] 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 highdimensional 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 [46] .We have now included a comparison with Amos et al. [3] 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 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. The results in Table 8 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). Guidelines:• The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).Justification: There is no societal impact of this work. Our paper proposes a new 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:• 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 a 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 as 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 or unintentional) misuse of the technology. • If 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). 


Decision: 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.

Reviews: review: paper_summary: paper_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.

The 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.
strength_weakness: strength_weakness: [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.
questions: 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.

- 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: limitations: Yes. Limitations of the method are mentioned and I agree with the author's claims that there are no direct potential societal impacts.
paper_summary: paper_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.

The 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.
strength_weakness: strength_weakness: While the paper has many strengths, there are also some potential weaknesses or areas that could be improved:

- 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.
- 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.
questions: 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: limitations: The authors have discussed a few limitations in the paper.
paper_summary: paper_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.
strength_weakness: strength_weakness: 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.
2. 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.
3. The idea is not particularly novel. The combination of TD-MPC and differentiable trajectory optimization is rather straightforward.
4. It would also be beneficial to have more validation on real robots.
questions: questions: see the weaknesses section.
limitations: limitations: yes
paper_summary: paper_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.
strength_weakness: strength_weakness: [EDIT] The concerns were resolved throughout the rebuttal. In particular, my initial assessment of the paper regarding the ignorance of TD-MPC2 was inaccurate.

Despite the thorough results presented in the paper, I have to say that it is not well contextualized relative to the existing literature.

* 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. 

* 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 …” 

[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.

[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).

[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.

[4] Sacks, Jacob, Rwik Rana, Kevin Huang, Alex Spitzer, Guanya Shi, and Byron Boots. "Deep model predictive optimization." arXiv preprint arXiv:2310.04590 (2023).

[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).
questions: 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])

[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: 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.

* As mentioned above, the contribution of this paper is quite limited especially in light of the recent development of TD-MPC2.

Comments: title: Thank you for your response!
Thank you for your response!
Thank you for your response!
Response to Authors and Acknowledgement of Inaccurate Feedback in Initial Review
References for the rebuttal
comment: We sincerely thank the reviewer for reading our rebuttal and for the prompt response. We are glad that our rebuttal has addressed your concerns. We want to thank you again for taking the time to review our paper and read our rebuttal. Your suggestions and feedback have greatly helped improve the quality of our paper.
Thanks for the detailed response. I don't have further concerns.
We sincerely thank the reviewer for reading our rebuttal and for the prompt response. 

> One suggestion I would like to make is that DiffTop should be clearly distinguished from the previous works mentioned in the related work section. Currently, the paper’s organization doesn’t make this distinction immediately clear. A dedicated section, perhaps with a simple diagram highlighting the conceptual differences, would be helpful.

Thank you for this suggestion! We agree that dedicated section or a diagram would be especially helpful to clearly distinguish our method from prior works. We will certainly incorporate this in the final version of the paper, in addition to the updated discussion of the closely related works in the initial rebuttal.  

We want to thank you again for taking the time to review our paper and read our rebuttal. Your suggestions and feedback have greatly helped improve the quality of our paper.
We want to sincerely thank the reviewer for the prompt response, and the detailed explanation. We are glad our rebuttal has addressed your concerns, and we really appreciate the update of the review and the revision of the score. 

> It should further clarify the uniqueness of DiffTOP when compared against prior approaches that share similarity.

Thank you for this suggestion! In addition to the discussion to the 4 related works in the initial rebuttal, we will also update the paper to further clarify the uniqueness of our method when compared to prior approaches. 

We want to thank the reviewer again for taking the time to review our paper and read our rebuttal. Your suggestions and feedback have greatly helped improve the quality of the paper.
Thank you for the detailed and well-written response! I particularly appreciate the authors' effort in conducting additional experiments within such a short period. Overall, I remain positive about the paper. One suggestion I would like to make is that DiffTop should be clearly distinguished from the previous works mentioned in the related work section. Currently, the paper’s organization doesn’t make this distinction immediately clear. A dedicated section, perhaps with a simple diagram highlighting the conceptual differences, would be helpful.
I sincerely thank the authors for their time and effort to prepare the rebuttal. Please find below my response.

* > We 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:

  * Thank you for conducting the thorough literature review. It should further clarify the uniqueness of DiffTOP when compared against prior approaches that share similarity.


* > We 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.

  * Thank you for your response. I appreciate the level of details provided in the exposition and apologize for mistakenly making an inaccurate claim in my original review; I had been confused in part by the following statement made in Section 3.1 of TD-MPC2 paper, which I cite below for transparency:
    * > However, accurately predicting raw future observations (e.g., images or proprioceptive features) over long time horizons is a difficult problem, and does not necessarily lead to effective control (Lambert et al., 2020). Rather than explicitly modeling dynamics using reconstruction, TD-MPC2 aims to learn a maximally useful model: a model that accurately predicts outcomes (returns) conditioned on a sequence of actions. 

    Here, Lambert et al. (2020) is the original objective mismatch paper. It implicitly states that TD-MPC2 takes care of the objective mismatch issue. However, based on your response (and upon closer look at Appendix A of TD-MPC2), I agree that the model objective of TD-MPC2 is largely the same as that of TD-MPC. Thus, the results presented in Figure 3 of the present paper already confirms the superiority of DiffTOP over TD-MPC/TD-MPC2 in terms of directly optimizing the task performance. In the revised paper, the authors may want to explicitly counter-argue the above statement made by the TD-MPC2 authors, rather than just stating that TD-MPC suffers from the objective mismatch issue (e.g., in Section 4.1), since these two statements are indeed conflicting with each other and may cause confusion for other people as well.

Again, I appreciate the authors for taking the review comments seriously. I have updated my review and revised the score.
[1] Ze, Y., et al.,  3d diffusion policy: Generalizable visuomotor policy learning via simple 3d representations, RSS, 2024  
[2] Amos, B., et al., Differentiable mpc for end-to-end planning and control. NeurIPS, 2018  
[3] Hafner, D., et al., Mastering diverse domains through world models. arXiv preprint arXiv:2301.04104, 2013  
[4] Hansen N., et al., Temporal Difference Learning for Model Predictive Control, ICML 2022  
[19] Jin, W., et al., Pontryagin differentiable programming: An end-to-end learning and control framework.NeurIPS, 2020  
[42] Xu, M., et al., Revisiting implicit differentiation for learning problems in optimal control, NeurIPS, 2023