# How to Implement Differentiable Simulation with Warp Tape in Newton

> Learn how to implement differentiable simulation with Warp Tape in Newton. Run simulations within wp.Tape() for automatic gradient computation on Warp arrays.

- Repository: [Newton Physics/newton](https://github.com/newton-physics/newton)
- Tags: how-to-guide
- Published: 2026-03-19

---

**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`](https://github.com/newton-physics/newton/blob/main/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`](https://github.com/newton-physics/newton/blob/main/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.

```python
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.

```python

# 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:

```python
@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:

```python

# 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`.

```python

# 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`](https://github.com/newton-physics/newton/blob/main/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:

```python
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`](https://github.com/newton-physics/newton/blob/main/newton/solvers.py)** – Public API documentation and high-level solver interfaces that demonstrate the differentiable workflow.

- **[`newton/_src/solvers.py`](https://github.com/newton-physics/newton/blob/main/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`](https://github.com/newton-physics/newton/blob/main/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`](https://github.com/newton-physics/newton/blob/main/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`](https://github.com/newton-physics/newton/blob/main/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`](https://github.com/newton-physics/newton/blob/main/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`](https://github.com/newton-physics/newton/blob/main/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`](https://github.com/newton-physics/newton/blob/main/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`](https://github.com/newton-physics/newton/blob/main/newton/tests/test_softbody.py) contains unit tests that verify gradient correctness through finite difference comparisons. Additionally, the documentation in [`newton/solvers.py`](https://github.com/newton-physics/newton/blob/main/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.