Note

This example is available as a jupyter notebook here.

Defining a custom Joint Type that supports dynamical simulation¤

In this notebook we will define a new joint type that is a hinge joint with a random joint axes direction.

It will also support dynamical simulation.

import ring
from ring import maths, base

import jax
import jax.numpy as jnp

import mediapy as media

from ring.algorithms.jcalc import _draw_rxyz

We will give this new joint type the identifier rr (random revolute). Although it actually already exists in the library, but we can overwrite it.

# we use such a `params` input to specify the joint-axes, if we later then randomize the attribute of the system object
# we will have the effect of a hinge joint with a randomized joint axes direction
# here we tell the library how it should initialize this `params` PyTree
def _draw_random_joint_axis(key):
    return maths.rotate(jnp.array([1.0, 0, 0]), maths.quat_random(key))

def _rr_init_joint_params(key):
    return dict(joint_axes=_draw_random_joint_axis(key))

# next, we tell the library how it can randomly draw a trajectory for its generalized coordinate; the hinge joint angle
def _rr_transform(q, params):
    # here we use this `params` object
    axis = params["joint_axes"]
    q = jnp.squeeze(q)
    rot = maths.quat_rot_axis(axis, q)
    return ring.Transform.create(rot=rot)

# this tells the library how to dynamically simulate the type of joint
def _motion_fn(params):
    return base.Motion.create(ang=params["joint_axes"])

# now, we can put it all together into a new `x_xy.JointModel`
rr_joint = ring.JointModel(_rr_transform, motion=[_motion_fn], rcmg_draw_fn=_draw_rxyz, init_joint_params=_rr_init_joint_params)

# and then we register the joint; Note that `overwrite`=True, because it already exists; that way you can e.g. overwrite the
# default joint types such as the free joint
ring.register_new_joint_type("rr", rr_joint, q_width=1, qd_width=1, overwrite=True)

xml_str = """
<x_xy>
<options dt="0.01" gravity="0 0 9.81"></options>
<worldbody>
<geom dim="0.1" type="xyz"></geom>
<body damping=".01" joint="rr" name="pendulum" pos="0 0 0.5">
<geom dim="0.1" type="xyz"></geom>
<geom dim="0.5 0.1 0.1" mass="0.5" pos="0.25 0 0" type="box"></geom>
</body>
</worldbody>
</x_xy>
"""

# this seed determines (among other things) the randomness of the joint-axes direction
# via the above specified `_rr_init_joint_params`
seed: int = 2
sys = ring.System.create(xml_str, seed=seed)

state = ring.State.create(sys)
xs = []
for t in range(500):
    state = jax.jit(ring.step)(sys, state)
    xs.append(state.x)

sys.links.joint_params

{'rr': {'joint_axes': Array([[ 0.41278404, -0.6329913 ,  0.65492845]], dtype=float32)},
 'default': Array([], shape=(1, 0), dtype=float32)}

def show_video(sys, xs: list[ring.Transform]):
    frames = sys.render(xs, render_every_nth=4)
    media.show_video(frames, fps=25)

show_video(sys, xs)

Rendering frames..: 100%|██████████| 125/125 [00:01<00:00, 94.24it/s]

the class x_xy.RCMG already has the built-in flag randomize_joint_params which can be toggled in order to use the user-provided logic _rr_init_joint_params for randomizing the joint parameters

(X, y), (key, q, x, _) = ring.RCMG(sys, randomize_joint_params=True, keep_output_extras=True).to_list()[0]

executing generators: 100%|██████████| 1/1 [00:02<00:00,  2.70s/it]


show_video(sys, x)

Rendering frames..: 100%|██████████| 1500/1500 [00:14<00:00, 103.09it/s]

but for dynamic_simulation flag to work we additional need to specify the function ring.JointModel.p_control_term

print(rr_joint.p_control_term)

None


try:
    (X, y), (key, q, x, _) = ring.RCMG(sys, randomize_joint_params=True, keep_output_extras=True, dynamic_simulation=True).to_list()[0]
except NotImplementedError:
    print("NotImplementedError: Please specify `JointModel.p_control_term` for joint type `rr`")

executing generators:   0%|          | 0/1 [00:00<?, ?it/s]

NotImplementedError: Please specify `JointModel.p_control_term` for joint type `rr`