Stabilizing Reinforcement Learning in Differentiable Multiphysics Simulation
Authors: Eliot Xing
Vernon Luk
Jean Oh
Venue: ICLR 2025 Spotlight (2025)
Keywords: reinforcement learning
differentiable simulation
Primary area: reinforcement learning

Abstract: Recent advances in GPU-based parallel simulation have enabled practitioners to collect large amounts of data and train complex control policies using deep reinforcement learning (RL), on commodity GPUs. However, such successes for RL in robotics have been limited to tasks sufficiently simulated by fast rigid-body dynamics. Simulation techniques for soft bodies are comparatively several orders of magnitude slower, thereby limiting the use of RL due to sample complexity requirements. To address this challenge, this paper presents both a novel RL algorithm and a simulation platform to enable scaling RL on tasks involving rigid bodies and deformables. We introduce Soft Analytic Policy Optimization (SAPO), a maximum entropy first-order model-based actor-critic RL algorithm, which uses first-order analytic gradients from differentiable simulation to train a stochastic actor to maximize expected return and entropy. Alongside our approach, we develop Rewarped, a parallel differentiable multiphysics simulation platform that supports simulating various materials beyond rigid bodies. We re-implement challenging manipulation and locomotion tasks in Rewarped, and show that SAPO outperforms baselines over a range of tasks that involve interaction between rigid bodies, articulations, and deformables. Additional details at https://rewarped.github.io/.

Published as a conference paper at ICLR 2025
STABILIZING REINFORCEMENT LEARNING IN DIFFERENTIABLE MULTIPHYSICS SIMULATION
Eliot Xing & Vernon Luk & Jean Oh
Camegie Mellon University
(etaoxi ng, vluk, jeanch @mu edu
ABSTRACT
Recent advances in GPU-based parallel simulation have enabled practitioners to collect large amounts of data and train complex control policies using deep re- inforcement learning (RL), on commodity GPUs. However, such successes for RL in robotics have been limited to tasks sufficiently simulated by fast rigid- -ody dynamics. Simulation techniques for soft bodies are comparatively several orders of magnitude slower, thereby limiting the use of RI. due to sample complexity requirements. To address this challenge, this paper presents both a novel RI al- gorithm and simulation platform to enable scaling RL on tasks involving rigid bodies and deformables. We introduce Soft Analytic Policy Optimization (SAPO), a maximum entropy first-order model-based actor-critic RL algorithm, which uses first-order analytic gradients from differentiable simulation to train stochastic tor to maximize expected return and entropy. Alongside our approach, we develop Rewarped, parallel differentiable multiphysics simulation platform that supports simulating various materials beyond rigid bodies. We re-implement challenging manipulation and locomotion tasks in Rewarped, and show that SAPO outper- forms baselines over a range of tasks that involve interaction between rigid bodies, articulations, and deformables. Additional details rewarped, gi github..... io.
INTRODUCTION
Progress in deep reinforcement learning (RL) has produced policies capable of impressive behavior, from playing games with superhuman performance (Silver et al., 2016; Vinyals al., 2019) to controlling robots for assembly (Tang et al., 2023), dexterous manipulation (Andrychowicz et al., 2020; Akkaya et al., 2019), navigation (Wijmans et al., 2020; Kaufmann et al., 2023), and locomo- tion (Rudin et al., 2021; Radosavovic et al., 2024). However, standard model-free RL algorithms are extremely sample inefficient. Thus, the main practical bottleneck when using RL is the cost of acquiring large amounts of training data.
To scale data collection for online RL, prior works developed distributed RL frameworks (Nair et al., 2015; Horgan al., 2018; Espeholt et al., 2018) that many processes across a large compute cluster, which is inaccessible to most researchers and practitioners. More recently, GPU-based parallel environments (Dalton et al., 2020; Freeman al., 2021; Liang et al., 2018: Makoviychuk et al., 2021; Mittal et al., 2023; Gu et al., 2023) have enabled training RL at scale on single consumer GPU,
However, such successes of scaling RL robotics have been limited tasks sufficiently simulated by fast rigid-bod body dynamics (Makoviychuk et al., 2021), while physics physics-based based simulation techniques for soft bodies are comparatively several orders of magnitude slower. Consequently for tasks involving deformable objects, such as robotic manipulation of rope (Nair et al., 2017; Chi et al., 2022), cloth (Ha Song, 2022; Lin et al., 2022), elasties (Shen al., 2022), liquids (Ichnowski et al., 2022; Zhou et al., 2023), dough (Shi al., 2022; 2023; Lin et al., 2023), granular piles (Wang et al., 2023; Xue et al., 2023), approaches based on motion planning, trajectory optimization, or model predictive control have been preferred over and outperform RL (Huang al., 2020; Chen et al., 2022).
How can we overcome this data bottleneck to scaling RL on tasks involving deformables? Model- based reinforcement learning (MBRL) has shown promise at reducing sample complexity, by leverag- ing some known model or learning a world model to predict environment dynamics and rewards (Mo- erland et al., 2023). In contrast to rigid bodies however, soft bodies have complex dynamics

---

Published as a conference paper at ICLR 2025
AntRun
HandReorient
RollingFlat
SoftJumper
HandFlip
FluidMove
Figure/ 1: Visualizations of tasks implemented in Rewarped. These are manipulation and locomo- tion tasks involving rigid and soft bodies. AntRun and HandReorient are tasks with articulated rigid bodies, while RollingFlat, SoftJumpe, HandFlip, and FluidMove are tasks with deformables.
and dighes-dimmesiionll state spaces. This makes learning model dynamics of deformables highly nontrivial (Lin et al., 2021), often requiring specialized systems architecture and material-speciffc assumptions such volume preservation connectivity.
Recent developments in differentiable physics- based simulators of deformables (Hu et al., 2019b; Du al., 2021; Huang et al., 2020; Zhou et al., 2023; Wang al., 2024; Liang et al., 2019; Qiao et al., 2021a; Li et al., 2022b; Heiden et al., 2023) have shown that first-order gradients from differentiable simulation can be used for gradient-based trajectory optimization and achieve low sample complexity. Yet such approaches are sensitive to initial conditions and get stuck in local optima due to non-smooth optimization landscapes or discontinuities induced by contacts (Li al., 2022a; Antonova al., 2023). Additionally, existing soft-body simulations are not easily parallelized, which limits scaling RL them. Overall, there no existing simulation platform that parallelized, differentiable, and supports interaction between articulated rigid bodies and deformables.
In this paper, we approach the sample efficiency problem using first-order model-based RL (FO- MBRL), which leverages first-order analytic gradients from differentiable simulation to accelerate policy learning, without explicitly learning world model. Thus far, FO-MBRL has been shown to achieve low sample complexity articulated rigid-body locomotion tasks (Freeman et al., 2021; Xu et al.. 2021), but has not yet been shown to work well for tasks involving deformables (Chen et al., 2022). We hypothesize that entropy regularization gradients from differentiable simulation, such as by smoothing the optimization landscape (Ahmed al., 2019). To this end, we introduce novel maximum entropy FO-MBRL algorithm, alongside parallel differentiable multiphysics simulation platform for RL.
stabilize policy optimization over analytic
Contributions. i) We introduce Soft Analytic Policy Optimization (SAPO), first-order MBRL algorithm based on the maximum entropy RL framework. We formulate SAPO as an on-policy actor- critic RL algorithm, where stochastic actor is trained maximize expected return and entropy using first-order analytic gradients from differentiable simulation. ii) We present Rewarped, a scalable and easy-to-use platform which enables parallclizing RL environments of GPU-accelerated differentiable multiphysics simulation and supports various materials beyond rigid bodies. iii) We demonstrate that parallel differentiable simulation enables SAPO to outperform baselines over range of challenging manipulation and locomotion tasks re-implemented using Rewarped that involve interaction between rigid bodies, articulations, and deformables such as elastic, plasticine, or fluid materials.
2 RELATED WORK
We refer the reader to (Newbury et al., 2024) for an overview of differentiable simulation. We cover non-parallel differentiable simulation and madel- RI. in Appendix A.
2

---

Published as a conference paper at ICLR 2025
Materials?
Simulator
V?
Rigid
Articulated
Elastic
Plasticine
Fluid
Isaac Gym
X
Isaac Lab/ Orbit
X
X
ManiSkill
X
X
y*
TinyDiffSim
EXXX
XX
X
Brax
AXXX
MJX
X
X
X
DaXBench
v*
DFlex
X
Rewarped (ours)
Table 1: Comparison of physics-based parallel simulation platforms for RL. We use to indicate incomplete feature support at the time of writing. i) Isaac Lab/ Orbit: Missing deformable tasks due to breaking API changes and poor simulation stability / scaling. ii) ManiSkill: The latest version ManiSkil13 does not yet support the soft body tasks introduced in v2. iii) MJX: Stability issues with autodifferentiation and gradients. iv) DaXBench: Plasticine task was omitted from benchmark and requires additional development. v) DFlex: While later work (Murthy et al., 2021; Heiden et al., 2023) has built on DFlex to support elastic and cloth materials, their simulations were not parallelized.
Parallel differentiable simulation. There are few prior works on parallel differentiable simulators capable of running many environments together, while also computing simulation gradients in batches. TinyDiffSim (Heiden et al., 2021) implements articulated rigid- dynamics and contact models in C+//UDA that can integrate with various autodifferentiation libraries. Brax (Freeman et al., 2021) implements parallel simulator in JAX for articulated rigid-body dynamics with simple collision shape primitives. Recently, MJX is building on Brax to provide a JAX re-implementation of MuJoCo (Todorov et al., 2012), physics engine widely used in RL and robotics, but does not have feature parity with MuJoCo yet. These aforementioned parallel differentiable simulators are only capable of modeling articulated rigid bodies. DaXBench (Chen et al., 2022) also uses JAX to enable fast parallel simulation of deformables such as rope and liquid by Material Point Method (MPM) or cloth by mass-spring systems, but does not support articulated rigid bodies. DFlex (Xu al., 2021) presents a differentiable simulator based on source-code transformation (Griewank & Walther, 2008; Hu et al., 2020) of simulation kernel code to C++/UUA, that integrates with PyTorch for tape-based autodifferenniaaion. Xu et al. (2021) use DFlex for parallel simulation of articulated rigid bodics for high-dimensional locomotion tasks. Later work (Murthy et al., 2021; Heiden al., 2023) also used DFlex to develop differentiable simulations of cloth and elastic objects, but these were not parallelized and did not support interaction with articulated rigid bodies. To the best of our knowledge, there is no existing differentiable simulation platform that is parallelized with multiphysics support for interaction between rigid bodies, articulations, and various deformables. this paper, we aim to close this gap with Rewarped, our platform for parallel differentiable multiphysics simulation, and in Table 1 we compare Rewarped against existing physics- based parallel simulation platforms.
Learning control with differentiable physics. Gradient-based trajectory optimization is commonly used with differentiable simulation of soft bodies (Hu et al., 2019b; 2020; Huang et al., 2020; Li et al., 2023a; Zhou et al., 2023; Wang et al., 2024; Si et al., 2024; Du et al., 2021; Li et al., 2022b; Rojas et al., 2021: Qiao et al., 2020; 2021a; Liu et al., 2024; Chen et al., 2022; Heiden et al., 2023). Differentiable physics can provide physical priors for control in end-to-end learning systems, such for quadruped locomotion (Song et al., 2024), drone navigation (Zhang et al., 2024), robot painting (Schaldenbrand al., 2023), or motion imitation (Ren et al., 2023). Gradients from differentiable simulation can also be directly used for policy optimization. PODS (Zamora et al., 2021) proposes a first and second order policy improvement based on analytic gradients of a value function with respect to the policy's is action outputs. APG (Freeman et al., 2021) uses analytic simulation gradients to directly compute policy gradients. SHAC (Xu et al., 2021) presents an actor-critic algorithm, where the actor is optimized over short horizon using analytic gradients, and terminal value function helps smooth the optimization landscape. AHAC (Georgiev et al., 2024) modifies SHAC to adjust the policy horizon by truncating stiff contacts based on contact forces the norm of the dynamics Jacobian. Several works also propose different ways to overcome bias and non-hh dynamics resulting from contacts, by reweighting analytic gradients (Gao et al., 2024; Son et al., 2024) or
3

---

Published as a conference paper at ICLR 2025
explicit smoothing (Suh et al., 2022; Zhang et al., 2023; Schwarke et al., 2024). In this work, propose a maximum entropy FO-MBRL algorithm to stabilize policy learning with gradients from differentiable simulation.
BACKGROUND
Reinforcement learning (RL) considers agent interacting with environment, formalized as In this work, we consider discrete-time, infinite-horizon MDPs with continuous action spaces, where are actions, P: S
Markov decision process (MDP) represented by tuple (S, AA, P, RG Po,
7). Po,
S are states, E A S x A + S is the transition function, R: S x A + R is a reward function, Po(s) is an initial state distribution, and 7 is the discount factor. We want to obtain policy S A which maximizes the expected discounted sum of rewards (return) E,, [Σi=u y're] with = R st. az), starting from state So ~ Pg- We also denote the state distribution Par (8) and state-action distribution Pa (s, a) for trajectories generated by a policy π Ta,))s;)
In practice, the agent interacts with the environment trajectory (So; (oo 81, ........ 1, 0T-1). We can define the H H-step return:
steps finite-length episode, yielding a
R-1
Ro.h (T)
Σ y YT
t=
and standard RL objective to optimize 0 parameterizing policy ny to maximize the expected return
( E sompa [Ro.T].
Typically, the policy gradient theorem (Sutton et al., 1999) provides a useful expression of VzJ(TT) that does not depend the derivative of state distribution (-)
[ pr(s)) (oT(als)(" a) da ds,
(e)(T) &
where (Tt, a.) 1E [Rar]] is the Q-function (state-action value function).
We proceed to review zeroth-order versus first-order estimators of the policy gradient following the discussion in (Suh et al., 2022; Georgiev et al., 2024). We denote a single zeroth- estimate
TO] Ro:r
Vo log π (arlst).
where the zeroth-order batched gradient (ZOBG) is the sample mean V J 6 and is an unbiased estimator, under some mild assumptions to ensure the gradients are well-defined. The ZOBG yields N N-sample Monte-Carlo estimate commonly known the REINFORCE estimator (Williams, 1992) in RL literature, or the score function likelihood-ratio estimator. Policy gradient methods may use different forms of Equation 4 to adjust the bias and variance of the estimator (Schulman et al., 2015b). For instance, baseline term can be used to reduce variance of the estimator, by substituting Ro:r with R&r RI.+11-
Differentiable simulation as the environment provides gradients for the transition dynamics P and rewards R, so directly obtain analytic value for V Ro:r under policy no- In this setting, for a single first-order estimate:
J(T) = Vo RoRT
then the first-order batched gradient (FOBG) is the sample mean V = A ( and is also known the pathwise derivative (Schulman et al., 2015a) reparameterization trick (Kingma & Welling, 2014; Rezende et al., 2014; Titsias & Gredilla, 2014).
First-order model-based RL (FO-MBRL) aims analytic gradients) as known differentiable model, in contrast to vanilla MBRL which either assumes given non-differentiable model learns world model of dynamics and rewards from data.
differentiable simulation (and its first-order
4

---

Published as a conference paper at ICLR 2025
Analytic Policy Gradient (APG, Freeman et al. (2021)) uses FOBG estimates to directly maximize the discounted return over a truncated horizon:
++
J(m) = Σ (ernaiipph"" 'n),
(6)
and is also referred to as Backpropagation Through Time (BPTT, Werbos (1990); Mozer (1995)). particularly when the horizon is the full episode length (Degrave et al., 2019; Huang et al., 2020).
Short-Horizon Actor-Critic (SHAC, Xu al. (2021)) FO-MBRL algorithm which learns policy To and (terminal) value function V
+-1-1
( Σ
1=t
1+H-1
c(V) Σ
Bar [I[V(s) (s))l)
(8)
where (s,) are value estimates for state 8, computed starting from time step over an H-step
horizon. TD(A) (Sutton, 1988) is used for value estimation, which computes X-returns G G as weighted average of value-bootstrappe
k-step returns Gea++
H-1-t
V(s:) G = =1 A)
Σ
X G-GGe + - + GutHH,
(9)
where Greth (E x + *8++8)) The policy and value function are optimized in
an alternating fashion per standard actor-critic formulation (Konda Tsitsiklis, 1999). The policy gradient is obtained by FOBG estimation, with single first-order estimate:
1.M J(T)) Vo(Ro. + * VVss)),
(10)
and the value function is optimized as usual by backpropagating V CVV) of the mean- squared loss in Eq. Combining value estimation with a truncated short-horizon window where H < T (Williams & Zipser, 1995), SHAC optimizes over a smoother surrogate reward landscape compared to BPTT over the entire Tsstep episode.
SOFT ANALYTIC POLICY OPTIMIZATION (SAPO)
Empirically we observe that SHAC, Mtaeeofftthe-ue FO-MBRL algorithm, still prone suboptimal convergence to local minima in the reward landscape (Appendix, Figure 5). We hypothesize that entropy regularization can stabilize policy optimization over analytic gradients from differentiable simulation, such as by smoothing the optimization landscape (Ahmed et al., 2019) or providing robustness under perturbations (Eysenbach & Levine, 2022).
We draw on the maximum entropy RI. framework (Kappen, 2005; Todorov, 2006; Ziebart et al., 2008; Toussaint, 2009; Theodorou et al., 2010; Haarnoja et al., 2017) to formulate Soft Analytic Policy Optimization (SAPO), maximum entropy FO-MBRL algorithm (Section 4.1). 1). To implement SAPO, we make several design choices, including modifications building on SHAC (Section 4.2). In Appendix B. 1, deseribe how visual observations differentiable simulation. Pseudocode for SAPO is shown Appendix B.2, and the computational graph of SAPO is illustrated in Appendix Figure 4.
visual encoders to learn policies from high-dimensional
4.1 MAXIMUM ENTROPY RL DIFFERENTIABLE SIMULATION
Maximum entropy RL (Ziebart et al., 2008; Ziebart, 2010) augments the standard (undiscounted) return maximization objective with the expected entropy of the policy over P. (st):
((T)
In + a azz st]],
(11)
5

---

Published as a conference paper at ICLR 2025
where [7/aasss] - क (aa log (Taels))da, is the continuous Shannon entropy of the action distribution, and the temperature a balances the entropy term versus the reward.
Incorporating the discount factor (Thomas, 2014; Haarnoja al., 2017), obtain the following objective which maximizes the expected return and entropy for future states starting from (st, a₂) weighted by its probability Pr under policy π
80
Maxent (T)
=
Σ
E6 (s happ
/ Bi Io + aH, | [
(12)
The soft Q-function is the expected value under of the discounted sum of rewards and entropy:
Vaore (st, ar) = re E (o
Σ I'n aHHlarls))
(13)
and the soft value function is:
V5sn (8t) = a log /
exp( Qmon (s, a))
Sapl-Qrott
(14)
1 a
When x (a)s) exp( !ud (Q (8, a) Vion (s))) - " then the soft Bellman equation yields the following relationship:
Thot (8t, az) Tr + (s....... Jup Vosn (8t+1)]: where we can rewrite the discounted maximum entropy objective Eq. 12:
(15)
Imaxent (T)
B(a, (s, + [a///a
(16)
Σ E (mn Ire + aHx [a,
|s + Vanit (st+1))
(17)
By Soft Policy Iteration (Haamojj et al., 2018a), the soft Bellman operator B defined by (B* (8t, a.) rt 7Bs,11pp (V(++1)) has unique contraction Q" BBQ (Fox al., 2016) and converges to the optimal policy .².
Our main observation is when the environment is differentiable simulation, estimates to directly optimize maxont (T), including discounted policy entropy. Consider the entropy- augmented H/stepp return:
FOBG
H-1 7 + aH an (ar(s:)),
Re.h (T) Σ
(18)
t=
then have single first-order estimate of Eq. 17:
Juaxennt (T) Vo + Vauit (sH)).
(19)
Furthermore, we can incorporate the angropy-augmented return into TD) estimates of Eq. 9 using soft value- boots--boottt k-step returns:
TRL+k =
Σ t=
(res aHa + + Voon
(20)
where (st) = I'A 1 and the value function is trained by minimizing Eq. 8 with Vsoft Veont. and Tr:++* substituted in place of V. V. and Get+a- We refer to this maximum entropy FO-MBRI. formulation as Soft Analytic Policy Optimization (SAPO).
Note that we instantiate SAPO as an actor-critic algorithm that learns the soft value function by TD learning with on-policy data. In comparison, Soft Actor-Critic (SAC), popular off-policy maximum entropy model-free RL algorithm, estimates soft Q-values by minimizing the soft Bellman residual with data sampled from replay buffer. Connections may also drawn between SAPO to maximum entropy variant SVG(H) (Heess et al., 2015; Amos et al., 2021), which uses rollouts from a learned world model instead of trajectories from differentiable simulation.

---

Published as a conference paper at ICLR 2025
4.2 DESIGN CHOICES
I. Entropy adjustment. In practice, we apply automatic temperature tuning (Haarnoja et al., 2018b)
to match a target entropy H via an additional Lagrange dual optimization step:
min E (-)p,)))pp [ave [t] H)]. 20
(21)
We use H = (d)m following (Ball et al., 2023).
II. Target entropy normalization. To mitigate non con-stationarity in target values (Yu et al.. 2022) and improve robustness across tasks with varying reward scales and action dimensions, we normalize entropy estimates. The continuous Shannon entropy not scale invariant (Marsh, 2013). In particular. we offset (Han & Sung, 2021) and scale entropy by H to be approximately contained within [0, +1}.
III. Stochastic policy parameterization. We use state-dependent variance, with squashed Normal distribution Th = =tanh(N(Ho(s)) (s))), ( which aligns with SAC (Haarnoja et al., 2018b). This enables policy entropy adjustment and captures aleatoric uncertainty the environment (Kendall & Gal, 2017; Chua et al., 2018). In contrast, SHAC uses state-independent variance, similar to the original PPO implementation (Schulman et al., 2017).
IV. Critic ensemble, no target networks. We use the clipped double critic trick (Fujimoto et al., 2018) and also remove the critic target network in SHAC, similar to (Georgiev et al., 2024). However when updating the actor, instead compute the average over the two value estimates include in the return (Eq. 19), while using the minimum to estimate target values in standard fashion, following (Ball al., 2023). While originally intended to mitigate overestimation bias in Q-learning (due to function approximation and stochastic optimization (Thrun & Schwartz, 2014)), prior work has shown that the value lower bound obtained by clipping can be overly conservative and cause the policy to pessimistically underexplore (Ciosek et al., 2019; Moskovitz et al., 2021).
Target networks (Mnih et al., 2015) are widely used (Lillicrap al., 2016; Fujimoto et al., 2018; Haarnoja et al., 2018b) to stabilize temporal difference (TD) learning, at the cost of slower training. Efforts have been made to eliminate target networks (Kim al., 2019; Yang et al., 2021; Shao et al., 2022; Gallici et al., 2024), and recently CrossQ (Bhatt et al., 2024) has shown that careful use of normalization layers can stabilize off-policy model-free RL to enable removing target networks for improved sample efficiency. CrossQ also reduces Adam 81 momentum from 0.9 to 0.5, while keeping the default 3/2 = 0.999. In comparison, SHAC uses 31 = 0.7 and B2 = 0.95. Using smaller momentum parameters decreases exponential decay (for the moving average estimates of the 1st and 2nd moments of the gradient) and effectively gives higher weight to more recent gradients, with less smoothing by past gradient history (Kingma & Ba, 2015).
V. Architecture and optimization. We use SiLU (Elfwing et al., 2018) instead of ELU for the activation function. We also switch the optimizer from Adam to AdamW (Loshchilov & Hutter, 2017), and lower gradient norm clipping from 1.0 to 0.5. Note that SHAC already uses LayerNorm (Ba al., 2016), which has been shown stabilize TD learning when not using target networks replay buffers (Bhatt et al., 2024; Gallici et al., 2024).
REWARPED: PARALLEL DIFFERENTIABLE MULTIPHYSICS SIMULATION
We aim to evaluate our approach on more challenging manipulation and locomotion tasks that involve interaction between articulated rigid bodies and deformables. To this end, we introduce Rewarped, our parallel differentiable multiphysics simulation platform that provides GPU-accelerated parallel environments for RI. and enables computing batched simulation gradients efficiently. We build Rewarped on NVIDIA Warp (Macklin, 2022), the successor to DFlex (Xu et al., 2021; Murthy et al., 2021: Turpin et al., 2022; Heiden et al., 2023).
We proceed to discuss high-level implementation details and optimization tricks to enable efficient parallel differentiable simulation. We develop parallelized implementation of Material Point Method (MPM) which supports simulating parallel envirouments of complex deformable materials, building on the MLS-MPM implementation by (Ma et al., 2023) used for non-par moraalllll simulation. Furthermore, we support nay coupling from kinematic articulated rigid bodies to MPM particles, based on the (non-parallel) MPM-based simulation from (Huang et al., 2020; Li al., 2023a).
7

---

Published as a conference paper at ICLR 2025
5.1 PARALLEL DIFFERENTIABLE SIMULATION
We implement all simulation code in NVIDIA Warp (Macklin, 2022), library for differentiable programming that converts Python code into CUDA kernels by runtime JIT compilation. Warp implements reverse-mode auto-differemtiation through the discrete adjoint method, using a tape to record kernel calls for the computation graph, and generates kernel adjoints to compute the backward pass. Warp uses source-code transformation (Griewank & Walther, 2008; Hu et al., 2020) to automatically generate kernel adjoints.
We use gradient checkpointing (Griewank & Walther, 2000; Qiao et al., 2021b) to reduce memory requirements. During backpropogation, we run the simulation forward pass again to recompute intermediate values, instead of saving them during the initial forward pass. This is implemented by capturing and replaying CUDA graphs, for both the forward pass and the backward pass of the simulator. Gradient checkpointing by CUDA graphs enables gradients over multiple time steps efficiently, when using more simulation substeps for simulation stability. We use a custom PyTorch autograd function to interface simulation data and model parameters between Warp and PyTorch while maintaining auto-differentiation functionality.
compute batched simulation
EXPERIMENTS
We evaluate our proposed maximum entropy FO-MBRI. al-MBRRaa Soft Analytic Policy Optimization (SAPO, Section 4), against baselines on a range of locomotion and manipulation tasks involving rigid and soft bodies. We implement these tasks in Rewarped (Section 5), our parallel differentiable multiphysics simulation platform. We also compare algorithms on DFlex rigid-body locomotion tasks introduced (Xu et al., 2021) in Appendix F.2.
Baselines. We compare vanilla model-free RL algorithms: Proximal Policy Optimization (PPO, Schulman et al. (2017)), an on-policy actor-critic algorithm: Soft Actor-Critic (SAC, Haarnoja et al. (2018b)) off-policy maximum entropy actor-critic algorithm. We use the implementations and hyperparameters from (L. et al., 2023b) for both, which have been validated to scale well with parallel simulation. Implementation details (network architecture, common hyperparameters, etc.) are standardized between methods for fair comparison, see Appendix C. We also compare against Analytic Policy Gradient (APG, Freeman al. (2021)) and Short-Horizon Actor-Critic (SHAC, Xu FO-MBRL algorithms that leverage first-order ana-
et al. (2021)), both of which are state-of-the lytic gradients from differentiable simulation for policy learning. Finally, we include gradient- based trajectory optimization (TrajOpt) as baseline, which uses differentiable simulation gradients to optimize for an open-loop action sequence that maximizes total rewards across environments.
Tasks. Using Rewarped, we re-implement a range of challenging manipulation and locomotion tasks involving rigid and soft bodies that have appeared in prior works. Rewarped enables training algorithms on parallel environments, and differentiable simulation to compute analytic simulation gradients through environment dynamics and rewards. We visualize these tasks in Figure 1. To 2500 particles per env. See Appendix E for more details.
simulate deformables,
AntRun - Ant locomotion task from DFlex (Xu et al., 2021), where the objective is to maximize the forward velocity of four-legged ant rigid-body articulation.
Reodiet Allegro hand manipulation task from Isaac Gym (Makoviychuk et al., 2021) 1), where the objective is to perform in-hand dexterous manipulation to rotate rigid cube given target pose. We replace non-differentiable terms of the reward function (ie. boolean comparisons) with differentiable alternatives to enable computing analytic gradients.
RollingFlat Rolling pin manipulation task from PlasticineLab (Huang et al., 2020), where the objective is to flatten a rectangular piece of dough using a cylindrical rolling pin.
SoftJumper - Soft jumping locomotion task, inspired by GradSim (Murthy et al., 2021) and DiffTaichi (Hu et al., 2020), where the objective is to maximize the forward velocity and height of high- high-dimensional actuated soft elastic quadruped.
HandFlip Shadow hand flip task from DexDeform (Li et al., 2023a), where the objective is to flip cylindrical piece of dough in half within the palm of a dexterous robot hand.
FluidMove - Fluid transport task from SoftGym (Lin et al., 2021), where the objective is to move container filled with fluid given target position, without spilling fluid out of the container.
8

---

Published as a conference paper at ICLR 2025
Note that {AntRun, HandReorient) are tasks that involve articulated rigid bodies only, with state- based observations. In contrast, [RollingFlat, SoftJumper. HandFlip, FluidMove] are tasks that also involve deformables, with both state-based and high-dimensional (particle-basse visual observations.
AntRun
Hendieo
RollingFlat
5k
250
100
Return
125
2.5k
87.5
0
-125 4
75
0
2M
4M
2M
4M
0
2M
4M
SoftJumper
HandFlip
FluidMove
2.4k
100
30
Return
1.2k
20
0-
0+
10
3M
6M
3M
6M
0
2M
4M
SAPO (ours)
SHAC
APG
SAC
PPO
Figure 2: Rewarped tasks training curves. Episode return as a function of environment steps in Rewarped AntRun (A C IR$), HandReorient (AA C R16), RollingFlat (A C RR², SoftJumper
(A c )222), HandFlip (AA C IR2",, and FluidMove (AA C I2) tasks. Smoothed using EWMA with
0.99. Mean and 95% Cls over 10 random seeds.
AntRun
HandReorient RollingFlat
SoftJumper
HandFlip
FluidMove
PPO
2048.7 36.6
5.9 4.9
81.2 + 0.1
261.5 + 12,4
7.3 + 1.1
27.3 : 0.2 28.2 0.7
SAC
2063.6 + 13.9
70.5 + 10.2
83.0 + 0.3
-161.8 + 2.5
4.6 1.1
TrajOpt
915.5 + 29.6
81.5 + 0.1
437.2
17.7
27.3 2.6
27.0 7,0++00 0.1
APG
258.7 20.3
1.9
86.9 - 0.4
956.6 + 15.6
38.2
26.3 + 0.3
SHAC
3621.0 54.4
-
86.8 = 0.4
853.3 10.2
32.7 2.9
21.7 0.4
SAPO (ours)
4535.9 = 24.5
221.7 + 9.5
100.4 + 0.4 1820.5 + 47.9 90.0 + 2.2 30.6 = 0.4
Table 2: Rewarped tasks tabular results. Evaluation episode returns for final policies after training.
Mean and 95% CIs over 10 random seeds with 2N episodes per seed for 32 or 64 parallel envs.
6.1 RESULTS ON REWARPED TASKS
We compare SAPO, proposed maximum entropy FO-MBRL algorithm, against baselines a range of challenging manipulation and locomotion tasks that involve rigid and soft bodies, re- implemented in Rewarped, our parallel differentiable multiphysics simulation platform. In Figure we visualize training curves to compare algorithms. SAPO shows better training stability across different random seeds, against existing FO-MBRL algorithms APG and SHAC. In Table 2, report evaluation performance for final policies after training. SAPO outperforms all baselines across all tasks we evaluated, given the same budget of total number of environment steps. We also find that on tasks involving deformables, APG outperforms SHAC, which is consistent with results in DaXBench (Chen et al., 2022) on their set of soft-body manipulation tasks. However, SHAC outperforms APG on the articulated rigid-body tasks, which agrees with the rigid-body locomotion results in DFlex (Xu et al., 2021) that we also reproduce ourselves in Appendix F.2.
In Appendix Figure 11, we visualize different trajectories produced by SAPO policies after training. We observe that SAPO learns to perform tasks with deformables that we evaluate on. For RollingFlat, SAPO controls the rolling pin to flatten the dough and spread it across the ground. For SoftJumper SAPO learns a locomotion policy that controls soft elastic quadruped to jump forwards. For HandFlip, SAPO capable of controlling high sogrree-ffffreddoo dexterous robot hand, to flip

---

Published as a conference paper at ICLR 2025
piece of dough in half within the palm of the hand. For FluidMove, SAPO learns policy to move the container of fluid with minimal spilling. Additionally, SAPO learns a successful locomotion policy for the articulated rigid-body locomotion task AntRun. For HandReorient however, SAPO is only capable of catching the cube and preventing it from falling to the ground. This is a challenging task that will likely require environment steps learn policies capable of re-orienting the cube to given target poses in succession.
6.2 SAPO ABLATIONS
We investigate which components of SAPO yield performance gains over SHAC, on the HandFlip task. We conduct several ablations SAPO: (a) w/o Vaoft, where instead the critic is trained in standard fashion without entropy in target values; (b) w/o H., where we do not use entropy-augmented returns and instead train the actor to maximize expected returns only; (c) w/o H₂ and Vaott, which corresponds to modifying SHAC by applying design choices {III, IV, V) described Section 4.2.
We visualize training curves in Figure 3, and in Table 3 report final evaluation performance as well as percentage change from SHAC's performance as the baseline. From ablation (b), we find that using analytic gradients to train policy maximize both expected return and entropy is critical to the performance of SAPO, compared to ablation (a) which only replaces the soft value function.
Additionally, we observe that ablation (c), where apply design choices {III, IV, V) onto SHAC, result in approximately half of the performance improvement of SAPO over SHAC on the HandFlip task. We also conducted this ablation these modifications to SHAC to have minimal impact in those settings, shown in Appendix F.3. We also conduct individual ablations for these three design choices in Appendix F.4.
the DFlex rigid-body locomotion tasks however, and found
HandFlip
HandFlip
(A%)
100
SAPO (ours) w/o Vist
SAPO (ours)
90 +
+ +172.7%
Return
w/o Viont
77 +
+ 133.3%
50
w/o He
w/o H
59 +
+78.8%
w/o He and Vin
w/o Hl and Viott
56 +
+69.7%
SHAC
0-
SHAC
33 + 3
0
3M 6M
Figure 3: SAPO ablations HandFlip training curves. Episode return as a function of environ- ment steps. Smoothed using EWMA with 0.99. Mean and 95% CIs over 10 random seeds.
Table 3: SAPO ablations HandFFip tabular
results. Evaluation episode returns for final policies after training. Mean and 95% CIs over 10 random seeds with 64 episodes per seed.
CONCLUSION
Due to high sample complexity requirements and slower runtimes for soft-body simulation, RL has had limited success on tasks involving deformables. To address this, we introduce Soft Analytic Policy Optimization (SAPO), first-order model-based actor-critic RL algorithm based on the maximum entropy RL framework, which leverages first-order analytic gradients from differentiable simulation to achieve higher sample efficiency. Alongside this approach, we present Rewarped, a scalable and easy-to-use platform which enables parallelizing RI. environments of GPU-based differentiable multiphysics simulation. We c--imlemmnn challenging locomotion and manipulation tasks involving rigid bodies, articulations, and deformables using Rewarped. On these tasks, we demonstrate that SAPO outperforms baselines in terms of sample efficiency as well as task performance given the same budget for total environment steps.
Limitations. SAPO relies on end-to-end learning using first-order analytic gradients from differen- tiable simulation. Currently, we use (non-occluded) subsampled particle states from simulation as observations to policies, which is infeasible to obtain in real-world settings. Future work may use differentiable rendering to provide more realistic visual observations for policies while maintain- ing differentiabilityy towards sim2real transfer of policies learned using SAPO. Another promising direction to consider is applications between differentiable simulation and learned world models.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENTS
The authors would like to thank Uksang Yoo, Ananye Agarwal, and David Held for helpful discussions on this work. We would also like to thank Miles Macklin, Eric Heiden, and the rest of the NVIDIA Warp team. This work was in part supported by the Technology Innovation Program (200181 12, Development of autonomous manipulation and gripping technology using imitation learning based on disual-tactile sensing) funded by the Ministry of Trade, Industry & Energy (MOTIE), Korea.
REFERENCES
Zafarali Ahmed. Nicolas Roux, Mohammad Norouzi, and Dale Schuurmans. Understanding the impact of entropy on policy optimization. In ICML, 2019.
Iige Akkaya, Marcin Andrychowicz, Maciek Chociej, Mateusz Litwin, Bob McGrew, Arthur Petron, Alex Paino, Matthias Plappert, Glenn Powell, Raphael Ribas, et al. Solving rubik's cube with a robot hand. arXiv preprint arXiv: arXiv:1910.0711 1910.071133 2019.
Brandon Amos, Samuel Stanton, Denis Yarats, and Andrew Gordon Wilson. On the model-based stochastic value gradient for continuous reinforcement learning. L4DC, 2021.
OpenAl: Marcin Andrychowiez, Bowen Baker, Maciek Chociej, Rafal Jozefowicz, Bob McGrew, Jakub Pachocki, Arthur Petron, Matthias Plappert, Glenn Powell, Alex Ray, et al. Learning dexterous in-hand manipulation. The International Journal of Robotics Research, 39(1):3-20, 2020.
Rika Antonova, Jingyun Yang, Krishna Murthy Jatavallabhula, and Jeannette Bohg. Rethinking optimization with differentiable simulation from global perspective. In CoRL, 2023.
Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. arXiv preprint arXiv: 2016.
Philip J Ball, Laura Smith, Ilya Kostrikov, and Sergey Levine. Efficient online reinforcement learning with offline data, ICML, 2023.
Aditya Bhatt, Daniel Palenicek, Boris Belousov, Max Argus, Artemij Amiranashvili, Thomas Brox, Peters. Crossq: Batch normalization in deep reinforcement learning for greater sample efficiency and simplicity. In ICLR, 2024.
Siwei Chen, Yiqing Xu, Cunjun Yu, Linfeng Li, Xiao Ma, Zhongwen Xu, and David Hsu. Daxbench: Bench- marking deformable object manipulation with differentiable physics. In ICLR, 2022.
Cheng Chi, Benjamin Burchfiel, Eric Cousineau. Siyuan Feng, and Shuran Song. Iterative residual policy for goal-conditioned dynamic manipulation of deformable objects. In RSS. 2022.
Kurtland Chua, Roberto Calandra, Rowan McAllister, and Sergey Levine. Deep reinforcement learning in handful trials using probabilistic dynamics models. NeurlPS, 2018.
Kamil Ciosek, Quan Vuong. Robert Loftin, and Katja Hofmann. Better exploration with optimistic actor critic. In NewrlPS, 2019.
Ignasi Clavera, Yao Fu, and Pieter Abbeel. Model-augmented actor-critic: Backpropagating through paths. In ICLR, 2020.
Steven Dalton, Michael Garland, and luri Frosio. Accelcrating reinforcement learning through gpu atari emulation. In NeurlPS, 2020.
Filipe de Avila Belbute-Peres, Kevin Smith, Kelsey Allen, Josh Tenenbaum, and J Zico Kolter. End-to-end differentiable physics for learning and control. In NewrlPS, 2018.
Jonas Degrave, Michiel Hermans, Joni Dambre, and Francis Wyffels. A differentiable physics engine for deep learning in robotics. Frontiers Neurorobotics, 13, 2019.
Mare Peter Deisenroth, Dieter Fox, Carl Edward Rasmussen. Gaussian processes for data-efficient learning in robotics and control. IEEE Transactions on Pattern Analysis and Machine Intelligence, 37(2):408 423, 2013.
Tao Du, Kui Wu, Pingchuan Ma. Sebastien Wah. Andrew Spielberg, Daniela Rus, and Wojciech Matusik. Diffipd: Differentiable projective dynamics. ACM Transactions Graphics (TOG), 41(: 2021.
Stefan Elfwing. Fiji Uchibe, and Kenji Doya. Sigmoid linear units for neural network function approximation in reinforcement leamning Neural networks, 107:3-11, 2018.
11

---

Published as a conference paper at ICLR 2025
Lasse Espeholt, Hubert Soyer, Remi Munos, Karen Simonyan. Vlad Mnih, Tom Ward, Yotam Doron, Vlad Firoiu, Tim Harley, lain Dunning, al. Impala: Scalable distributed deep-rl with importance weighted actor-learner architectures. In ICML, 2018.
Benjamin Eysenbach and Sergey Levine. Maximum entropy rl (provably) solves some robust problems. In ICLR, 2022.
Vladimir Feinberg, Alvin Wan, Ion Stoica, Michael Jordan, Joseph E Gonzalez, and Sergey Levine. Model- based value estimation for efficient model-free reinforcement learning, arXiv preprint arXiv: 1rXiv:1033000101 2018.
Roy Fox, Ari Pakman, and Naftali Tishby. Taming the noise in reinforcement learning via soft updates. In UAI, pp. 202-211, 2016.
Daniel Freeman, Erik Frey, Anton Raichuk, Sertan Girgin, Igor Mordatch, and Olivier Bachem. Brax-a differentiable physics engine for large scale rigid body simulation. In NeurIPS Datasets and Benchmarks Truck, 2021.
Scott Fujimoto, Herke Hoof, and David Meger. Addressing function approximation error in actor-critic methods. In I,ML,, 2018.
Matteo Gallici, Mattie Fellows, Benjamin Ellis, Bartomeu Pou, Ivan Masmitja, Jakob Nicolaus Foerster, and Mario Martin. Simplifying deep temporal difference learning. arXiv preprint arXiv:2407.048) 2024.
Feng Gao, Liangzhi Shi, Shenao Zhang, Zhaoran Wang, and Wu. Adaptive-gradient policy optimization: Enhancing policy learning in non-smooth differentiable simulations. In ICML, 2024.
Moritz Geilinger, David Hahn, Jonas Zehnder, Moritz Bächer, Bernhard Thomaszewski, and Stelian Coros. Add: Analytically differentiable dynamics for multi-body systems with frictional contact. ACM Transactions on Graphics (TOG), 39(6): 15, 2020.
Ignat Georgiev, Krishnan Srinivasan, Jie Xu, Eric Heiden, and Animesh Garg. Adaptive horizon actor-critic for policy leamingg in contact-rich differentiable simulation. In ICML, 2024.
Andreas Griewank and Andrea Walther. Revolve: implementation of checkpointing for the reverse adjoint mode of computational differentiation. ACM Transactions on Mathematical Software (TOMS), 26(1):19 45, 2000.
Andreas Griewank and Andrea Walther. Evaluating derivatives: principles and techniques of algorithmic differentiation. SIAM, 2008.
Jiayuan Gu, Fanho Xiang, Xuanlin Li, Zhan Ling, Xiqiang Liu, Tongzhou Mu, Yihe Tang, Stone Tao, Xinyue Wei, Yunchao Yao, Xiaodi Yuan, Pengwei Xie. Zhiao Huang. Rui Chen, and Hao Su. Maniskill2: A unified benchmark for generalizable manipulation skills. In ICLR, 2023.
Huy Ha and Shuran Song. Flingbot: The unreasonable effectiveness of dynamic manipulation for cloth unfolding. In CoRL. 2022.
Tuomas Haarnoja. Haoran Tang. Pieter Abbeel, and Sergey Levine. Reinforcement learning with deep energy- based policies. In ICMI., 2017.
Tuomas Haarnoja, Aurick Zhou, Pieter Abhcel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor. In ICML,, 2018a.
Tuomas Haarnoja, Aurick Zhou, Kristian Hartikainen, George Tucker, Schoon Ha, Jie Tan, Vikash Kumar, Henry Zhu, Abhishek Gupta, Pieter Abbeel, and Sergey Levine. Soft actor-critic algorithms and applications. arXiv preprint 1812 2.09055 2018b.
Danijar Hafner, Timothy Lillicrap, Jimmy Ba, and Mohammad Norouzi. Dream control: Learning behaviors by latent imagination. In ICLR, 2020.
Seungyul Han and Youngchul Sung. max-min entropy framework for reinforcement learning. In NeurlPS, 2021.
Nicolas Heess, Gregory Wayne, David Silver, Timothy Lillicrap, Tom Erez, and Yuval Tassa. Learning continuous control policies by stochastic value gradients. In NeurIPS, 2015.
Eric Heiden, David Millard, Erwin Coumans, Yizhou Sheng, and Gaurav S Sukhatme. Neuralsim: Augmenting differentiable simulators with neural networks. In ICRA, 2021.
Eric Heiden, Miles Macklin, Yashraj Narang, Dieter Fox, Animesh Garg, and Fabio Ramos. Disect: differentiable simulator for parameter inference and control in robotic cutting. Autonomous Robots, 47()): 549-578, 2023.
12

---

Published as a conference paper at ICLR 2025
Dan Horgan, John Quan, David Budden, Gabriel Barth-Maron, Matteo Hessel, Hado van Hasselt, and David Silver. Distributed prioritized experience replay. In ICLR, 2018.
Taylor A Howell, Simon Le Cleac'h, Jan Brüdigam, Zico Kolter, Mac Schwager, and Zachary Manchester.
Dojo: A diffcrentiable physics engine for robotics. arXiv preprint arXiv:2203 00806, 2022.
Yuanming Hu, Yu Fang. Ziheng Ge, Ziyin Qu, Yixin Zhu, Andre Pradhana, and Chenfanfu Jiang. A moving least squares material point method with displacement discontinuity and two-way rigid body coupling. ACM Trunsactions Graphics (TOG), 37(4):1-14, 2018.
Yuanming Hu, Tzu-Mao Li, Luke Anderson, Jonathan Ragan-Kelley, and Frédo Durand. Taichi: language for high-performance computation on spatially sparse data structures. ACM Transactions on Graphics (TOG), 38 (6):1-16, 2019a.
Yuanming Hu, Jiancheng Liu, Andrew Spielberg, Joshua Tenenbaum, William Freeman, Jiajun Wu, Daniela Rus, and Wojciech Matusik. Chainqueen: A real-time differentiable physical simulator for soft robotics. In ICRA. 2019b.
Yuanming Hu, Luke Anderson, Tzu-Mao Li, Qi Sun, Nathan Carr, Jonathan Ragan-Kelley, and Fredo Durand.
Difftaichi: Differentiable programming for physical simulation. In ICIR, 2020.
Zhiao Huang, Yuamming Hu, Tao Du, Siyuan Zhou, Hao Su, Joshua B Tenenbaum, and Chuang Gan. Plas- ticinelab: soft-hody manipulation benchmark with differentiable physies. In ICIR,, 2020.
Jeffrey Ichnowski, Yahav Avigal, Yi Liu, and Ken Goldberg. Gomp-fit: Grasp-optimized motion planning for fast inertial transport. ICRA, 2022.
Hilbert Kappen. Path integrals and symmetry breaking for optimal control theory. Journal of statistical mechanicx: theory and experiment, 2005(11):PI 2005.
Elia Kaufmann, Leonard Bauersfeld, Antonio Loquercio, Matthias Müller, Vladlen Koltun, and Davide Scara- muzza. Champion-level drone racing using deep reinforcement learning. Nature, 620(7976):822 9877 2023.
Alex Kendall and Yarin Gal. What uncertaintics do we need in bayesian deep learning for computer vision? In NewrlPS, 2017.
Seungchan Kim, Kavosh Asadi, Michael Littman, and George Konidaris. Deepmellow: removing need for target network in deep q-learning. In IJCAI, 2019.
Diederik Kingma and Jimmy Ba. Adam: method for stochastic optimization. ICLR, 2015.
Diederik Kingma and Max Welling. Auto-encoding variational bayes. ICLR, 2014.
Vijay Konda and John Tsitsiklis. Actor-critic algorithms. In NewrIPS, 1999.
Sizhe Li, Zhiao Huang, Tao Du, Hao Su, Joshua B. Tenenbaum, and Chuang Gan. Contact points discovery for soft-body manipulations with differentiable physics. In ICLR, 2022a.
Sizhe Li, Zhiao Huang. Tao Chen, Tao Du, Hao Su, Joshua Tenenbaum, and Chuang Gan. Dexdeform: Dexterous deformable object manipulation with human demonstrations and differentiable physics. In ICLR, 2023a.
Yifei Li, Tao Du, Kui Wu, Jie Xu, and Wojeiech Matusik. Diffeloth: Differentiable cloth simulation with dry frictional contact. ACM Transactions on Graphics (TOG), 42(1):1-20, 2022b.
Zechu Li, Tao Chen, Zhang-Wor Wei Hong, Anurag Ajay, and Pulkit Agrawal. Parallel g-learning: Scaling off-policy reinforcement learning under massively parallel simulation. In ICMI., 2023b.
Jacky Liang, Viktor Makoviychuk, Ankur Handa, Nuttapong Chentanez, Miles Macklin, and Dicter Fox. Gpu-accelerated robotic simulation for distributed reinforcement learning. In CoRL, 2018.
Junbang Liang, Ming Lin, and Vladlen Koltun. Differentiable cloth simulation for inverse problems. In NeurlPS, 2019.
Timothy P Lillicrap. Jonathan J Hunt, Alexander Pritzel, Nicolas Heess. Tom Erez, Yuval Tassa, David Silver, and Daan Wierstra. Continuous control with deep reinforcement learning. ICLR, 2016.
Xingyu Lin, Yufei Wang, Jake Olkin, and David Held. Softgym: Benehmarking deep reinforcement learning for deformable object manipulation. CoRL, 2021.
Xingyu Lin, Yufei Wang, Zixuan Huang, and David Held. Learning visible connectivity dynamics for cloth smoothing. In CoRL, 2022.
13

---

Published as a conference paper at ICLR 2025
Xingyu Lin, Carl Qi. Yunchu Zhang. Zhiao Huang. Katerina Fragkiadaki, Yunzhu Li, Chuang Gan, and David Held. Planning with spatial-temporal abstraction from point clouds for deformable object manipulation. In CoRL, 2023.
Min Liu, Gang Yang. Siyuan Luo, and Lin Shao. Softmac: Differentiable soft body simulation with forecast- based contact model and two-way coupling with articulated rigid bodies and clothes. In IROS, 2024.
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. arXiv preprint arXiv: 1711.05001, 2017.
Pingchuan Ma, Peter Yichen Chen, Bolei Deng, Joshua B Tenenbaum, Tao Du, Chuang Gan, and Wojciech Matusik. Learning neural constitutive laws from motion observations for generalizable pde dynamics. In ICML, 2023.
Miles Macklin. Warp: A hegh-performance python framework for gpu simulation and graphics. https: //github. com/ avidia, diw/warp. March 2022. NVIDIA GPU Technology Conference (GTC).
Viktor Makoviychuk, Lukasz Wawrzyniak, Yunrong Guo, Michelle Lu, Kier Storey, Miles Macklin, David Hoeller, Nikita Rudin, Arthur Allshire, Ankur Handa, et al. Isaac gym: High performance gpu hased physics simulation for robot learning. In NeurlPS Datasets and Benchmarks Track, 2021.
Charles Marsh. Introduction to continuous entropy. Department of Computer Science, Princeton University, 1034, 2013.
Mayank Mittal, Calvin Yu, Qinxi Yu, Jingzhou Liu, Nikita Rudin, David Hoeller, Jia Lin Yuan, Ritvik Singh, Yunrong Guo, Hammad Mazhar, Orbit: A unified simulation framework for interactive robot learning environments. IEEE Robotics and Automation Letters, 86))3374-37747, 2023.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A Rusu, Joel Veness, Marc G Bellemare, Alex Graves, Martin Riedmiller, Andreas K Fidjeland, Georg Ostrovski, et al. Human-level control through deep reinforcement learning. Nature, 51 518(7540):529 533, 2015.
TM Moerland, DJ Brockens, Aske Plaat, and CM Jonker. Model-based reinforcement learning: survey. Foundations and Trends in Machine Learning, 16(1), 2023.
Ted Moskovitz, Jack Holder, Aldo Pacchiano, Michael Arbel, and Michael Jordan. Tactical optimism and pessimism for deep reinforcement learning. In NeurIPS, 2021.
Michael Mozer. focused backpropagation algorithm for temporal pattern recognition. In Backpropagation: theory, architectures, and applications, pp. 137-169. Psychology Press, 1995.
J. Krishna Murthy, Miles Macklin, Florian Golemo, Vikram Voleti, Linda Petrini, Martin Weiss, Breandan Consi- dine, Jérôme Parent- -évesque, Kevin Xie, Kenny Erleben, Liam Paull, Florian Shkurti, Derek Nowrouzezahrai, and Sanja Fidler. gradsim: Differentiable simulation for system identification and visuomotor control. In ICLR. 2021.
Arun Nair, Praveen Srinivasan, Sam Blackwell, Cagdas Alcicek, Rory Fearon, Alessandro De Maria, Vedavyas Pannecrshelvam, Mustafa Suleyman, Charles Beattie, Stig Petersen, et al. Massively parallel methods for deep reinforcement lcarning. arXiv preprint arXiv: 1507.0426, 2015.
Ashvin Nair, Dian Chen, Pulkit Agrawal, Phillip Isola, Picter Abbeel, Jitendra Malik, and Sergey Levine. Combining self-supervised learning and imitation for vision-based rope manipulation. In ICRA, 2017.
Rhys Newbury. Jack Collins, Kerry He. Jiahe Pan, Ingmar Posner, David Howard, and Akansel Cosgun. A review of differentiable simulators. IEEE Access, 2024.
Paavo Parmas, Carl Edward Rasmussen, Jan Peters, and Kenji Doya. Pipps: Flexible model-based policy search robust to the curse of chaos. JCML, 2018
Andrew Patterson, Samuel Neumann, Martha White, and Adam White. Empirical design in reinforcement learning. arXiv preprint arXhv:2304.01315, 2023.
Yi-Ling Qiao, Junbang Liang, Vladlen Koltun, and Ming Lin. Scalable differentiable physics for learning and control. In ICML, 2020.
Yi-ling Qiao, Junbang Liang, Vladlen Koltun, and Ming Lin. Differentiable simulation of soft multi body systems. In NewrlPS, 2021a.
Yi-Ling Qiao, Junhang Liang, Vladlen Koltun, and Ming C Lin. Efficient differentiable simulation of articulated bodies. In ICML., 202
Ilija Radosavovie, Tete Xiao, Bike Zhang, Trevor Darrell, Jitendra Malik, and Koushil Sreenath. Real-world humanoid Incomotion with reinforcement learning. Science Robotics, (89)eea),957,, 2024.
14

---

Published as a conference paper at ICLR 2025
Jiawei Ren, Cunjun Yu, Siwei Chen, Xiao Ma, Liang Pan, and Ziwei Liu. Diffmmimic: Efficient motion mimicking with differentiable physics. ICLR, 2023.
Danilo Jimenez Rezende, Shakir Mohamed, and Daan Wierstra. Stochastic backpropagation and approximate inference in deep generative models. In ICML, 2014.
Junior Rojas, Eftychios Sifakis, and Ladislav Kavan. Differentiable implicit soft-body physics, arXiv preprint arXiv:2102.0559) 2021.
Nikita Rudin, David Hoeller, Philipp Reist, and Marco Hutter. Learning to walk in minutes using massively
parallel deep reinforcement learning. In CoRL, 2021.
Peter Schaldenbrand, James McCann, and Jean Oh. Frida: A collaborative robot painter with differentiable. real2sim2real planning environment. In ICRA, 2023.
John Schulman, Nicolas Heess, Theophane Weber, and Picter Abbeel. Gradient estimation using stochastic computation graphs. NewrIPS, 2015a.
John Schulman, Philipp Moritz, Sergey Levine, Michael Jordan, and Pieter Abbeel. High figh-dimensional continuous control using generalized advantage estimation. arXiv preprint 1506.0550002438 2015b.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv: 170.006777 2017.
Clemens Schwarke, Victor Klemm, Jesus Tordesillas, Jcan-Pierre Sleiman, and Marco Hutter. Learning quadrupedal locomotion via differentiable simulation. arXiv preprint arXiv:2404.02888 2024.
Lin Shao, Yifan You, Mengyuan Yan, Shenli Yuan, Qingyun Sun, and Jeannette Bohg. Grae: Self- and self-regularized actor-critic. In CoRL, 2022.
Bokui Shen, Zhenyu Jiang, Christopher Choy, Leonidas J. Guibas, Silvio Savarese. Anima Anandkumar, and Yuke Zhu. Acid: Action-conditional implicit visual dynamics deformable object manipulation. RSS, 2022.
Haochen Shi, Huazhe Xu, Zhiao Huang, Yunzhu Li, and Jiajun Wu. Robocraft: Learning to see, simulate, and shape clasto- clasto-plastic objects with graph networks. In RSS, 2022.
Haochen Shi, Huazhe Xu, Samuel Clarke, Yunzhu Li, and Jiajun Wu. Robocook: Long- Long-horizon clasto- clasto-plastic
object manipulation with diverse tools. In CoRL, 2023.
Zilin Si, Gu Zhang, Qingwei Ben, Branden Romero, Zhou Xian, Chao Liu, and Chuang Gan. Difftactile: A physies-based differentiable tactile simulator for contact-rich robotic manipulation. In ICLR, 2024.
David Silver, Aja Huang, Chris J Maddison, Arthur Guez, Laurent Sifre, George Van Den Driessche, Julian Schrittwieser, Ioannis Antonoglou, Veda Panncershelvam, Mare Lanctot, et al. Mastering the game of go with deep neural networks and tree search. Nature, 529(75577:44444444 2016.
Breannan Smith, Fernando De Goes, and Theodore Kim. Stable neo-hookean flesh simulation. ACM Transactions on Graphics (TOG), 15,, 2018.
Sanghyun Son, Laura Zheng, Ryan Sullivan, Yi-Ling Qiao. and Ming Lin. Gradient informed proximal policy optimization. In NewrlPS, 2024.
Yunlong Song. Sangbae Kim, and Davide Scaramuzza. Learning quadruped locomotion using differentiable simulation. In CoRI. 2024.
Hyung Suh, Max Simchowitz, Kaiqing Zhang, and Russ Tedrake. Do diffcrentiable simulators give better policy gradients? In ICML,, 2022.
Richard S Sutton. Learning to predict by the methods of temporal differences. Machine learning, 3:9-44, 1988.
Richard S Sutton. Integrated architectures for lcarning, planning, and reacting based on approximating dynamic programming. In ICML, 1990.
Richard S Sutton, David McAllester, Satinder Singh, and Yishay Mansour. Policy gradient methods for reinforcement learning with function approximation. In NeurlPS, 1999.
Bingjie Tang, Michael A Lin, Iretiayo Akinola, Ankur Handa, Gaurav S Sukhatme, Fabio Ramos. Dieter Fox, and Yashraj Narang. Industrcal: Transferring contact-rich assembly tasks from simulation reality. In RSS, 2023.
Evangelos Theodorou, Jonas Buchli, and Stefan Schaal. generalized path integral control approach to reinforcement leamning. The Journal of Machine Learning Research, 11:3137-3181, 2010.
Philip Thomas. Bias in natural actor critic algorithms. In ICML, 2014.
15

---

Published as a conference paper at ICLR 2025
Sebastian Thrun and Anton Schwartz. Issues in using function approximation for reinforcement learning. In Proceedings of the 1993 Connectionist Models Summer School, pp. 255. Psychology Press, 2014,
Michalis Titsias and Miguel Lázaro- Lázaro-Gredilla, Doubly stochastic variational bayes for non-conjugate inference. In /CML, 2014.
Emanuel Todorov. Lincarly-solvable markov decision problems. NeurIPS, 2006.
Emanuel Todorov, Tom Erez, and Yuval Tassa, Mujoco: physics engine for model- based control. In IROS, 2012.
Marc Toussaint. Robot trajectory optimization using approximate inference. In ICML, 2009.
Dylan Turpin, Liquan Wang, Eric Heiden, Yun-Chun Chen, Miles Macklin, Stavros Tsogkas, Sven Dickinson, and Animesh Garg. Grasp'd d: Differentiable contact-rich grasp synthesis for multi-fingered hands, In ECCV, 2022.
Oriol Vinyals. Igor Babuschkin, Wojeiech M Czarnecki, Michaël Mathieu, Andrew Dudzik, Junyoung Chung, David Choi, Richard Powell, Timo Ewalds, Petko Georgiev, al. Grandmaster level starcraft using multi-agent reinforcement learning. Nature, 575(7782) 354, 2019.
Yian Wang. Juntian Zheng, Zhehuan Chen, Zhou Xian, Gu Zhang, Chao Liu, and Chuang Gan. Thin-shell object manipulations with differentiable physics simulations. In ICI.R, 2024.
Yixuan Wang, Yunzhu Li, Katherine Driggs-Campher Li Fei-Fei, and Jiajun Wu. Dynamic-resolutionn model learning for object pile manipulation. In RSS, 2023.
Paul J Werbos. Backpropagation through time: what it does and how to do it. Proceedings of the IEEE, 78(10): 1550-1560, 1990.
Keenon Werling, Dalton Omens, Jeongseok Lee, Ioannis Exarchos, and C. Karen Liu. Fast and feature-complete differentiable physics engine for articulated rigid bodies with contact constraints. RSS, 2021.
Erik Wijmans, Abhishek Kadian, Ari Morcos, Stefan Lee, Irfan Essa, Devi Parikh, Manolis Savva, and Dhruv Batra. Dd-ppo: Learning near-perfect pointgoal navigators from 2.5 billion frames. In ICLR, 2020.
Ronald Williams. Simple statistical gradient-following algorithms for connectionist reinforcement learning. Machine learning, 8:229-22 256, 1992.
Ronald Williams and David Zipser. Gradient- learning algorithms for recurrent networks and their computational complexity. In Backpropagation: theory, architectures, and applications, pp. 433-486 Psychology Press, 1995.
Jie Xu, Viktor Makoviychuk, Yashraj Narang, Fabio Ramos, Wojeiech Matusik, Animesh Garg, and Miles Macklin, Accelerated policy learning with parallel differentiable simulation. ICLR, 2021.
Shangjie Xue, Shuo Cheng, Pujith Kachana, and Danfei Xu. Neural field dynamics model for granular object piles manipulation. In CoRL, 2023.
Ge Yang, Anurag Ajay, and Pulkit Agrawal. Overcoming the spectral bias of neural value approximation. In ICLR, 2021.
Haonan Yu, Haichao Zhang, and Wei Xu. Do need the entropy reward (in practice)? arXiv preprint arXiv:2201. 122344 2022.
Miguel Zamora, Momehil Peychev, Sehoon Ha, Martin Vechey, and Stelian Coros. Pods: Policy optimization differentiable simulation. In ICML, 2021.
Yanjic Zc, Gu Zhang, Kangning Zhang, Chenyuan Hu, Muhan Wang, and Huazhe Xu. 3d diffusion policy: Generalizable visuomotor policy learning via simple 3d representations. RSS, 2024.
Shenao Zhang, Wanxin Jin, and Zhaoran Wang. Adaptive barrier smoothing for first-order policy gradient with contact dynamics. In JCML, 2023.
Yuang Zhang, Yu Hu, Yunlong Song, Danping Zou, and Weiyao Lin. Back to newton's laws: Learning vision-based agile flight via differentiable physics. arXiv preprint arXiv:2407.10645 2024.
Xian Zhou, Bo Zhu, Zhenjia Xu, Hsiao-Yt Yu Tung, Antonio Torralba, Katerina Fragkiadaki, and Chuang Gan. Fluidlab: A differentiable environment for benchmarking complex fluid manipulation. In ICLR, 2023.
Brian D Ziebart. Modeling purposeful adaptive behavior with the principle of maximum causal entropy. Carnegic Mellon University, 2010.
Brian Ziebart, Andrew L Maas, Andrew Bagnell, and Anind Dey. Maximum entropy inverse reinforcement learning. In AAAI, 2008
16

---

Published as a conference paper at ICLR 2025
A EXTENDED RELATED WORK
Differentiable simulators. We review prior works on non-parallel differentiable simulators for roboties and control problems. Degrave et al. (2019) demonstrate that gradient- based optimization can be used to train neural network controllers by implementing rigid-body physics engine in a modern auto-differentiation library to compute analytic gradients. Avila Belbute-Peres et al. 2018) show how to analytically differentiate through 2D rigid-body dynamics, by defining linear complementarity problem (LCP) and using implicit differentiation. Exploiting sparsity to differentiate through the LCP more efficiently, Nimble (Werling et al., 2021) presents differentiable rigid-body physics engine that supports complex 3D contact geometries. Dojo (Howell et al., 2022) implements a differentiable rigid-body physies engine with hard contact model (based nonlinear friction cones) via a nonlinear complementarity problem (NCP), and a custom interio-pooint solver to reliably solve the NCP. Also using implicit differentiation, Qiao et al. (2020) adopt meshes as object representations to support both rigid and deformable objects and adopt a sparse local collision handling scheme that enables coupling between objects of different materials. Geilinger al. (2020) presents unified frictional contact model for rigid and deformable objects, while analytically computing gradients through adjoint sensitivity analysis instead. DiffTaichi (Hu et al., 201 2020), used within several differentiable soft-body simulators (Hu et al., 2019b; Huang et al., 2020; Zhou et al., 2023; Si et al., 2024; Wang et al., 2024), introduces a differentiable programming framework for tape-based auto-differentiation using source-code transformation to generate kernel adjoints. We build Rewarped using NVIDIA Warp (Macklin, 2022), differentiable programming framework similar to DiffTaichi.
Model-based reinforcement learning. Generally, approaches in MBRL (Moerland et al., 2023) either assume a known rewards from data. We will focus the latter, learned models can be differentiated through for policy learning, similar to FO-MBRI. which computes analytic policy gradients by backpropagating rewards from some known differentiable model (ie. differentiable simulator). Value gradient methods estimate policy gradients by directly backpropagating through the learned model (Deisenroth et al., 2013; Heess et al., 2015; Feinberg al., 201 8; Parmas al., 2018; Clavera al., 2020; Hafner et al., 2020). Learned models can also be used for data augmentation by generating imagined trajectories from the learned model, as first proposed in DYNA (Sutton, 1990). In this paper, we consider FO-MBRL (Freeman et al., 2021; Xu et al., 2021) which does analytic policy optimization with data and gradients from a differentiable simulator, without explicitly learning a world model.
model or learn a world model by predicting dynamics and
Actoo
Critic
Actor MLP
Critic
Actor
Critic
M
T
T
T
Critic Encoder
Encoder
Encader
Encoder
Encoder
Encoder
5,
Jnri
DIff. Sim.
DIfT. Sim.
Diff. Sim.
Figure 4: Diagram of computational graph of SAPO. Note that J(T) only updates the actor, although gradients are still propagated through the critic.
17

---

Published as a conference paper at ICLR 2025
B SAPO ALGORITHM DETAILS
B. VISUAL ENCODERS DIFFERENTIABLE SIMULATION
We use separate visual encoders for the actor To (a and critic Vi (f)))): to enable learning on deformable tasks with high-dimensional (particle-based) visual inputs. To maintain differentiability to compute analytic gradients and reduced memory requirements, use downsampled particle states of the simulation as point clouds for high-dimensional visual observations. For runtime efficiency, PointNet-like DP3 encoder (Ze al., 2024) encode point cloud observation into lower-dimensional latent vector. We leave combining differentiable rendering (of RGB or depth image observations) with differentiable simulation, like in (Murthy al., 2021), future work.
B.2 PSEUDOCODE
Algorithm 1: Soft Analytic Policy Optimization (SAPO)
Initialize network parameters 0, Φ, this Ci
Model components
to 4- 0
Actor
To (ar(((((8)))
repeat
Actor encoder
f0(8t)
Create buffer for 1 to
Critic
VA, (fc.
H do a, ~ na (=)))(((s))
Critic encoder
(st)
Critic index
= 1
h4 Hr (1(s))
Hyperparameters
h4 (he + (217)
Horizon
H
1r, t enven. stevaste((..)
Entropy temperature
Ui41 < Vis (fcc
(0)
Target entropy
H
df then
TD trace decay
A
to
Discount
& Add data buffer
Learning rates
7J
B < BU [(8t, ar Tı, de h4 (141)))
Num critics
C
Mini-epochs
K
to + to + (H
Update actor using Eq. 17, with normalized
entropy he and mean values 1 C. Ei
(0, () (9, σ) 1) (0.0) (m.x))mmxeen Imaxent (T)
> Detach data from differentiable simulation autograd:
+ stoograd(())
Update entropy temperature using Eq. 21, with unnormalized entropy h4
< a-a- all E/ ah,, H)]
> Compute TD(A) value targets via Eq. using soft returns of Eq. 20, with normalized entropy h and min values min,=1...C (())
0, +....
for K updates do
Sample (st. i.)) B
> Update crities using Eq. 8 with clipped
soft value targets ;:
(Ws, (i) + (vr, (3) v V (W..))(((V)
until converged;
18

---

Published as a conference paper at ICLR 2025
C HYPERPARAMETERS
We run all algorithms on consumer workstations with NVIDIA RTX 4090 GPUs. Each run uses single GPU, which we run both the GPU-accelerated parallel simulation and optimization loop. We use a recent high-performance implementation standard model-free RL algorithms which has been validated for parallel simulation (Li et al., 2023b). We aim to common hyperparameter values among algorithms where applicable, such as for discount factor, network architecture, etc.
For TrajOpt, initialize single T-length trajectory of zero actions. This action is repeated across N = 16 parallel environments (N = 32 for [AntRun, HandReorient)). We optimize this trajectory for 50 epochs with horizon H of 32 steps. We Adam W the optimizer, with learning rate of 0.01, decay rates (2), = (0.7, 0.95), and gradient norm clipping of 0.5. For evaluation, we playback this single trajectory across parallel environments, each with different random initial states.
shared
PPO
SAC
APG
SHAC
SAPO
Num envs
64
Batch size
2048
Horizon H
32
Mini-epochs K
5
8
16
16
Discount
0.99
TD/GAE
0.95
0.95
0.95
Actor y
5e 4
56
50-4 4
2e-3
2e-3
2e 3
Critic y
5e
50 4
5e-4
4
Entropy 7
5e-3
-e-3
1) schedule
KL(0.008)
linear Adam (0.7, 0.95) 1.0
linear Adam (0.7.0.95) 1.0
linear
Optim type
AdamW
Optim (B1, B2)
(0.9, 0.999)
(0.7,0.95)
Grad clip
0.5
Norm type
LayerNorm
Act type
SiL.U
ELU no
ELU no
Actor (s)
yes
Actor log(o)
log 1.0]
[-5, 2]
[-5, 2]
Num critics C
2
2
Critic t
0.995
0.995
Replay buffer
106
Target entropy H.
dim(.A)22
dim (A)/2
Init temperature
1.0
1.0
Table Shared hyperparameters. Algorithms use hyperparameter settings in the shared column unless otherwise specified in individual column.
Hopper
Ant
Humanoid
SNUHumanoid
Actor ML.
(128, 64, 32)
(128, 64,32)
(256, 128)
(512, 256) (256, 256)
Critic MIP
(64, 64) (64, 64)
(128, 128)
Table 5: DFlex spec-fe hyperpararameters. All algorithms the same actor and critic
network architecture.
shared
AntRun
HandReorient
Num envs N
32
64
64
Batch size
1024
2048
2048
Actor MLP
(512, 256)
(128, 64, 32)
Critic MLP
(256, 128)
(64, 64)
Table 6: Rewarped task-specific hy peprararraeetrrs All algorithms use the same actor and critic network architecture. Algorithms use hyperparameter settings in the shared column unless otherwise specified in an individual column.
19

---

Published as a conference paper at ICLR 2025
D REWARPED PHYSICS
We review the simulation techniques used to simulate various rigid and soft bodies in Rewarped. Our discussion is based on (Xu et al., 2021; Heiden et al., 2023; Murthy et al., 2021; Macklin, 2022), as well as (Hu et al., 2019b; Huang et al., 2020; Ma et al., 2023) for MPM.
To backpropagate analytic policy gradients, we need to compute simulation gradients. Following reverse order, observe
T-1
VOJ(T) V
Σ
YT=
(22)
dr: or: da:
Ør. dsg
(23)
ae da. 00
0s d8
ds 0 f a²) ds:
0 0f(sc. az)
da:
+
(24)
d0 0s: dl Øs:
00
with a no(-(s)), 8111 f(s at).
(25)
We use the reparameterization trick to compute gradients @tale) for stochastic policy тт to obtain
3e
da, Note that the simulation gradient is @f(a) which we compute by the adjoint method through 80
0s:
auto-differentiation (AD).
While we could estimate the simulation gradient with finite differences (FD) through forms of (x(x) R [(*th)-f(a) FD has two major limitations. Consider f RV - R-M where N is the input dimension, M is the output dimension, and L is the cost of evaluating S. Then FD has time complexity O((N + 1) M L). In comparison, reverse-mode AD has time complexity O(M L). and also benefits from GPU-acceleration. Furthermore, FD only computes numerical approximations of the gradient with precision depending on the value of h, while AD yields analytic gradients. For simulation platform built on a general differentiable programming framework, computing gradients for different simulation dynamics f is straightforward using AD. While alternatives to compute analytic simulation gradients through implicit differentiation have been proposed (see Appendix A), they are less amenable to batched gradient computation for parallel simulation.
h
D.1 COMPOSITE RIGID BODY ALGORITIIM (CRBA)
Consider the following rigid body forward dynamics equation to solve
Mî = J" + cla. d) + T a),
(26)
where q. J. are joint coordinates, velocities, accelerations, FF are external forces, includes Coriolis forces, T are joint-space actuations, M is the mass matrix, and is the Jacobian. Featherstone's 's composite rigid body algorithm (CRBA) is employed to solve for articulation dynamics. After obtaining joint accelerations q, a semi-implicit Euler integration step is performed to update the system state (). We use the same softened contact model as (Xu et al., 2021).
D.2 FINITE ELEMENT METHOD (FEM)
To simulate dynamics, finite element model (FEM) employed based on tetrahedral discretization of the solid's mesh. A stable neo-Hookean constitutive model (Smith et al., 2018) is used to model elastic deformable solids with per-element actuation
(q, +) = 1/ (Ic - 3)
# > [J - - log(Te (Tc + 1);
(27)
where (A, (4) are the Lamé parameters which control each tetrehedral element's resistance to shear and strain, a a constant, J det(F) (F) is the relative volume change, Ic = tr tr(FTF)) and the deformation gradient. Integrating V over each tetrahedral element yields the total elastic potential energy. to then compute Falastic from the energy gradient, and finally update the system state using a semi-implicit Fuler integration step. We use the same approach as (Murthy et al., 2021).
20

---

Published as a conference paper at ICLR 2025
D.3 MATERIAL POINT METHOD (MPM)
The moving least squares material point method (MLS-MPM) (Hu et al., 2018) can efficiently simulate variety of complex deformables, such elastoplastics and liquids. Consider the two equations for conservation of momentum and mass
p6 V P+ pb
(28)
Dp DI -V-
(29)
where è is the velocity, is the acceleration, p is the density, and is the body force. For this system to be well-defined, constitutive laws must define P, see (Ma et al., 2023). We simulate elastoplastics using an elastic constitutive law with von Mises yield criterion:
P(F) = U2pu + Atr(E))UU
(30)
20y
87 = ||2||
(31)
2/1
F
8 87 0,
P(F) U exp(e
(32)
BY VT
where P(F) is projection back into the elastic region for stress that violates the yield criterion, and F = UEV* the singular value decomposition (SVD). We simulate fluids as weakly compressible,
using a plastic constitutive law with the fixed corotated elastic model
P(F) -J(J 1)F
(33)
P(F) = J31,
(34)
where P(F) is a projection back into the plastic region and J = det(F)) (F).
We the same softened contact model from PlasticineLab (Huang et al., 2020), with one-way coupling between rigid bodies to MPM particles. For any grid point with signed distance d to the nearest rigid body, we compute smooth collision strength min(exp(-adi 1). The grid point velocity before and after collision projection is linearly interpolated using Coupling is implemented by the Compatible Particle-In- In-Clll (CPIC) algorithm, see steps 2-4 of (Hu al., 2018).
21

---

Published as a conference paper at ICLR 2025
E REWARPED TASKS
We visualize each task Figure 1, and high-level task descriptions are provided in Section 6.
E.1 TASK DEFINITIONS
For each task, we define the observation space O, action space A, reward function R, termination d, episode length 7, initial state distribution po, and simulation method used for transition function P. We denote the goal state distribution p² the task defines it. We also report important physics hyperparameters, including both shared ones and those specific to the simulation method used.
We use i, & for standard unit vectors and proj Uy to denote the x--corddinate of vector u.
for the projection of u onto We use [w]
We use for the indicator function and U(b) = U (-b,b)) b) for the uniform distribution. We denote an additive offset -5 multiplicative scaling x 6, from some prior value.
For notation, we use z-axis up. Let q, q. q be joint coordinates, velocities, accelerations. Let (p, 0) be world positions and orientations, (u,w) w) be linear and angular velocities, (a, a) be linear and angular accelerations, derived for root links such as the center of mass (CoM).
units. We use for the frame time, and each frame computed using S substeps, the physics time step is At S.. The gravitational acceleration is g.
All physical quantities are reported
FEM. We x' to denote subsampled set of FEM particle positions æ. where ø' is the CoM of the subsampled particles (i.e. average particle position for uniform density). Similarly for velocities v. We also report the following quantities: number of particles Na. number of tetrahedron Niet- particles' density Lamé parameters (X, p4) and damping stiffness Kdamp-
MPM. We use x' to denote subsampled set of MPM particle positions where X' is the CoM of the subsampled particles (i.e. average particle position for uniform density). We also report the following quantities: friction coefficient for rigid bodies μs, grid size Ng. number of particles N.. particles' density p. Young's 's modulus E, Poisson's ratio v, yield stress dy.
22

---

Published as a conference paper at ICLR 2025
E.1.1 ANTRUN
R377
[p2, 0, v, w, q, Map Wheadings ac-1)
R absolute joint torques
Uz + R0 + Rheading + Rheight)
d
1(p= hterm)
1000
Pu
+ 6p U(.1)), +50 (24)
+ōg (0.2))
~~ (0.25)
CRBA
h
1/60
16
Table 7: AntRun task definition. Based on Ant from (Xu et al., 2021).
The derived quantities correspond to the ant's CoM. We use termination height herm 0.27. We denote Hup = projg and Wheading = proj; p. The reward function maximizes the forward velocity Ux,, with auxiliary reward terms Rap upp encourage vertical stability, Rreading Uhonding to encourage running straight, and Rheight = P> hterm to discourage falling. Initial joint angles, joint velocities, and CoM transform of the randomized.
23

---

Published as a conference paper at ICLR 2025
E.1.2 HANDREORIENT
R272 "full", see (Makoviychuk et al., 2021)
A
R16 absolute joint torques т
R
see (Makoviychuk et al., 2021)
d
see (Makoviychuk et al., 2021)
600
Po
see (Makoviychuk et al., 2021)
P
CRBA
1/120
32
Table 8: HandReorient task definition. Based on AllegroHand from (Makoviychuk al., 2021).
See (Makoviychuk al., 2021). Initial joint angles of the hand, CoM transform of the cube, and target orientations for the cube are randomized.
24

---

Published as a conference paper at ICLR 2025
E.1.3 ROLLINGFLAT
[ R2, R- [a', x', (Pas Ps. 0z)
A
R relative positions Pa.Pa Ps and orientation 0=
R
Ra + flat
d
300
Pu
8x ~ (0.9, 1.1)
Z (O01)), + tdz UO,, 0.05)
P
MPM
S 5e-5
S
40
μbb
0.9
N
48
N
2592
p
1.0
E
5000.
0.2
0y
50.
Table 9: RollingFlat task definition. Based on RollingPin from (Huang al., 2020).
The derived quantities correspond to the rolling pin's CoM. We use target flatten height f = 0.125. For the reward function, the first term Ra minimizes the difference between the particles' CoM and hiaa, while the second term maximizes the smoothness the particles R
d
hiflat
Rd 1+ d Var[x']]
[1d 0.33) 1+ (dd s 0.33) 2]
Rfiat
Initial volume and CoM transform of the particles are randomized.
25

---

Published as a conference paper at ICLR 2025
E.1.4 SOFTJUMPER
20243 R2, 1 22 : |x', x', v2, ac-1
A
R absolute tetrahedral volumetric activations (subsampled)
R
d'z + (3.0Rup 0.0001 Σ ||al
d
300
Po
dæ ~ (0.9, 1.1)
u (0.8), +ox, U 0.44
P
FEM
1/60
S
80
N
204
Niet
444
P
1.0
1000.
1000.
kdamp
1.0
Table 10: SoftJumper task definition. Inspired by (Murthy et al., 2021; Hu et al., 2020).
Denote the default initial height ho E. 0.2. The reward function maximizes the forward velocity v' of the quadruped's CoM, with two auxiliary reward terms Rup x'= ho to encourage jumping and the other to encourage energy-efficient policies by minimizing action norms. For the action space, we downsample by factor of 2 and upsample actions to obtain the full set of activations. Initial volume and CoM transform of the quadruped randomized.
26

---

Published as a conference paper at ICLR 2025
E.1.5 HANDFLIP
[ R3, [22 |,', gl
A
R24 relative joint angles
(
d
T
300
Po
dæ ~ (0.95, 1.05)
8 2,1
U (0.1), + oxz U (0,0.1)
P
MPM
h
S 5e-5
S
40
P
0.5
Ng
48
N
2500
P
1.0
E
4000.
0.2
By
130.
Table 11: HandFlip task definition. Based on Flip from (Li et al., 2023a).
We define [1(r)= clamp( Za=r" 1..0, 1.0) as the normalized improvement of from initial ro. The
To reward function maximizes the NI of Rtiip where the term Rytip = ||[x'|,, [='],,,]l l...ll, minimizes the average Euclidean distance between two mirrored halves hr, ha of the particles. Initial volume and CoM transform of the particles are randomized.
27

---

Published as a conference paper at ICLR 2025
E.1.6 FLUIDMOVE
25000 IR, 12 RR [x', x', P.p'
A
R$ relative positions p Rd + Repill
R
d
300
Po
6²
p¹ f- | [ppp, (0.2), 8p, (0.00.0.)
P
MPM
S 5e-4
S
40
μb
0.5
Ng
48
N
2880
P
1c3
E
le5
0.3
Table 12: FluidMove task definition. Based on Transport from (Lin et al., 2021).
For the reward function, the first term Ra minimizes the distance to the target position p", while the second term penalizes particles which are spilled outside the container
d = 1p 2
Ra = [1 > 0.023 + ((d < 0.02)
Repilt- see (Lin al., 2021)
Target position for the container randomized.
28

---

Published as a conference paper at ICLR 2025
F EXTENDED EXPERIMENTAL RESULTS
F.1 EXAMPLE SHAC GETTING STUCK LOCAL OPTIMA
We reproduce the original DFlex Ant results from SHAC (Xu et al., 2021), and in Figure we visualize individual runs for insight (Patterson et al., 2023). From this, we observe that one of the runs quickly plateaus to suboptimal policy after 1M steps and does not improve.
DFlex Ant
9k
Return
6k
3k
0-
0
1M
2M
3M
4M
Figure Example SHAC getting stuck in local minima. Episode return function of environment steps in DFlex Ant (AA C (R²). One run (colored in red) quickly plateaus after 1M steps and does not improve. random seeds.
F.2
RESULTS DFLEX LOCOMOTION
DFlex Hopper
DFlex Ant
DFlex Humanoid
DFlex SNUHumanoid
6k
9k1
Return
3k-
5k-
4.5k
3.5k
5M
10M
5M
10M
5M
10M
5M
10M
SAPO (ours)
SHAC
APG
SAC
PPO
Figure 6: DFlex locomotion training curves. Episode return as a function environment steps in DFlex Hopper (A C RR), Ant (A C RR²), Humanoid (A4 C TR?2),, and SNUHumanoid (AA C RR152) locomotion tasks. Mean and 95% CIs over 10 random seeds.
Hopper
Ant
Humanoid
SNUHumanoid
PPO
3155 + 30
3883 + 60
414 + 45
135 + 3
SAC
3833 + 50
3366 + 25
4628 + 120
846 + 44
APG
590 3
368 11
783 + 16
149 1
SHAC
4939 + 3
7779 - 70
8256 + 74
5755 = 67
SAPO (ours)
5060 = 18
8610 - 40
8469 =
6427 53
Table 13: DFlex locomotion tabular results. Evaluation episode returns for final policies after training. Mean and 95% CIs over 10 random seeds with 128 episodes per seed.
29

---

Published as a conference paper at ICLR 2025
F.3
SAPO ABLATIONS ON DFLEX LOCOMOTION
DFlex Hopper
DFlex Ant
DFlex Humanooi
DFlex SNUHumanoid
6k
10k
9k
7k
Return
0-
0
10M
SAPO (ours)
w/o Vise
w/o Hg and Vin
SHAC
Figure 7: SAPO ablations DFlex locomotion training curves. Episode return as function of environment steps in DFlex Hopper (A C I89), Ant (A RR), Humanoid (AA C " and SNUHumanoid (AA C IR 162) locomotion tasks. Mean and 95% CIs over 10 random seeds.
Hopper
Ant
Humanoid
SNUHumanoid
(avg %)
SAPO (ours)
5060 + 18
8610 + 42
8469 + 59
6427 + 52
6.8%
w/o Vsaft
4882
7729 +
8389 76
6392 54
2.7%
w/o Hr and Vooft
5036 + 2
7897 +: 30
7731
+ 91
6032 + 58
0.5%
SHAC
4939 + 7779 + 8256 76
5755 66
Table 14: SAPO ablations DFlex locomotion tabular results. Evaluation episode returns for final
policies after training. Mean and 95% CIs over 10 random seeds with 128 episodes per seed.
F.4
ADDITIONAL SAPO ABLATIONS FOR DESIGN CHOICES {III, IV, V)
HandFlip
HandFlip
(A%)
SAPO (ours)
SAPO (ours)
90
+172.7%
Return
w/o V
w/o V
71 - 3
+115.2%
50
w/o
w/o IV
64 + 3
+93.9%
w/o III
w/o III
49 + 3
+48.5%
SHAC
SHAC
33 = 3
3M
6M
Figure SAPO ablations (ext.) HandFlip train-
Table 15: SAPO ablations - (ext.) HandFlip
ing curves. Episode return function of environ- ment steps. Smoothed using EWMA with a = 0.99. Mean and 95% Cls over 10 random seeds.
tabular results. Evaluation episode returns for final policies after training. Mean and 95% Cls over 10 random seeds with 64 episodes per seed.
See Section 4.2 for descriptions these design choices.
30

---

Published as a conference paper at ICLR 2025
F.5 RUNTIME AND SCALABILITY OF REWARPED
We report all timings on a consumer workstation with am AMD Threadripper 5955WX CPU, NVIDIA RTX 4090 GPU, and 128GB DDR4 3200MHz RAM. We the scalability of Rewarped's parallel MPM implementation over an increasing number of environments, on the HandFlip task. See Appendix E. 1.5 for details on environment settings.
We run 5 full-length trajectories, and report the averaged runtime of forward simulation. As refer- ence, the non-parallel MPM implementation from (Li et al., 2023a) runs with single environment at 70 FPS using 3.51 GB of GPU memory (when N 2500, N₂ = 48%).
With one environment, Rewarped's s MPM implementation is 3x faster and uses 2/3 less memory compared to DexDeform's. s. Using 32 environments, Rewarped achieves 20x total speedup com- pared to DexDeform for simulating decreases because overhead is amortized across multiple environments that are simulated in same underlying physics scene. In Figure 9, we observe that the total FPS drastically increases until plateauing beyond N 32 environments. This occurs because GPU utilization is maxed out, more GPUs or a GPU with more CUDA cores would be needed to continue scaling.
HandFlip task. Furthermore, the GPU memory used per
HandFlip
1,500
20
GPU Memory (GB)
Total FPS
15
1,000
10
500
5
0
1 8
16
32
64
128
# of Envs.
Figure 9: Rewarped runtimes HandFlip curves. Total FPS (bluc) and GPU memory (green) shown. Total FPS averaged over 5.N full-length trajectories with random actions.
# of Envs. Total FPS
Relative Speedup GPU Memory (GB)
210
1.00
1.09
373
1.78
1.28
617
2.94
1.58
978
4.66
2.25
16
1,237
5.89
3.59
32
1,394
6.64
6.16
64
1,506
7.17
11.40
128
1,569
7.47
21.86
Table 16: Rewarped runtimes HandFlip tabular results. Total FPS averaged over 5N full-length trajectories with random actions.
31

---

Published as a conference paper at ICLR 2025
HandFlip
200
20
150
GPU Memory (GB)
Total FPS
15
100
10
50
0
16 # of Envs.
32
Figure 10: Rewarped runtimes -HanFFlip (N₂ = 30000) curves. Total FPS (blue) and GPU
memory (green) are shown. Total FPS averaged over 5.V full-length trajectories with random actions.
# of Envs.
Total FPS
Relative Speedup GPU Memory (GB)
85
1.00
1,50
118
1.39
2.20
150
1.76
3.12
8
169
1.99
5.37
172
2.02
10.22
32
179
2.11
18.71
Table 17: Rewarped runtimes HandFlip (N, 30000) tabular results. Total FPS averaged over 5N full-length trajectories with random actions.
In Table 18, we report training runtimes for algorithms on the HandFlip task, averaged over 10 random seeds. We also note that runtimes may be further optimized using PyTorch compile and CUDA graphs, but we did not implement this.
Algorithm
Runtime (hours)
PPO
1.46
SAC
1.67
APG
7.43
SHAC
7.66
SAPO (ours)
7.77
Table 18: Training runtimes HandFlip task. Algorithm training runtimes over 6M environment steps. Mean over 10 random seeds.
32

---

Published as a conference paper at ICLR 2025
F.6 VISUALIZATIONS OF TRAJECTORIES REWARPED TASKS
AntRun
4B
HandReorient
RollingFlat
Soft
HandFlip
FluidMove
time
Figure 11: Visualizations of trajectories from policies learned by SAPO in Rewarped tasks. The camera view fixed between different time steps.
33

---

Published as a conference paper at ICLR 2025
F.7 VISUALIZATIONS OF OPTIMIZATION LANDSCAPE
Following (Xu et al., 2021), we visualize loss surfaces to compare different algorithms and analyze the smoothness of the optimization landscape. We undiscounted episodic returns for the loss values, and apply symlog scale for visualization. We sample two filter-normalized random directions in policy parameter space and evaluate perturbations final policies after training. We compare algorithms trained on DFlex Ant (Figure 13) and Rewarped HandFlip (Figure 12), observing that loss surfaces for SAPO appear to be flatter and smoother, compared to those of APG SHAC.
APG
SHAC
SAPO
Figure 12: Loss surface comparison between algorithms DFlex Ant.
APG
SHAC
SAPO
Figure 13: Loss surface comparison between algorithms Rewarped HandFlip.
34


Reviews:
[{"review_id": "PdFGk7A9fa", "reviewer": "Reviewer_WBQk", "summary": "This manuscript introduces Soft Analytic Policy Optimization (SAPO), a learning framework for sequential policy optimization that combines reinforcement learning, particularly maximum entropy reinforcement learning, and differentiable simulation. The authors differentiate the soft state value function estimator and make a set of hyper-parameter and architecture modifications for a better and more stable algorithm with practical implementation. To demonstrate the effectiveness of their method, the authors implement six different simulated environments with highly heterogeneous physics, including elastic bodies, articulated bodies, and fluids. These environments are used for both locomotion and manipulation tasks with a highly efficient GPU-based implementation using Warp.", "questions": "1. The authors declare that the design choices have larger impact to SAPO than to SHAC: what is the reason behind it? Could the authors provide some insights?\n2. I am confused by the TrajOpt baseline. I would guess that the environments are randomly initialized, which makes the trajectory optimization statistically impossible to work since it only replays the mean of learned actions. However, I found it better than some of the baselines. I might need more explanation here.\n3. APG baseline also looks weird to me. Typically my experience with APG gives me impression that it instantly improves in a very short time and plateaus. However, the figures show it constantly improves. Is there a reason behind it? Or my impression is wrong?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 3, "strengths": "1. The paper is well-written. I did not find any mathematical mistakes or factual errors.\n2. The approach of differentiating the maximum entropy reinforcement learning is valid and I personally found it interesting.\n3. The results look promising with good visualizations.\n4. The release of nicely implemented simulation environments can be useful for the robotics community.", "weaknesses": "1. The implementation details of the simulated environments are missing. How is the robotic ant simulated and actuated? What are the Young's modulus and Poisson's ratio of the elastic body? And fluid? What about the reward function for each environment?\n2. I found the visual encoder part to be extremely confusing. Why a PointNet is used for visual encoder? Does the word \"visual\" mean point cloud instead of images here? Since the title of appendix A.1 is \"Learning Visual Encoders in Differentiable Simulation\", I would assume the visual encoder is learned from scratch. Why not use a pre-trained model?\n\nMinor typo:\n    Line 50: scaling -> scale", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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": "85tQ3VUWVn", "reviewer": "Reviewer_wHaW", "summary": "This paper introduces Rewarped, a differentiable multiphysics simulator, and SAPO (Soft Analytic Policy Optimization), a maximum entropy RL method designed to leverage differentiable simulation for efficient policy learning. SAPO uses Rewarped’s capabilities to compute on-policy, first-order policy gradients, maximizing a maximum entropy RL objective. Built on NVIDIA's Warp, Rewarped includes a parallelized Material Point Method (MPM)-based soft body simulator with two-way coupling, providing an advanced simulation environment that supports a variety of rigid and soft body manipulation tasks. Benchmark results show that SAPO outperforms current state-of-the-art methods across challenging locomotion and manipulation environments.", "questions": "1. A comprehensive ablation study would provide valuable insights into the effectiveness of different design choices. For instance, comparing the performance of SAPO with different activation functions (e.g., ELU vs. SiLU).\n2. I am curious about the improvements brought by the paralleled MPM. how much acceleration can be achieved with the paralleled MPM, compared with the version in DexDeform? How much GPU memory is required?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The paper is well-written and the ideas are presented clearly, making it accessible and easy to follow.\n2. The parallelized soft body simulator in Rewarped is a notable advancement. Integrating soft body manipulation environments within a unified Warp-based framework provides an excellent foundation for further research on differentiable physics and soft body RL.\n3.  SAPO's combination of on-policy soft policy learning with an off-policy critic is an effective method for differentiable physics-based RL. This design mitigates the complexity of offline policy gradient evaluation while still benefiting from the simulator’s analytical gradients.\n4. SAPO demonstrates strong performance on complex soft body manipulation tasks, such as dexterous hand manipulation, consistently outperforming baseline methods.\n5. The authors’ commitment to releasing the dataset and simulation code enhances the reproducibility and accessibility of their work.", "weaknesses": "1. While combining maximum entropy with differentiable physics is a promising approach, it may appear incremental, as similar techniques have been explored in prior RL literature.\n2. If I understand correctly, only one-way coupling is considered in MPM simulator. It would be nice if two way coupling could be added to allow impact back on the rigid body as in Maniskill2.\n3. The paper lacks a discussion of the efficiency gains from the parallelized MPM implementation. Including an analysis of wall-time efficiency or runtime comparisons would provide valuable insight", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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": "kXUuHgQuHQ", "reviewer": "Reviewer_W9i4", "summary": "This paper releases a new algorithm called SAPO which acts better on Soft rigid body dynamics and also a simulator, with the main contribution being sample efficient RL on these kinds of problems. SAPO is a model-based, actor-critic method that leverages analytic gradients from a differentiable simulator to optimize a stochastic policy with entropy maximization.  Rewarped is the simulator, a  differentiable multi-physics simulation platform capable of simulating soft bodies which usually have slow simulation times. They demonstrate that SAPO with Rewarped achieves superior performance on manipulation and locomotion tasks that require handling interactions between rigid structures and deformables outperforming baseline approaches.", "questions": "1. How does the simulation speed scale with number of particles in the soft body simulation ? A discussion on this can greatly help in understanding the applicability of Rewarped. A potential experiment of interest can be simulation speed vs number of particles.\n2. Similarly, what is the effect of number of parallel environments on the simulation speed ? Similar to the above experiment, a table/plot showing simulation speed vs num of parallel environments on a standard consumer grade GPU (eg. 4090)", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The paper is well motivated. The authors take inspiration from max entropy model free RL and extend this to SOTA first order model based RL. \n2. The contribution of a accelerated differentiable simulator for soft body dynamics is welcome.\n3. The paper is comprehensive in providing an ablation over components of their method and also provides explanations for their design choices.", "weaknesses": "1. Since this work is inspired by model-based reinforcement learning (MBRL) and maximum entropy RL, and the authors present a first-order MBRL algorithm based on the max-entropy RL framework, they should clearly distinguish their approach from prior work, such as this paper(https://arxiv.org/abs/2112.01195) and this one (https://ieeexplore.ieee.org/document/9892381). While it can be argued that these works use a world model rather than a differentiable simulation, this distinction should be explicitly discussed. A brief discussion on the potential inter-applicability of world model and gradient simulator-based methods could open up promising directions for future research. A comparison of the novelty, applications, and results of these works would strengthen the contribution of the paper.\n2. It is not directly intuitive as to why the review of Zeroth Order Batched Gradient is necessary in the main manuscript in Section 3 Equation 4. It is my understanding that the paper builds on top of first order model based RL methods instead. If this section is indeed required please elaborate on it's importance.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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": "Fvh73IVz3N", "reviewer": "Reviewer_cLkG", "summary": "This work proposes a first-order model-based reinforcement learning (FO-MBRL) method that combines differentiable simulation and entropy maximization. It also provides a new parallel differentiable simulator that supports various physical systems. The paper demonstrates the efficacy of its proposed FO-MBRL on typical test environments implemented in the new differentiable simulator and reports its performance gain over several representative baselines from reinforcement learning and differentiable simulation.", "questions": "I am generally happy with this work, and my major concern has been listed in the weakness section. Several very minor questions:\n- I am under the impression that there should have been a log function in Eqn. 4 (somewhat similar to Eqn. 2 in Georgiev 24). Could you confirm your Eqn. 4 is correct?\n- Eqn 12: where is (sl, al) in the inner expectation sampled from? Does this notation miss a probability distribution?\n- How is your simulation gradient backprogagated through the stochastic policy tanh (line 332) exactly?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "- The paper is well-organized and nicely written. In particular, I like the way the paper organizes the background section and the preamble in the experiment section;\n-  The discussion on the design decisions (Sec. 4.2) is informative for practitioners. The paper is upfront with these specific choices in the framework, which I really appreciate;\n- The quality of the simulation environments seems respectable, and the rendering is well-polished (Fig. 1). I am inclined to believe this differentiable simulator has the potential to be a major player in the future, but I’d like to check a few more details (See weaknesses).\n- The experimental results and the ablation study look convincing to me.", "weaknesses": "One contribution claimed by this work is the new differentiable simulator that simultaneously accommodates GPU parallelization, multiphysics, and differentiability.\n- I didn’t find detailed specifications of the environments, or maybe I missed them. The claim of this simulator is quite ambitious, and verifying this claim requires scrutinizing many technical details. In particular, I am looking for at least the following information: 1) the simulation problem size, including grid and particle numbers in MPM and the time step sizes; 2) the treatment of multiphysics coupling and their gradients, e.g., rigid-soft contact and solid-fluid coupling as shown in Fig. 1.\n- Table 1 and the related work should also discuss DiffTaichi, which I believe supports GPU parallelization, deformable solids and fluids, and gradient computation.\n- While I highly appreciate the engineering effort behind Rewarped, I don’t see much technical/algorithmic novelty in Sec. 5: the core simulation algorithms behind rigid/soft/fluids are from existing (differentiable-)MPM papers and GPU parallelization is offered by Warp. Grinding existing differentiable simulators in Table 1 and the related work section is a bit unnecessary :)", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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:

Decision: Accept (Spotlight)