[BUG] Scene remains still with `sim_substeps` being an odd number in `example_soft_body` when using CUDA Graph

[ricetwice]

Bug Description

Description:
I've encountered an issue where the example_soft_body example in example.sim remains in its initial state and does not move when sim_substep is set to an odd number. I suspect that this bug is related to CUDA Graph because everything works correctly when CUDA Graph is not used. Additionally, I encountered the same issue when writing my own code with CUDA Graph, based on the examples in example.sim.

While this issue is specifically observed in example_soft_body, other examples might also be affected, but I haven't extensively tested them.

Steps to Reproduce:

1. Set sim_substep to an odd number in example_soft_body.py.
2. Run the example.
3. Observe that the simulation does not progress and remains in the initial state.

System Information

Warp version: 1.2.2
CUDA version: 12.5
GPU: NVIDIA GeForce GTX 1060
OS: Ubuntu 22.04
Python version: 3.10

[nvlukasz]

Hi @ricetwice, thanks for reporting this issue. Looks like it's related to state swaps in the Example.simulate() method:

# swap states
(self.state_0, self.state_1) = (self.state_1, self.state_0)

To support odd substep counts, one solution is to use two graphs. One that starts with the original state_0 and a second one that starts with state_1, then alternate between them in Example.step().

We'll submit a fix shortly, but in the meantime this is one possible solution:

In Example.__init__(), replace:

        if self.use_cuda_graph:
            with wp.ScopedCapture() as capture:
                self.simulate()
            self.graph = capture.graph

with

        if self.use_cuda_graph:
            # capture two graphs to account for state swaps in simulate()
            # this supports both even and odd substep counts
            with wp.ScopedCapture() as capture_0:
                self.simulate()
            with wp.ScopedCapture() as capture_1:
                self.simulate()
            self.graph_0 = capture_0.graph
            self.graph_1 = capture_1.graph

In Example.step(), replace:

            if self.use_cuda_graph:
                wp.capture_launch(self.graph)

with:

            if self.use_cuda_graph:
                # alternate betwen the two captured graphs to account for even or odd substep counts
                wp.capture_launch(self.graph_0)
                (self.graph_0, self.graph_1) = (self.graph_1, self.graph_0)
                # if the substep count is odd, we need to swap states so that the latest appears as state_0
                if self.sim_substeps % 2 > 0:
                    (self.state_0, self.state_1) = (self.state_1, self.state_0)