[BUG] Missing rigid body update in VBD Integrator

[Hantao-Ye]

Bug Description

Hi,

Firstly, I want to express my sincerest gratitude for your great work in developing warp and bringing VBD Integrator into the community. I have been playing with the newest VBD integrator release at 1.6.0 and found it really stable for simulating deformable cloth.

But I found some issues when using it:

1. Modelbuilder.color() will report an indexing error when there is no trimesh graph for coloring in the scene
2. Missing rigid body integration (for vel/pos update) in VBD integrator

To repeat the issue, I follow the example_cloth_self_contact.py to create a similar env:

import numpy as np
import warp as wp

from warp.sim import ModelBuilder, VBDIntegrator
from warp.sim.render import SimRenderer


class VBDEnv:
    def __init__(
        self,
        num_frames: int = 60,
        stage_path: str = None,
    ):
        fps = 60
        self.frame_dt = 1.0 / fps
        self.num_substeps = 10
        self.iterations = 10
        self.dt = self.frame_dt / self.num_substeps

        self.num_frames = num_frames
        self.sim_time = 0.0
        self.profiler = {}

        builder = ModelBuilder(up_vector=wp.vec3(0.0, 0.0, 1.0))

        box_pos_x = 2.8e-1
        box_pos_y = 2.4e-1
        box_pos_z = 1.0e-1

        box_x_length = 2.0e-1
        box_y_length = 1.8e-1
        box_z_length = 4.0e-2

        box_mass = 1.0
        box_pos = wp.vec3(box_pos_x, box_pos_y, box_pos_z)
        box_inertia = (
            1
            / 12
            * box_mass
            * np.array(
                [
                    [box_y_length**2 + box_z_length**2, 0, 0],
                    [0, box_x_length**2 + box_z_length**2, 0],
                    [0, 0, box_x_length**2 + box_y_length**2],
                ]
            )
        )

        box_id = builder.add_body(
            origin=wp.transform(p=box_pos, q=wp.quat_identity()),
            com=wp.vec3(0.0, 0.0, 0.0),
            I_m=wp.mat33(box_inertia),
            m=box_mass,
        )
        builder.add_shape_box(
            body=box_id,
            hx=box_x_length / 2.0,
            hy=box_y_length / 2.0,
            hz=box_z_length / 2.0,
            ke=1.0e4,
            kd=1.0e2,
            kf=1.0e1,
        )

        # Add a cloth for vbd integrator
        # cloth_pos_x = 2.8e-1
        # cloth_pos_y = 2.4e-1
        # cloth_pos_z = 4.0e-1

        # cloth_x_length = 1.0e-1
        # cloth_y_length = 1.0e-1

        # builder.add_cloth_grid(
        #     pos=wp.vec3(cloth_pos_x, cloth_pos_y, cloth_pos_z),
        #     rot=wp.quat_identity(),
        #     vel=wp.vec3(0.0, 0.0, 0.0),
        #     dim_x=10,
        #     dim_y=10,
        #     cell_x=cloth_x_length / 10,
        #     cell_y=cloth_y_length / 10,
        #     mass=1.0,
        #     particle_radius=1.0e-2,
        # )

        builder.color()

        self.model = builder.finalize()
        self.model.ground = True
        self.model.soft_contact_ke = 1.0e4
        self.model.soft_contact_kd = 1.0e2
        self.model.soft_contact_mu = 0.2

        self.integrator = VBDIntegrator(
            self.model, self.iterations
        )

        self.state0 = self.model.state()
        self.state1 = self.model.state()

        if stage_path:
            self.renderer = SimRenderer(self.model, stage_path, scaling=400.0)
        else:
            self.renderer = None

    def step(self):
        with wp.ScopedTimer("step", print=False, dict=self.profiler):
            wp.sim.collide(self.model, self.state0)
            for _ in range(self.num_substeps):
                self.state0.clear_forces()

                self.integrator.simulate(self.model, self.state0, self.state1, self.dt)

                (self.state0, self.state1) = (self.state1, self.state0)

                self.sim_time += self.dt

    def render(self):
        if self.renderer is None:
            return

        with wp.ScopedTimer("render", print=False):
            self.renderer.begin_frame(self.sim_time)
            self.renderer.render(self.state0)
            self.renderer.end_frame()


if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser(
        formatter_class=argparse.ArgumentDefaultsHelpFormatter
    )
    parser.add_argument(
        "--device", type=str, default=None, help="Override the default Warp device."
    )
    parser.add_argument(
        "--stage_path",
        type=lambda x: None if x == "None" else str(x),
        default="test.usd",
        help="Path to the output USD file.",
    )
    parser.add_argument(
        "--num_frames", type=int, default=60, help="Total number of frames."
    )

    args = parser.parse_known_args()[0]

    with wp.ScopedDevice(args.device):
        example = VBDEnv(stage_path=args.stage_path, num_frames=args.num_frames)

        for i in range(example.num_frames):
            example.step()
            example.render()
            print(f"[{i:4d}/{example.num_frames}]")

        frame_times = example.profiler["step"]
        print(
            "\nAverage frame sim time: {:.2f} ms".format(
                sum(frame_times) / len(frame_times)
            )
        )

        if example.renderer:
            example.renderer.save()

By running the provided python file, it will fail to run by reporting the following traceback, which is the first issue I mentioned:

Traceback (most recent call last):
  File "<frozen runpy>", line 198, in _run_module_as_main
  File "<frozen runpy>", line 88, in _run_code
  File "test_env.py", line 149, in <module>
    example = VBDEnv(stage_path=args.stage_path, num_frames=args.num_frames)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "test_env.py", line 85, in __init__
    builder.color()
  File ".venv/lib/python3.12/site-packages/warp/sim/model.py", line 4483, in color
    self.particle_coloring = color_trimesh(
                             ^^^^^^^^^^^^^^
  File ".venv/lib/python3.12/site-packages/warp/sim/graph_coloring.py", line 176, in color_trimesh
    graph_edge_indices = wp.array(trimesh_edge_indices[:, 2:], dtype=int, device="cpu")
                                  ~~~~~~~~~~~~~~~~~~~~^^^^^^^
IndexError: too many indices for array: array is 1-dimensional, but 2 were indexed

After adding the cloth to the environment, the program can run, but the box shape added to the scene holds its initialized position instead of falling during the simulation. I have checked the VBDIntegrator implementation, it feels like it computes force and applied integrations only on particles (within the colored graph), but the force/position on the rigid body is not computed and updated.

Please correct me if the previous suspects are wrong. If my guess is true, is there any plan for integrating rigid body dynamics in your developing timeline. I would appreciate your clarification with it.

System Information

$ lsb_release -a
Distributor ID: Ubuntu
Description:    Ubuntu 22.04.5 LTS
Release:        22.04
Codename:       jammy

$ python3 --version
Python 3.12.8

$ pip list | grep warp
warp-lang           1.6.0

[Hantao-Ye]

The color function will also report the similar traceback if replace the cloth with a rectangular tetrahedral FEM grid (created by function add_soft_grid).