Providing a template class for Simulations

[JonathanKuelz]

Hi all,

I have started working with warp a couple months ago, and I remember that in the beginning, setting up just the framework of a simulation has already been a significant hurdle. I believe this effort can be reduced by providing a template class for simulations.

What do I mean by that? A lot of simulations share a similar pattern:

1. Construct a builder
2. Finalize the builder
3. Given number of frames and substeps, compute dt, simulation time, and other derived attributes
4. Set up initial states and potentially a controller
5. Set up the integrator and renderer
6. (Optional): Set up an optimizer
7. Set up graph capture
8. For n frames, run m steps each
9. (Optional) In each step/frame/simulation, compute a loss and optimize.

While some simulations may deviate from this pattern, I noticed that I was writing a lot of duplicate code in the beginning. I ended up writing a WarpSimulation template class for rigid body simulations that formalizes steps 2-9 from the above structure and a bit more.

I would be happy to share that code if there's interest in adding something like that to the warp codebase, but I could not spend a significant amount of time on extending or maintaining it for a longer time. Do you see the benefit of such a template class, or is this intentionally not provided?

[DINHQuangDung1999]

I have also recently started using Warp, and I found your idea very helpful. It was a struggle for me in the first 2-3 weeks getting used to the syntaxes as a freshman. Having a template class with some explaination would definitely help ones avoid spending so much time figuring out how to write a proper simulation class.

[JonathanKuelz]

I can share this version for now -- not sure if it works out of the box as I had to copy-past a couple of dependencies. Let me know if anything's missing:

from __future__ import annotations

import abc
import functools
from pathlib import Path
from typing import Collection, Dict, Optional, Tuple, Union, List

import numpy as np
import warp as wp
import warp.sim
import warp.sim.render
import torch


class WarpSim(HasWarpModel, abc.ABC):
    """A template class to set up Warp Simulations."""

    def __init__(self,
                 builder: wp.sim.ModelBuilder,
                 num_frames: int,
                 substeps: int,
                 fps: int,
                 name: str,
                 integrator: str = 'SemiImplicit',
                 requires_grad: bool = False,
                 stage_path: Optional[str] = None,
                 capture_cuda_graph: bool = False,
                 verbose: Union[bool, str, int] = False,
                 *,
                 loss_context: str = 'simulation',
                 requires_mesh_gradients: bool = False,
                 render_kwargs: Optional[dict] = None,
                 render_com: bool = False,
                 renderer: Optional[wp.sim.render.SimRenderer] = None,
                 store_gradients: bool = False,
                 store_gradients_for: Collection[str] = (),
                 store_intermediate_states: Optional[bool] = None,
                 verbosity_level_timing: int = 1,
                 joint_attach_kd: Optional[float] = None,
                 joint_attach_ke: Optional[float] = None,
                 ignore_kernels_in_tape: Optional[List[str]] = None,
                 ):
        """
        Sets up a generic warp simulation.

        :param builder: A builder for the model to be used for simulation.
        :param num_frames: The total number of frames to be simulated. If the _breaking_condition method is overridden,
            this will be the maximum number of frames for which states will be pre-allocated in case that gradients are
            required or store_intermediate_states is set to True.
        :param substeps: The number of substeps per frame. Together with fps, this determines the simulation time step.
        :param fps: How many frames per second should be simulated and possibly rendered.
        :param name: Name for this simulation.
        :param integrator: Choose between 'SemiImplicit', 'Featherstone', and 'XPBD'.
        :param requires_grad: Should the model and states be set up to compute gradients?
        :param stage_path: If provided, the simulation will render a USD stage to this path. If you start it with
            "usd:", this will automatically resolve to the default storage path of this package.
        :param capture_cuda_graph: Should the CUDA graph be captured, stored and reused?
        :param verbose: Either set the verbosity level directly or provide a string with 'v' for each level. True is 1.
        :param loss_context: The context in which the loss function should be computed. Can be 'substep', 'frame', or
            'simulation'.
        :param store_gradients: If true, the latest gradients from simulation will be stored in the latest_gradients
            property of this instance.
        :param store_gradients_for: Only relevant if store_gradients is true. A list of attributes or arrays for which
            gradients should be stored.
        :param requires_mesh_gradients: If true, enables computations of gradients w.r.t. mesh vertices.
        :param render_kwargs: Additional keyword arguments for the renderer.
        :param render_com: If true, the renderer will show the com's of all bodies.
        :param store_intermediate_states: When computing gradients, it is strictly necessary to store intermediate
            states. This will be set to requires_grad if not provided.
        :param verbosity_level_timing: Which verbosity level should be used to time the simulation?
        :param joint_attach_kd: If provided, this will override the default damping for the joint attach spring.
        :param joint_attach_ke: If provided, this will override the default stiffness for the joint attach spring.
        :param ignore_kernels_in_tape: If provided, all tapes will ignore kernels with the given keys (names).
        """
        self._verbosity_level_timing: int = verbosity_level_timing
        self.ignore_kernels_in_tape = ignore_kernels_in_tape

        if isinstance(verbose, str):
            verbose = verbose.lower().count('v')
        elif isinstance(verbose, bool):
            verbose = int(verbose)
        self.verbose: int = verbose
        if store_intermediate_states is None:
            store_intermediate_states = requires_grad
        self._store_intermediate_states: bool = store_intermediate_states
        if render_kwargs is None:
            render_kwargs = {}
        self.render_kwargs = render_kwargs
        self.render_com: bool = render_com
        self.name: str = name

        self.num_frames: int = num_frames
        self.substeps: int = substeps
        self.current_frame: int = 0
        self.current_substep: int = 0
        self.fps: float = fps
        self._requires_grad: bool = requires_grad
        self.store_gradients: bool = store_gradients
        self.store_gradients_for = list(store_gradients_for)
        self.latest_gradients: Dict[wp.array, torch.Tensor] = {}

        self.builder: wp.sim.ModelBuilder = builder
        self.model: wp.sim.Model = builder.finalize(requires_grad=requires_grad)

        if requires_mesh_gradients:
            parsed_hashes = dict()  # Warp does not finalize meshes if another one with same hash was parsed before.
            for shape_src in (s for s in self.shape_meshes):
                if hash(shape_src) in parsed_hashes:
                    shape_src.mesh = parsed_hashes[hash(shape_src)].mesh
                else:
                    parsed_hashes[hash(shape_src)] = shape_src
                    shape_src.mesh.points.requires_grad = True
        self._requires_mesh_gradients: bool = requires_mesh_gradients

        self.inertia_values = wp.empty(shape=(self.model.body_count, 6), dtype=float, requires_grad=requires_grad)
        self.body_inertia = wp.zeros(shape=(self.model.body_count,), dtype=wp.mat33f, requires_grad=True)
        wp.launch(extract_inertia_values, dim=self.model.body_count, inputs=[self.model.body_inertia],
                  outputs=[self.inertia_values])

        self.control: wp.sim.Control = self.model.control()
        self.states: List[List[wp.sim.State]] = [[]]
        self._reset_states()

        if joint_attach_kd is not None:
            self.model.joint_attach_kd = joint_attach_kd
        if joint_attach_ke is not None:
            self.model.joint_attach_ke = joint_attach_ke

        self.integrator: wp.sim.Integrator = self._get_integrator(integrator)

        if self.model.joint_count > 0:
            wp.sim.eval_fk(self.model, self.model.joint_q, self.model.joint_qd, None, self.current_state)

        if renderer is not None:
            self._do_render: bool = True
            self.stage_path: Optional[str] = renderer.stage.GetRootLayer().identifier
            self.renderer = renderer
        elif stage_path is not None:
            assert not capture_cuda_graph, "Currently, capturing the CUDA graph is not supported when rendering."
            if not (stage_path.endswith('.usd') or stage_path.endswith('.usda')):
                raise ValueError("Please provide a path to a USD stage file with the extension '.usda'.")
            stage_path = resolve_usd_path(stage_path, self.name)
            self._do_render: bool = True
            self.stage_path: Optional[str] = stage_path
            Path(stage_path).unlink(missing_ok=True) # It seems like usda sometimes just appends instead of overwriting.
            self.renderer = wp.sim.render.SimRenderer(self.model, stage_path, **self.render_kwargs)
        else:
            self._do_render: bool = False
            self.stage_path: Optional[str] = None
            self.renderer = None

        if self._requires_grad:
            self.loss = wp.array((0.,), dtype=float, requires_grad=True)
        else:
            self.loss = None

        self.loss_context = loss_context.lower()
        self._graph = None
        if capture_cuda_graph:
            with wp.ScopedCapture() as capture:
                self.launch()
            self._graph = capture.graph

    @abc.abstractmethod
    def simulation_step(self):
        """The simulation step for a single simulation step (substep) needs to be implemented by inheriting classes."""
        pass

    @property
    def current_state(self) -> wp.sim.State:
        """Returns the currently active state."""
        if not self._store_intermediate_states:
            return self.states[0][0]
        return self.states[self.current_frame][self.current_substep]

    @property
    def current_cumulative_substep(self):
        """The current substep, not reset at every frame."""
        return self.current_frame * self.substeps + self.current_substep

    @property
    def frame_dt(self) -> float:
        """The time that passes between two frames."""
        return 1 / self.fps

    @property
    def loss_context(self) -> str:
        """The current context for loss computation."""
        return self._loss_context

    @loss_context.setter
    def loss_context(self, value: str):
        """Sets the context for loss computation."""
        assert value in ('substep', 'frame', 'simulation', 'deactivated'), "Unknown loss interval."
        self._loss_context = value

    @property
    def next_state(self) -> wp.sim.State:
        """Returns the state that the integrator will be writing to next."""
        if not self._store_intermediate_states:
            return self.states[0][1]
        if self.current_substep < self.substeps - 1 or self.current_frame == self.num_frames - 1:
            return self.states[self.current_frame][self.current_substep + 1]
        else:
            return self.states[self.current_frame + 1][0]

    @property
    def sim_dt(self) -> float:
        """The time that passes within one integratino step."""
        return self.frame_dt / self.substeps

    @property
    def sim_time(self) -> float:
        """The current simulation time."""
        return self.current_frame * self.frame_dt + self.current_substep * self.sim_dt

    @property
    def shape_meshes(self) -> List[wp.sim.Mesh]:
        """Returns all shape meshes in the model."""
        return [s for s in self.model.shape_geo_src if isinstance(s, wp.sim.Mesh)]

    @property
    def shape_pose(self) -> wp.array(dtype=wp.transform):
        """The current transform of all shape coordinate origins."""
        tf = wp.empty(shape=(self.model.shape_count,), dtype=wp.transform)
        wp.launch(get_shape_poses, dim=self.model.shape_count,
                  inputs=[self.model.shape_body, self.current_state.body_q, self.model.shape_transform],
                  outputs=[tf]
                  )
        return tf

    @property
    def total_number_substeps(self) -> int:
        """The total number of substeps that will be simulated."""
        return self.num_frames * self.substeps

    @property
    def unique_shape_meshes(self) -> List[Tuple[int, wp.sim.Mesh]]:
        """Returns all unique shape meshes in the model with their corresponding shape index."""
        meshes = self.shape_meshes
        hashes = [hash(m) for m in meshes]
        return [(i, m) for i, (h, m) in enumerate(zip(hashes, meshes)) if h not in hashes[:i]]

    @property
    def verbose(self):
        """Gets the verbosity level for the simulation."""
        return self._verbose

    @verbose.setter
    def verbose(self, value):
        """Sets the verbosity level for the simulation."""
        self._verbose = value
        self._timer_active = value >= self._verbosity_level_timing

    # ---------- Meta ----------

    def clone(self,
              *args,
              **kwargs) -> WarpSim:
        """
        Create a new simulation that is a clone of this one with the provided changes.
        """
        new_kwargs = dict(
            builder=self.builder,
            num_frames=self.num_frames,
            substeps=self.substeps,
            fps=self.fps,
            name=self.name,
            integrator=self.integrator.__class__.__name__.replace('Integrator', ''),
            requires_grad=self._requires_grad,
            stage_path=self.stage_path,
            capture_cuda_graph=self._graph is not None,
            verbose=self.verbose,
            loss_context=self.loss_context,
            render_kwargs=self.render_kwargs,
            renderer=self.renderer,
            store_intermediate_states=self._store_intermediate_states,
            verbosity_level_timing=self._verbosity_level_timing,
            joint_attach_kd=self.model.joint_attach_kd,
            joint_attach_ke=self.model.joint_attach_ke,
        )
        new_kwargs.update(kwargs)
        return WarpSim(**new_kwargs)

    # ---------- Simulation ----------

    def apply_force(self, body_name: str, force: np.ndarray, state: wp.sim.State):
        """Applies the force to the body at the center of mass and writes it to state."""
        body_idx = self.body_index(body_name)
        wp.launch(add_wrench_body, dim=1, inputs=[wp.spatial_vector(*force), body_idx], outputs=[state.body_f])

    def get_body_com(self) -> np.array:
        """Get an array with all current body centers of mass."""
        body_com = self.model.body_com
        body_q = self.current_state.body_q
        com = wp.empty(shape=body_com.shape, dtype=wp.vec3)
        wp.launch(transform_points, dim=self.model.body_count, inputs=[body_q, body_com], outputs=[com])
        return com.numpy()

    def launch(self):
        """Top level entry point to launch the simulation."""
        if self._do_render:  # I like to render the initial state as well.
            self.render()

        with self._get_timer(name='simulation'):
            sim_tape = self._get_tape(level='simulation')
            with sim_tape:
                self._fix_inertial_properties(sim_tape)
                f = 0
                while not self._break_condition():
                    self.current_frame = f
                    with self._get_timer(name='frame'):
                        self.step_frame()
                    f += 1
                self.possibly_compute_loss('simulation')
            self.possibly_optimize('simulation', sim_tape)

        if self._do_render:
            self.renderer.save()

    def render(self):
        """Render the current frame."""
        assert self._do_render, "This simulation is currently not set up to render."
        with self._get_timer(name='render'):
            self.renderer.begin_frame(self.sim_time)
            self.renderer.render(self.current_state)
            self._render_callback()
            if self.render_com:
                self._render_com()
            self.renderer.end_frame()

    def step_frame(self):
        """Perform a single simulation step (substep) and render the frame."""
        frame_tape = self._get_tape(level='frame')
        with frame_tape:
            self._fix_inertial_properties(frame_tape)
            for ss in range(self.substeps):
                self.current_substep = ss
                step_tape = self._get_tape(level='substep')
                with step_tape:
                    self._fix_inertial_properties(step_tape)
                    self.simulation_step()
                    self.possibly_compute_loss('substep')
                self.possibly_optimize('substep', step_tape)
                if not self._store_intermediate_states:
                    s0, s1 = self.current_state, self.next_state
                    self.states = [[s1, s0]]
            self.possibly_compute_loss('frame')

        self.possibly_optimize('frame', frame_tape)
        if self._do_render:
            self.render()

    def loss_fn(self):
        """The loss function to be minimised."""
        if self._requires_grad:
            raise NotImplementedError("Loss function needs a custom implementation.")
        else:
            raise ValueError("You should not end up in the loss function if not requiring gradients.")

    # ---------- Optimisation ----------
    
    def opt_fn(self):
        """The optimisation function to be called after the loss has been computed."""
        if self._requires_grad:
            raise NotImplementedError("Optimization routine needs a custom implementation.")
        else:
            raise ValueError("You should not end up in the optimization routine if not requiring gradients.")
    
    def possibly_compute_loss(self, level: str):
        """
        Computes the loss if the current context is the loss context.
        """
        if not self._requires_grad:
            return
        elif level.lower() != self.loss_context:
            return
        with self._get_timer(name='loss'):
            self.loss_fn()
        if self.verbose >= 1:
            print(f"Loss: {self.loss.numpy()[0]}")

    def possibly_optimize(self, level: str, tape: Union[wp.Tape, DummyTape]):
        """
        Backprops the gradient and performs optimization if the current context is the loss context.
        """
        if not self._requires_grad:
            return
        if level.lower() != self.loss_context:
            return
        with self._get_timer(name='backward'):
            if self.ignore_kernels_in_tape is not None:
                if self.verbose >= 1:
                    print('Warning: Experimental feature: ignoring kernels in tape.')
                tape.launches = [l for l in tape.launches if not (isinstance(l, list) and l[0].key in self.ignore_kernels_in_tape)]
            tape.backward(self.loss)
        if self.store_gradients:
            self._store_gradients(tape.gradients)
        self.opt_fn()
        tape.zero()
        self.loss.zero_()

    def _break_condition(self):
        """Called after every frame to see if the simulation should stop."""
        return self.current_frame >= self.num_frames - 1

    # ---------- Analysis ----------

    def get_state_attr_array(self, attr: str, idx: int) -> wp.array:
        """
        Creates a new, flat warp array that contains all the state attributes at idx in the history of this simulation.

        Useful for kernel-based loss computation after a simulation has terminated.
        :param attr: The name of the state attribute.
        :param idx: The index of the attribute to be copied.
        """
        if not self._store_intermediate_states:
            raise ValueError("This method doesn't make sense if intermediate states were not stored.")
        arr = wp.empty((self.total_number_substeps,), getattr(self.current_state, attr).dtype, requires_grad=True)
        for i in range(self.total_number_substeps):
            frame = i // self.substeps
            substep = i % self.substeps
            state = self.states[frame][substep]
            wp.launch(array_select_copy, dim=1, inputs=[getattr(state, attr), i, idx], outputs=[arr])
        return arr

    # ---------- Private Methods -----------

    def _fix_inertial_properties(self, tape: Union[wp.Tape, DummyTape]):
        """Recomputes the inverse inertial properties of the model."""
        if isinstance(tape, DummyTape):
            return
        with self._get_timer(name='inertial_properties'):
            if self._requires_mesh_gradients:
                recompute_mesh_inertial_contributions(self.model)
            wp.launch(write_symmetric_inertia, dim=self.model.body_count, inputs=[self.inertia_values],
                      outputs=[self.body_inertia])
            self.model.body_inertia = self.body_inertia
            wp.launch(recompute_inverse_inertial_properties, dim=self.model.body_count,
                      inputs=[self.model.body_mass, self.body_inertia],
                      outputs=[self.model.body_inv_mass, self.model.body_inv_inertia]
                      )
            if isinstance(self.integrator, wp.sim.FeatherstoneIntegrator):
                # The featherstone integrator stores some auxiliarry varia
                self.integrator.allocate_model_aux_vars(self.model)

    def _get_integrator(self, name: str) -> wp.sim.Integrator:
        """Returns the integrator class that will be used, based on the namestring."""
        if name.lower() == 'semiimplicit':
            integrator: wp.sim.Integrator = wp.sim.SemiImplicitIntegrator()
        elif name.lower() == 'featherstone':
            integrator: wp.sim.Integrator = wp.sim.FeatherstoneIntegrator(self.model)
        elif name.lower() == 'xpbd':
            integrator: wp.sim.Integrator = wp.sim.XPBDIntegrator()
        else:
            raise ValueError(f"Unknown integrator: {name}")
        return integrator

    def _get_tape(self, level: str) -> Union[wp.Tape, DummyTape]:
        """
        Returns a context manager that is either a warp tape or does exactly nothing, depending on whether the current
        scope is the same as the loss computation level.
        """
        if level.lower() == self.loss_context:
            return wp.Tape()
        return DummyTape()

    def _get_timer(self, name: str, **kwargs) -> wp.ScopedTimer:
        """Return a timer that's active if the verbosity of this class is set accordingly."""
        kwargs.setdefault('active', self._timer_active)
        return wp.ScopedTimer(name, **kwargs)

    def _render_callback(self):
        """An optional callback to be implemented to render additional information each frame."""
        pass

    def _render_com(self):
        """Renders the com of all bodies."""
        com = self.get_body_com()
        for i, p in enumerate(com):
            name = self.model.body_name[i].replace(' ', '_')
            self.renderer.render_sphere(name=name, pos=p, radius=0.05, rot=wp.quat_identity(), color=(1., 0., 0.))

    def _reset_states(self):
        """Sets all the states of the simulation to the initial state."""
        if not self._store_intermediate_states:
            self.states: List[List[wp.sim.State]] = [[self.model.state() for _ in range(2)]]
        else:
            self.states: List[List[wp.sim.State]] = [[self.model.state() for i in range(self.substeps)]
                                                     for _ in range(self.num_frames)]
            self.states[self.num_frames - 1].append(self.model.state())  # A final state to integrate into at the end.
        for fr in self.states:
            for s in fr:
                s.clear_forces()

    def _store_gradients(self, g: Dict):
        """Figures out which gradients to store and stores them."""
        arrays = list()
        for attr in self.store_gradients_for:
            if isinstance(attr, wp.array):
                arrays.append(attr)
            elif not isinstance(attr, str):
                raise ValueError("Attribute must be a string or array.")
            elif attr.startswith('model.'):
                arrays.append(functools.reduce(getattr, attr.split('.'), self))
            elif attr == 'mesh_vertices':
                vertices = [m.mesh.points for m in self.model.shape_geo_src if isinstance(m, wp.sim.Mesh)]
                arrays.extend(vertices)
            else:
                raise ValueError("Unknown attribute to store gradients for:", attr)
        for a in arrays:
            self.latest_gradients[a] = torch.tensor(a.grad).detach().clone()

@wp.kernel
def get_shape_poses(
    body_indices: wp.array(dtype=int),
    body_q: wp.array(dtype=wp.transform),
    shape_tf: wp.array(dtype=wp.transform),
    result: wp.array(dtype=wp.transform),
):
    tid = wp.tid()
    body_tf = body_q[body_indices[tid]]
    result[tid] = wp.transform_multiply(body_tf, shape_tf[tid])

@wp.kernel
def extract_inertia_values(
    inertia_matrix: wp.array(dtype=wp.mat33f),
    unique_values: wp.array(ndim=2, dtype=float),
):
    """Write the inertia matrix to the result array, explicitly setting Iyx=Ixy, Izx=Ixz, Izy=Iyz."""
    tid = wp.tid()
    inertia = inertia_matrix[tid]
    unique_values[tid, 0] = inertia[0, 0]
    unique_values[tid, 1] = inertia[1, 1]
    unique_values[tid, 2] = inertia[2, 2]
    unique_values[tid, 3] = inertia[0, 1]
    unique_values[tid, 4] = inertia[0, 2]
    unique_values[tid, 5] = inertia[1, 2]

@wp.kernel
def write_symmetric_inertia(
    unique_values: wp.array(ndim=2, dtype=float),
    result: wp.array(dtype=wp.mat33f),
):
    tid = wp.tid()
    result[tid] = wp.mat33f(
        unique_values[tid, 0], unique_values[tid, 3], unique_values[tid, 4],
        unique_values[tid, 3], unique_values[tid, 1], unique_values[tid, 5],
        unique_values[tid, 4], unique_values[tid, 5], unique_values[tid, 2],
    )


@wp.kernel
def add_wrench_body(
        wrench: wp.spatial_vectorf,
        body_id: int,
        f_body: wp.array(dtype=wp.spatial_vectorf)
):
    """Apply a force to a body."""
    f_body[body_id] += wrench


@wp.kernel
def array_select_copy(a: wp.array(dtype=Any), idx_copy: int, idx_select: int, out: wp.array(dtype=Any)):
    """Copy element at idx_select from array a to idx_copy in array out."""
    out[idx_copy] = a[idx_select]


@wp.kernel
def recompute_inverse_inertial_properties(
    m: wp.array(dtype=float),
    I: wp.array(dtype=wp.mat33),
    inv_m: wp.array(dtype=float),
    inv_I: wp.array(dtype=wp.mat33),
):
    """Computes inv_m and inv_I from m and I."""
    tid = wp.tid()
    inv_m[tid] = 1. / m[tid]
    inv_I[tid] = wp.inverse(I[tid])


@wp.kernel
def transform_points(
    transforms: wp.array(dtype=wp.transform),
    points: wp.array(dtype=wp.vec3),
    result: wp.array(dtype=wp.vec3),
):
    tid = wp.tid()
    result[tid] = wp.transform_point(transforms[tid], points[tid])

def resolve_usd_path(usd_path: str, sim_name: str = 'simulation') -> str:
    rel_path = Path(sim_name)
    rel_dir = str(rel_path.parent) + '/'
    base_name = rel_dir + re.sub(r'_\d+$', '', rel_path.name)
    sim_name = rel_path.name

    pth = re.sub(r'^usd:', str(usd_dir) + '/', usd_path)
    pth = pth.format(name=sim_name, base_name=base_name)
    pth = Path(pth)
    return str(pth.expanduser().resolve())

class HasWarpModel(abc.ABC):
    """A utility class that provides shortcuts assuming it has a model property."""

    model: wp.sim.Model

    def body_index(self, name: str) -> int:
        """Shortcut to get the index of a body by name."""
        return self.model.body_name.index(name)

    def body_q(self, name: str) -> np.ndarray:
        """Shortcut to get the 7D-coordinates (body_q) of a body by name."""
        return self.model.body_q.numpy()[self.body_index(name)]

    def joint_control_index(self, name: str) -> int:
        """Shortcut to get the (start) index of a joint in control.joint_act by name."""
        name_idx = self.model.joint_name.index(name)
        return self.model.joint_axis_start.numpy()[name_idx]


class DummyTape:
    """
    This defines a custom context manager that does nothing.

    Assume there's a function you possibly want to call with a tape in warp:
    tape = wp.Tape()
    with tape:
        fun()

    Now if at times, this should NOT be called with a warp tape, it is cumbersome to duplicate all the code
    in an if/else context. Instead, you can do:
    tape = wp.Tape() if gradients else DummyTape()
    with tape:
        fun()
    """

    def backward(self, *args, **kwargs):
        """If someone tries to use backward on this tape, something went wrong."""
        raise RuntimeError('This is a dummy tape, you should not call backward on it.')

    def __enter__(self):
        """Do nothing."""
        pass

    def __exit__(self, exc_type, exc_val, exc_tb):
        """Do nothing."""
        pass


def recompute_mesh_inertial_contributions(model: wp.sim.Model):
    """
    This function recomputes the inertial properties of all shapes in the model that are Meshes.

    This function is useful because it allows recomputing inertial properties in a differentiable manner inside a tape.
    The original builder.finalize() method provided by warp does not allow differentiation.
    :param model: The model to update
    """
    geo_types = model.shape_geo.type.numpy()
    shape_body = model.shape_body.numpy()
    for i in range(model.shape_count):
        if geo_types[i] != wp.sim.GEO_MESH:
            continue

        com = wp.zeros(1, dtype=wp.vec3, requires_grad=True)
        inertia = wp.zeros(1, dtype=wp.mat33, requires_grad=True)
        mass = wp.zeros(1, dtype=float, requires_grad=True)
        # volume = wp.zeros(1, dtype=float, requires_grad=True)

        shape = model.shape_geo_src[i]
        wp_mesh = shape.mesh
        vertices = wp_mesh.points
        faces = wp_mesh.indices

        wp.launch(array_mean_kernel, dim=vertices.shape[0], inputs=[vertices], outputs=[com])
        wp.launch(
            kernel=compute_mesh_inertia,
            dim=faces.shape[0] // 3,
            inputs=[com, 0.25, faces, vertices],
            outputs=[mass, inertia],
        )
        wp.launch(correct_inertia, dim=1,
                  inputs=[shape_body[i], i, mass, com, inertia, model.shape_transform, shape.mass, shape.com, shape.I, 1.0],
                  outputs=[model.body_mass, model.body_inv_mass, model.body_com,
                           model.body_inertia, model.body_inv_inertia],
                  )
        shape.mass = float(mass.numpy()[0])
        shape.com = wp.vec3(com.numpy())
        shape.I = wp.mat33(inertia.numpy())

[JonathanKuelz]

For future reference: The rewarped environment could also be a great starting point. Looks amazing at first sight

Code: https://github.com/rewarped/rewarped/blob/main/rewarped/environment.py
Paper: https://arxiv.org/abs/2412.12089

[DINHQuangDung1999]

Agree. I have been working with it for over 2 weeks and I think this is quite well-built.