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