How to Implement Differentiable Simulation with Warp Tape in Newton

Newton enables differentiable physics by running simulations inside a wp.Tape() context, computing a scalar loss from the final state, and calling tape.backward(loss) to automatically populate gradients on all Warp arrays marked with requires_grad=True.

Newton is a high-performance physics simulation library built on NVIDIA Warp that makes implementing differentiable simulation with Warp Tape straightforward. By leveraging Warp's automatic differentiation capabilities, Newton allows you to compute gradients through complex physics simulations for gradient-based optimization of control policies, material parameters, and initial conditions.

The Four-Step Differentiable Simulation Pipeline

Implementing differentiable simulation in Newton follows a consistent four-step workflow. This pipeline works across all solver types—including semi-implicit, implicit MPM, and Featherstone articulated body algorithms—because Newton's solvers in newton/_src/solvers.py use standard Warp kernels that the tape automatically records.

Step 1: Building a Differentiable Model with ModelBuilder

First, construct your simulation model and enable gradient tracking. In newton/_src/model.py, the ModelBuilder class provides the finalize() method which accepts a requires_grad parameter.

When requires_grad=True, Newton allocates gradient buffers (.grad fields) for all mutable buffers including particle positions, velocities, and material parameters.

import warp as wp
import newton

builder = newton.ModelBuilder()
builder.add_particle(
    pos=wp.vec3(0.0, 0.0, 0.0),
    vel=wp.vec3(1.0, 0.0, 0.0),
    mass=1.0
)

# Enable gradient tracking for all simulation buffers

model = builder.finalize(requires_grad=True)

Step 2: Initializing the Solver and Simulation State

Create a solver instance and allocate simulation states. Newton's solvers are agnostic to differentiability—they use the same API regardless of whether gradients are required. The differentiability comes from the Warp arrays and tape context, not from solver configuration.

# Initialize solver - no special configuration needed for gradients

solver = newton.solvers.SolverSemiImplicit(model)

# Allocate mutable simulation states

state_in = model.state()
state_out = model.state()
control = model.control()

Step 3: Recording the Forward Pass with wp.Tape

The core of differentiable simulation is the wp.Tape() context. Inside the with tape: block, every Warp kernel launch and every read/write operation on gradient-enabled arrays is recorded. This includes the solver's integration steps and any custom loss kernels.

Define a loss kernel that computes a scalar value from the simulation state:

@wp.kernel
def loss_kernel(
    particle_q: wp.array(dtype=wp.vec3),
    target: wp.vec3,
    loss: wp.array(dtype=float)
):
    delta = particle_q[0] - target
    loss[0] = wp.dot(delta, delta)

Execute the forward simulation and loss computation inside the tape context:

# Initialize loss array with gradient tracking

loss = wp.zeros(1, dtype=float, requires_grad=True)
target = wp.vec3(0.25, 0.0, 0.0)

tape = wp.Tape()
with tape:
    # Clear forces and step physics

    state_in.clear_forces()
    solver.step(state_in, state_out, control, None, 1.0/60.0)
    
    # Compute loss from final particle positions

    wp.launch(
        loss_kernel,
        dim=1,
        inputs=[state_out.particle_q, target],
        outputs=[loss],
    )

Step 4: Backpropagating Gradients and Accessing Results

After recording the forward pass, call tape.backward(loss) to trigger reverse-mode automatic differentiation. This populates the .grad fields of all arrays that participated in the computation with requires_grad=True.

# Backpropagate gradients from loss to all inputs

tape.backward(loss)

# Access gradient of initial velocity

velocity_gradient = state_in.particle_qd.grad
print(velocity_gradient)

These gradients can then be fed into any gradient-based optimizer, such as warp.optim.SGD or warp.optim.Adam, to optimize initial conditions, control sequences, or material parameters.

Real-World Example: Optimizing Soft Body Material Parameters

For practical applications, you often need to optimize material parameters rather than just initial states. The example in newton/examples/diffsim/example_diffsim_soft_body.py demonstrates end-to-end optimization of tetrahedral mesh material parameters using Newton's differentiable pipeline.

This implementation creates a training loop that wraps the forward-backward pass:

class SoftBodyOptimizer:
    def __init__(self):
        # Initialize model and solver

        self.model = self.create_model()
        self.solver = newton.solvers.SolverSemiImplicit(self.model)
        
        # Allocate simulation states for entire trajectory

        self.states = [self.model.state() for _ in range(self.sim_steps)]
        self.control = self.model.control()
        
        # Material parameters to optimize (Young's modulus, Poisson's ratio)

        self.material_params = wp.array(
            self.model.tet_materials.numpy()[0, :2].flatten(),
            dtype=wp.float32,
            requires_grad=True,
        )
        
        # Warp optimizer for gradient descent

        self.optimizer = warp.optim.SGD([self.material_params], lr=1e7)
    
    def capture(self):
        tape = wp.Tape()
        with tape:
            # Forward simulation

            self.solver.step(
                self.states[0], 
                self.states[-1],
                self.control, 
                None, 
                self.sim_dt
            )
            
            # Compute center of mass

            wp.launch(
                com_kernel, 
                dim=self.model.particle_count,
                inputs=[self.states[-1].particle_q], 
                outputs=[self.com]
            )
            
            # Compute loss against target position

            wp.launch(
                loss_kernel,
                dim=1,
                inputs=[self.target, self.com, self.pos_error, self.loss],
                outputs=[]
            )
        
        # Backpropagate and optimize

        tape.backward(self.loss)
        self.optimizer.step()

This pattern—allocating a wp.Tape, running the simulation forward, computing a loss, and calling tape.backward()—is the standard approach for implementing differentiable simulation with Warp Tape in Newton across all solver types.

Key Implementation Files in Newton

Understanding the source code structure helps when debugging gradient flow or extending functionality:

• newton/solvers.py – Public API documentation and high-level solver interfaces that demonstrate the differentiable workflow.

• newton/_src/solvers.py – Core implementation of SolverBase, SolverSemiImplicit, and other integration schemes. All solvers use standard Warp kernels, making them automatically compatible with wp.Tape().

• newton/_src/model.py – Contains ModelBuilder and the finalize(requires_grad=True) logic that allocates gradient buffers for particle positions, velocities, and material parameters.

• newton/examples/diffsim/example_diffsim_soft_body.py – Complete working example of material parameter optimization using the Warp Tape workflow.

• newton/tests/test_softbody.py – Unit tests validating gradient correctness through finite differences and analytical comparisons.

Summary

Implementing differentiable simulation with Warp Tape in Newton follows a consistent pattern across all physics solvers:

• Enable gradients by passing requires_grad=True to builder.finalize() when constructing your model.
• Use standard solvers like SolverSemiImplicit without modification—they automatically record to Warp Tape.
• Wrap simulation steps in a wp.Tape() context to record all kernel operations and array accesses.
• Compute scalar losses using custom Warp kernels that operate on simulation states.
• Backpropagate by calling tape.backward(loss) to populate .grad fields on all participating arrays.
• Optimize parameters using Warp's built-in optimizers like warp.optim.SGD or warp.optim.Adam.

This architecture allows you to compute gradients through complex contact dynamics, soft body deformations, and articulated rigid body systems without manual derivative calculations.

Frequently Asked Questions

What is Warp Tape and how does it enable automatic differentiation in Newton?

Warp Tape is NVIDIA Warp's mechanism for recording computational graphs during the forward pass of a simulation. When you instantiate wp.Tape() and execute code within its context, Warp records every kernel launch and array operation involving gradient-enabled buffers. Calling tape.backward(loss) then traverses this recorded graph in reverse, applying the chain rule to compute gradients for all inputs. In Newton, this means any simulation step—whether rigid body, soft body, or MPM—automatically becomes differentiable without code modification.

Do I need to modify Newton's solver code to support gradient computation?

No. Newton's solvers in newton/_src/solvers.py are designed to be agnostic to differentiability. Whether you use SolverSemiImplicit, SolverImplicitMPM, or Featherstone articulated body algorithms, the solvers invoke standard Warp kernels for integration. Since Warp Tape records all kernel operations automatically, the same solver code works for both forward-only and differentiable simulations. You only need to ensure your model is built with requires_grad=True and that you wrap the solver steps in a wp.Tape() context.

How do I optimize material parameters instead of initial states in Newton?

To optimize material parameters—such as Young's modulus or Poisson's ratio—you must ensure these parameters are stored in Warp arrays with requires_grad=True. In Newton, material properties are typically stored in the model's tetrahedral or particle buffers. You can extract these into optimizable arrays, as shown in newton/examples/diffsim/example_diffsim_soft_body.py, where self.material_params is created from model.tet_materials with gradient tracking enabled. Pass this array to a Warp optimizer like warp.optim.SGD, run the forward simulation inside wp.Tape(), backpropagate the loss, and call optimizer.step() to update the material properties based on the computed gradients.

Where can I find working examples of differentiable simulation in Newton?

The Newton repository includes several reference implementations. The primary example is newton/examples/diffsim/example_diffsim_soft_body.py, which demonstrates end-to-end optimization of soft body material parameters to match a target shape. For validation and testing, newton/tests/test_softbody.py contains unit tests that verify gradient correctness through finite difference comparisons. Additionally, the documentation in newton/solvers.py (lines 58-66) provides a minimal working example of a differentiable rollout with a custom loss kernel, showing the exact pattern for recording simulation steps and backpropagating gradients.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}