Mujoco slide joint type

[cemlyn007]

Feature: Mujoco slide joint type

I am currently in the process of recreating a Mujoco environment in Brax but I am having a bit of an issue because their code base uses an XML asset which uses the slide joint type.

Does Brax intend on supporting this joint and if so I will collaborate on a PR?

Looking forward to hearing from you all 👋

[btaba]

Hi @cemlyn007

I'm not sure if we'll support this in the current version of brax, but it will certainly be supported in the next major version. You'll be able to load the mujoco xml directly into brax. For now, a slide joint can be implemented roughly as:

@pytree.register
class Prismatic(Joint):
  """A prismatic joint."""

  def apply_reduced(self, qp_p: QP, qp_c: QP) -> Tuple[Q, Q]:
    """Constructs the position constraint update for a parent and child."""
    # positional constraints
    pos_p, pos_c, _ = self.get_offset(qp_p, qp_c)
    dq_p, dq_c = self.apply_position_update(
        qp_p, qp_c, pos_p, pos_c)

    # torque from qp_p.rot to qp_c.rot so that rotations are not allowed
    axis = math.rotate(self.axis_p[0], qp_p.rot)
    axis_c = math.rotate(self.axis_c[0], qp_c.rot)
    dq_1 = jp.cross(axis, axis_c)

    ref_p = math.rotate(self.axis_p[2], qp_p.rot)
    ref_c = math.rotate(self.axis_c[2], qp_c.rot)
    rot_axis = dq_1
    psi = math.signed_angle(rot_axis, axis_c, axis)
    fixrot = math.quat_rot_axis(rot_axis, psi)
    ref_c_n = math.rotate(ref_c, fixrot)
    dq_2 = jp.cross(ref_p, ref_c_n)

    v_apply = jp.vmap(self.apply_angle_update, [False, False, True])
    dq_p_ang, dq_c_ang = v_apply(qp_p, qp_c, jp.array([dq_1, dq_2]))
    dq_p += Q(
        pos=dq_p_ang.pos[0] + dq_p_ang.pos[1],
        rot=dq_p_ang.rot[0] + dq_p_ang.rot[1])
    dq_c += Q(
        pos=dq_c_ang.pos[0] + dq_c_ang.pos[1],
        rot=dq_c_ang.rot[0] + dq_c_ang.rot[1])

    return dq_p, dq_c  # pytype: disable=bad-return-type

  def get_offset(self, qp_p, qp_c):
    """Gets offset btwn parent and child along prismatic joint axis."""
    pos_p, _ = qp_p.to_world(self.off_p)
    pos_c, _ = qp_c.to_world(self.off_c)
    axis_p = math.rotate(self.axis_p[0], qp_p.rot)

    pos_c_in_axis = pos_p + axis_p * axis_p.dot(pos_c - pos_p)

    pos_c_in_axis = jp.where(axis_p.dot(pos_c - pos_p) > self.limit[0][1],
                             pos_p + self.limit[0][1] * axis_p,
                             pos_c_in_axis)
    pos_c_in_axis = jp.where(axis_p.dot(pos_c - pos_p) < self.limit[0][0],
                             pos_p + self.limit[0][0] * axis_p,
                             pos_c_in_axis)
    return pos_c_in_axis, pos_c, axis_p.dot(pos_c_in_axis)

  def axis_angle(self, qp_p: QP, qp_c: QP) -> Tuple[Any, Any]:
    """Returns axes and angles of a single joint."""
    axis = math.rotate(self.axis_p[0], qp_p.rot)
    _, _, dist = self.get_offset(qp_p, qp_c)
    return (axis,), (dist,)

  def damp(self, qp: QP) -> P:
    def damp_reduced(self, qp_p, qp_c):
      axis_p = math.rotate(self.axis_p[0], qp_p.rot)
      force = -1. * self.angular_damping * axis_p * (qp_p.vel - qp_c.vel)
      dp = 1. / self.body_p.mass * force
      dc = -1. / self.body_c.mass * force
      return dp, dc

    qp_p = jp.take(qp, self.body_p.idx)
    qp_c = jp.take(qp, self.body_c.idx)
    dp, dc = jp.vmap(damp_reduced)(self, qp_p, qp_c)  # pytype: disable=attribute-error

    # sum together all impulse contributions across parents and children
    body_idx = jp.concatenate((self.body_p.idx, self.body_c.idx))
    dq = jp.concatenate((dp, dc))  # pytype: disable=attribute-error
    dq = jp.segment_sum(dq, body_idx, qp.ang.shape[0])

    return P(vel=dq, ang=jp.zeros_like(dq))

Hope this helps for now

[btaba]

The actuator would look something like:

@pytree.register
class Position(Actuator):
  """Applies force to satisfy a target position."""

  def apply_reduced(self, act: jp.ndarray, qp_p: QP, qp_c: QP) -> Tuple[P, P]:
    axis, pos = self.joint.axis_angle(qp_p, qp_c)
    axis, pos = jp.array(axis), jp.array(pos)

    limit_min, limit_max = self.joint.limit[:, 0], self.joint.limit[:, 1]
    act = jp.clip(act, limit_min, limit_max)
    force = (act - pos) * self.strength
    force = jp.sum(jp.vmap(jp.multiply)(axis, force), axis=0)

    dp = -1. / self.joint.body_p.mass * force
    dc = 1. / self.joint.body_c.mass * force
    return dp, dc

  def apply(self, qp: QP, act: jp.ndarray) -> P:
    qp_p = jp.take(qp, self.joint.body_p.idx)
    qp_c = jp.take(qp, self.joint.body_c.idx)
    act = jp.take(act, self.act_index) * self.act_mask
    dp, dc = jp.vmap(type(self).apply_reduced)(self, act, qp_p, qp_c)

    # sum together all impulse contributions across parents and children
    body_idx = jp.concatenate((self.joint.body_p.idx, self.joint.body_c.idx))
    dp = jp.concatenate((dp, dc))
    dp = jp.segment_sum(dp, body_idx, qp.pos.shape[0])
    return P(vel=dp, ang=jp.zeros_like(qp.ang))

[cemlyn007]

Hi @btaba, hope you are having a good weekend!

Thank you very much for showing code snippets! Am quite excited to see what is in your next major release 🔥

[btaba]

Thanks @cemlyn007 , converted to discussion so others can benefit as well