About the stability of undamped systems

[simon-bachhuber]

Dear Brax authors,

thanks for open-sourcing this code, it's quite illuminating.
For mostly self-educational purposes i have built my own simulator / engine, but whilst i do use different conventions the underlying concepts and math is similar or identical.

Currently, i have hit the following stability problem where i forward simulate a double pendulum with no damping (i.e. perfect hinge joints). There exists gravity and so initially there is a certain amount of potential energy in the system. However, when the simulation time step is too long, the simulation becomes unstable and energy clearly increases. Decreasing the simulation time step makes the simulation stable again.

Animation 1: dt = 0.01. At a certain breakpoint the simulation becomes unstable.

Animation 2: dt = 0.001. The simulation stays stable for the entire time duration. Note that the trajectory in both animations is identical only during the first few seconds.

---

First: What surprised me immediately was that, whilst the trajectories may differ, i thought that using an semi-implicit euler scheme would guarantee that energy does not increase over time. Clearly i was wrong..

Second: I checked out your source and how you forward integrate. I tried both versions of exactly and approximately inverting the mass matrix, and both yielded very similar stability properties. However, i also can not "implicitly integrate the DOF damping" as my DOFs do not have any damping to them.

My question would then be as follows:

• Am i doing something wrong? (if so, could you provide me with some pointers)
• Or, am i simply asking for too much?

Thanks for your time, and for reading this.

PS: I created the animations using vispy. I would probably be able to write my vispy_render.py module interface general enough such that it can also be used easily by brax. However, i would only do this if there is interest. :)

[btaba]

Hi @SimiPixel thanks for taking a look at brax! Semi-implicit Euler does not guarantee stability, it's just "more" stable. The dt should be sufficiently small s.t. the integration does not become unstable. mujoco has several other integrators implemented that you can take a look at if you're interested https://mujoco.readthedocs.io/en/stable/computation.html#numerical-integration

A vispy renderer would be awesome, the viz looks pretty cool!

[simon-bachhuber]

yes that's what i figured but good to confirm it

[erikfrey]

+1 we would be very happy to review a vispy renderer PR!

[simon-bachhuber]

All right then i will close this and create a new discussion to talk about a vispy renderer.

One final thought: You currently do not define some abstract interface what a renderer backend would have to implement?
If not, that's probably what i will do first.

[btaba]

Thanks @SimiPixel !! We currently have this renderer https://github.com/google/brax/blob/main/brax/io/image.py

I think the interface would be pretty similar to this: render(sys: base.System, state: List[base.State], width: int, height: int, ...)

[simon-bachhuber]

Hmm, i was thinking something like this:

Edit: Please think of this as Pseudocode for now ;)

from abc import ABC, abstractmethod
from typing import Optional, TypeVar

import jax
import jax.numpy as jnp

Camera = TypeVar("Camera")
Geometry = TypeVar("Geometry")
Visual = TypeVar("Visual")


@jax.jit
@partial(jax.vmap, in_axes=(None, 0, 0))
def get_transform(
    x_links: "Transform", x_link_to_geom: "Transform", geom_link_idx
) -> jax.Array:
    "This can easily account for possible convention differences"
    return x_links[geom_link_idx].do(x_link_to_geom).to_4x4()


class Renderer(ABC):
    """
    Example:
        >> renderer = Renderer()
        >> renderer.populate(sys.geoms)
        >> for x in xs:
        >>   renderer.update_scene(x)
        >>   image = renderer.render()
    """

    @abstractmethod
    def _get_camera(self) -> Camera:
        pass

    @abstractmethod
    def _set_camera(self, camera: Camera) -> None:
        pass

    @abstractmethod
    def _render(self) -> jax.Array:
        pass

    def render(
        self, camera: Optional[Camera | list[Camera]] = None
    ) -> jax.Array | list[jax.Array]:
        if camera is None:
            camera = self._get_camera()

        if isinstance(camera, Camera):
            self._set_camera(camera)
            return self._render()

        images = []
        for cam in camera:
            self._set_camera(cam)
            images.append(self._render())
        return images

    @staticmethod
    def _add_box(geom: "Box") -> Visual:
        raise NotImplementedError

    @staticmethod
    def _add_sphere(geom: "Sphere") -> Visual:
        raise NotImplementedError

    def populate(self, geoms: list[Geometry]):
        self.geom_link_idx = []
        self.geom_transform = []
        self.visuals = []
        for geom in geoms:
            self.geom_link_idx.append(geom.link_idx)
            self.geom_transform.append(geom.transform)
            if isinstance(geom, "Box"):
                visual = self._add_box(geom)
            elif isinstance(geom, "Sphere"):
                visual = self._add_sphere(geom)
            self.visuals.append(visual)

        # stack them; Pseudocode here; transforms are ofc PyTrees
        self.geom_link_idx = jnp.concatenate(self.geom_link_idx)
        self.geom_transform = jnp.concatenate(self.geom_transform)

    @abstractmethod
    def _mutate_visual(self, visual: Visual) -> None:
        pass

    def update_scene(self, x: "Transform") -> None:
        "`x` are (n_links,) Transforms."
        # step 1: pre-compute all required transforms (uses jax.vmap)
        transform_per_visual = get_transform(x, self.geom_transform, self.geom_link_idx)

        # step 2: update visuals
        for i, visual in enumerate(self.visuals):
            self._mutate_visual(visual, transform_per_visual[i])

[btaba]

Hi @SimiPixel thanks for the extensive example! I think at this point we're not looking to be too prescriptive w.r.t. a render API. The simplest MVP to plug into brax right now is to implement brax.io.image_vispy.render(sys, state: List[State]). Let us know if there's interest there!

[simon-bachhuber]

Hey, yeah sure that should be easy enough.

I have made a quick somewhat compatible version, however, conventions of your transforms are quite a bit different to what i have been using. Would be good if you could look over the method here

The test_vispy_render.py in the root of the project can be executed, ideally this ofc should give us a nice inverted pendulum..

[btaba]

Took a quick look at the line you pointed to, looks good to me