def randomize_limited_and_rotational_joints()

in local_dm_control_suite/utils/randomizers.py [0:0]


def randomize_limited_and_rotational_joints(physics, random=None):
  """Randomizes the positions of joints defined in the physics body.

  The following randomization rules apply:
    - Bounded joints (hinges or sliders) are sampled uniformly in the bounds.
    - Unbounded hinges are samples uniformly in [-pi, pi]
    - Quaternions for unlimited free joints and ball joints are sampled
      uniformly on the unit 3-sphere.
    - Quaternions for limited ball joints are sampled uniformly on a sector
      of the unit 3-sphere.
    - The linear degrees of freedom of free joints are not randomized.

  Args:
    physics: Instance of 'Physics' class that holds a loaded model.
    random: Optional instance of 'np.random.RandomState'. Defaults to the global
      NumPy random state.
  """
  random = random or np.random

  hinge = mjbindings.enums.mjtJoint.mjJNT_HINGE
  slide = mjbindings.enums.mjtJoint.mjJNT_SLIDE
  ball = mjbindings.enums.mjtJoint.mjJNT_BALL
  free = mjbindings.enums.mjtJoint.mjJNT_FREE

  qpos = physics.named.data.qpos

  for joint_id in range(physics.model.njnt):
    joint_name = physics.model.id2name(joint_id, 'joint')
    joint_type = physics.model.jnt_type[joint_id]
    is_limited = physics.model.jnt_limited[joint_id]
    range_min, range_max = physics.model.jnt_range[joint_id]

    if is_limited:
      if joint_type == hinge or joint_type == slide:
        qpos[joint_name] = random.uniform(range_min, range_max)

      elif joint_type == ball:
        qpos[joint_name] = random_limited_quaternion(random, range_max)

    else:
      if joint_type == hinge:
        qpos[joint_name] = random.uniform(-np.pi, np.pi)

      elif joint_type == ball:
        quat = random.randn(4)
        quat /= np.linalg.norm(quat)
        qpos[joint_name] = quat

      elif joint_type == free:
        quat = random.rand(4)
        quat /= np.linalg.norm(quat)
        qpos[joint_name][3:] = quat