Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Accelerating Visual-Policy Learning through Parallel Differentiable Simulation

Haoxiang You, Yilang Liu, Ian Abraham

TL;DR

This work tackles the high computational cost of learning visual control policies by introducing D.Va, a decoupled visual-based analytical policy gradient that omits differentiating through rendering and leverages parallel differentiable simulation. By decomposing the analytical policy gradient into a decoupled term and a control-regularization term, and linking the decoupled gradient to open-loop trajectory distillation, the method combines trajectory optimization with policy learning in a scalable on-policy framework. Empirical results show substantial improvements in wall-clock time and final returns across both 2D and 3D visual control tasks, including humanoid locomotion where a running policy emerges within hours on a single GPU. The approach outperforms standard model-free and model-based baselines and avoids reliance on differentiable renderers, providing a robust and memory-efficient path toward end-to-end visual policy learning in simulated environments. Overall, D.Va demonstrates that decoupling perception from policy computation and distilling from short-horizon trajectory optimization yields practical gains for visual control with differentiable simulators.

Abstract

In this work, we propose a computationally efficient algorithm for visual policy learning that leverages differentiable simulation and first-order analytical policy gradients. Our approach decouple the rendering process from the computation graph, enabling seamless integration with existing differentiable simulation ecosystems without the need for specialized differentiable rendering software. This decoupling not only reduces computational and memory overhead but also effectively attenuates the policy gradient norm, leading to more stable and smoother optimization. We evaluate our method on standard visual control benchmarks using modern GPU-accelerated simulation. Experiments show that our approach significantly reduces wall-clock training time and consistently outperforms all baseline methods in terms of final returns. Notably, on complex tasks such as humanoid locomotion, our method achieves a $4\times$ improvement in final return, and successfully learns a humanoid running policy within 4 hours on a single GPU.

Accelerating Visual-Policy Learning through Parallel Differentiable Simulation

TL;DR

This work tackles the high computational cost of learning visual control policies by introducing D.Va, a decoupled visual-based analytical policy gradient that omits differentiating through rendering and leverages parallel differentiable simulation. By decomposing the analytical policy gradient into a decoupled term and a control-regularization term, and linking the decoupled gradient to open-loop trajectory distillation, the method combines trajectory optimization with policy learning in a scalable on-policy framework. Empirical results show substantial improvements in wall-clock time and final returns across both 2D and 3D visual control tasks, including humanoid locomotion where a running policy emerges within hours on a single GPU. The approach outperforms standard model-free and model-based baselines and avoids reliance on differentiable renderers, providing a robust and memory-efficient path toward end-to-end visual policy learning in simulated environments. Overall, D.Va demonstrates that decoupling perception from policy computation and distilling from short-horizon trajectory optimization yields practical gains for visual control with differentiable simulators.

Abstract

In this work, we propose a computationally efficient algorithm for visual policy learning that leverages differentiable simulation and first-order analytical policy gradients. Our approach decouple the rendering process from the computation graph, enabling seamless integration with existing differentiable simulation ecosystems without the need for specialized differentiable rendering software. This decoupling not only reduces computational and memory overhead but also effectively attenuates the policy gradient norm, leading to more stable and smoother optimization. We evaluate our method on standard visual control benchmarks using modern GPU-accelerated simulation. Experiments show that our approach significantly reduces wall-clock training time and consistently outperforms all baseline methods in terms of final returns. Notably, on complex tasks such as humanoid locomotion, our method achieves a improvement in final return, and successfully learns a humanoid running policy within 4 hours on a single GPU.
Paper Structure (54 sections, 1 theorem, 25 equations, 17 figures, 11 tables, 1 algorithm)

This paper contains 54 sections, 1 theorem, 25 equations, 17 figures, 11 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Background
  3. Problem formulation
  4. Analytical policy gradient
  5. Short-horizon actor critic (SHAC)
  6. Method
  7. Decoupled policy gradient
  8. Distilling from open-loop Trajectories
  9. Control regularization
  10. Experimental validation
  11. Decoupled visual based analytical policy gradient
  12. Experiment
  13. Comparison to RL methods
  14. Baseline
  15. Results
  16. ...and 39 more sections

Key Result

Theorem 1

The decoupled trajectory gradient in eq: decoupled policy gradient equals the negative gradient of the behavior cloning loss in Equation eq: behavior cloning loss, i.e., $\tilde{\nabla}_{\boldsymbol{\theta}} \mathcal{J}(\mathbf{s}_0, \boldsymbol{\theta}, \mathcal{E}) = -\nabla_{\boldsymbol{\theta}}

Figures (17)

  • Figure 1: Fast training of humanoid running policy from pixel input. Our method learns a stable running gait in 4 hours on a single RTX 4080 GPU.
  • Figure 2: Computation graphs of decoupled policy gradient (DPG) and APG. The policy gradient is traced from any reward $r_t$, propagated backward through the graph to the shared parameters $\boldsymbol{\theta}$. APG backpropagates through the entire pipeline, whereas the DPG prevents gradient flow through the rendering process.
  • Figure 3: Comparison of APG $\nabla_{\boldsymbol\theta} \mathcal{V}$ and DPG $\tilde{\nabla}_{\boldsymbol\theta} \mathcal{V}$ on Hopper with full state observation. Both gradients are computed from the same set of $\boldsymbol\theta$ values collected during SHAC training. The x-axis shows the return for each $\boldsymbol\theta$. Left: Cosine similarity is generally positive, indicating $\tilde{\nabla}_{\boldsymbol\theta} \mathcal{V}$ is a valid ascent direction. Right: Gradient norm between APG and DPG. In this experiment, conducted in state space where $g=\texttt{identity}$, the control regularization term $\mathcal{B}$ acts as a residual connection. As a result, APG generally exhibits a smaller norm compared to DPG.
  • Figure 4: Comparison with RL: our method achieves substantial speedups and significantly higher rewards across all tasks. Each curve shows the average performance over five random seeds, with shaded areas indicating standard deviation. In the humanoid task, dashed lines represent the final rewards attained by each algorithm at the end of training. The top row highlights wall-clock efficiency; the bottom row illustrates sample efficiency, with curves truncated at the maximum number of simulation steps for better visualization.
  • Figure 5: Training time for Ant over 1M steps. Left: phase percentages, where "planning" in DreamerV3 refers to rollouts by learned world model. Right: absolute times used per 1M steps. Most time in visual RL is spent fitting neural networks.
  • ...and 12 more figures

Theorems & Definitions (2)

  • Theorem 1
  • proof