def initialize_episode_mjcf(self, random_state):
super(RepeatSingleGoalMazeAugmentedWithTargets,
self).initialize_episode_mjcf(random_state)
subtarget_positions = self._maze_arena.target_positions
for pos, subtarget in zip(subtarget_positions, self._subtargets):
mjcf.get_attachment_frame(subtarget.mjcf_model).pos = pos
self._subtarget_rewarded = [False] * len(self._subtargets)
def set_walker(physics,
walker,
qpos,
qvel,
offset=0,
null_xyz_and_yaw=False,
position_shift=None,
rotation_shift=None):
"""Set the freejoint and walker's joints angles and velocities."""
qpos = np.array(qpos)
if null_xyz_and_yaw:
qpos[:2] = 0.
euler = tr.quat_to_euler(qpos[3:7], ordering='ZYX')
euler[0] = 0.
quat = tr.euler_to_quat(euler, ordering='ZYX')
qpos[3:7] = quat
qpos[:3] += offset
freejoint = mjcf.get_attachment_frame(walker.mjcf_model).freejoint
physics.bind(freejoint).qpos = qpos[:7]
physics.bind(freejoint).qvel = qvel[:6]
physics.bind(walker.mocap_joints).qpos = qpos[7:]
physics.bind(walker.mocap_joints).qvel = qvel[6:]
if position_shift is not None or rotation_shift is not None:
walker.shift_pose(physics,
position=position_shift,
quaternion=rotation_shift,
rotate_velocity=True)
def get_features(physics, walker):
"""Get walker features for reward functions."""
walker_bodies = walker.mocap_tracking_bodies
walker_features = {}
root_pos, root_quat = walker.get_pose(physics)
walker_features['position'] = np.array(root_pos)
walker_features['quaternion'] = np.array(root_quat)
joints = np.array(physics.bind(walker.mocap_joints).qpos)
walker_features['joints'] = joints
freejoint_frame = mjcf.get_attachment_frame(walker.mjcf_model)
com = np.array(physics.bind(freejoint_frame).subtree_com)
walker_features['center_of_mass'] = com
end_effectors = np.array(
walker.observables.end_effectors_pos(physics)[:]).reshape(-1, 3)
walker_features['end_effectors'] = end_effectors
if hasattr(walker.observables, 'appendages_pos'):
appendages = np.array(
walker.observables.appendages_pos(physics)[:]).reshape(-1, 3)
else:
appendages = np.array(end_effectors)
walker_features['appendages'] = appendages
xpos = np.array(physics.bind(walker_bodies).xpos)
walker_features['body_positions'] = xpos
xquat = np.array(physics.bind(walker_bodies).xquat)
walker_features['body_quaternions'] = xquat
root_vel, root_angvel = walker.get_velocity(physics)
walker_features['velocity'] = np.array(root_vel)
walker_features['angular_velocity'] = np.array(root_angvel)
joints_vel = np.array(physics.bind(walker.mocap_joints).qvel)
walker_features['joints_velocity'] = joints_vel
return walker_features
def initialize_episode_mjcf(self, random_state):
super(
ManyHeterogeneousGoalsMaze, self).initialize_episode_mjcf(random_state)
for target in itertools.chain(*self._active_targets):
target.detach()
target_positions = list(self._maze_arena.target_positions)
random_state.shuffle(target_positions)
all_targets = self._get_targets(len(target_positions), random_state)
for pos, target in zip(target_positions, itertools.chain(*all_targets)):
self._maze_arena.attach(target)
mjcf.get_attachment_frame(target.mjcf_model).pos = pos
target.initialize_episode_mjcf(random_state)
self._active_targets = all_targets
self._target_rewarded = [[False] * len(targets) for targets in all_targets]
def initialize_episode_mjcf(self, random_state):
self._reward = 0.0
self._discount = 1.0
self._should_terminate = False
self._prop.detach()
if self._proto_modifier:
trajectory = self._trajectory.get_modified_trajectory(
self._proto_modifier)
self._prop = trajectory.create_props(
priority_friction=self._priority_friction)[0]
self._arena.add_free_entity(self._prop)
# set the bucket position for this episode
bucket_distance = 1.*random_state.rand()+self._bucket_offset
mjcf.get_attachment_frame(self._bucket.mjcf_model).pos = [bucket_distance,
0, 0]
def test_walker_velocity(self):
initializer = initializers.UniformInitializer()
env = _env(_home_team(2) + _away_team(2), initializer=initializer)
root_joints = []
non_root_joints = []
for player in env.task.players:
attachment_frame = mjcf.get_attachment_frame(player.walker.mjcf_model)
root_joints.extend(
attachment_frame.find_all("joint", immediate_children_only=True))
non_root_joints.extend(player.walker.mjcf_model.find_all("joint"))
# Assign a non-zero sentinel value to the velocities of all root and
# non-root joints.
sentinel_velocity = 3.14
env.physics.bind(root_joints + non_root_joints).qvel = sentinel_velocity
# The initializer should zero the velocities of the root joints, but not the
# non-root joints.
initializer(env.task, env.physics, env.random_state)
np.testing.assert_array_equal(env.physics.bind(non_root_joints).qvel,
sentinel_velocity)
np.testing.assert_array_equal(env.physics.bind(root_joints).qvel, 0.)
def set_pose(self, physics, position=None, quaternion=None):
"""Sets position and/or orientation of this entity relative to its parent.
If the entity is attached with a free joint, this method will set the
respective DoFs of the joint. If the entity is either fixed or attached with
a different joint type, this method will update the position and/or
quaternion of the attachment frame.
Note that the semantics differ slightly between the two cases: the DoFs of a
free body are specified in global coordinates, whereas the position of a
non-free body is specified in relative coordinates with respect to the
parent frame. However, for entities that are either attached directly to the
worldbody, or to other entities that are positioned at the global origin
(e.g. the arena), there is no difference between the two cases.
Args:
physics: An instance of `mjcf.Physics`.
position: (optional) A NumPy array of size 3.
quaternion: (optional) A NumPy array of size 4.
Raises:
RuntimeError: If the entity is not attached.
"""
root_joint = mjcf.get_frame_freejoint(self.mjcf_model)
if root_joint:
if position is not None:
physics.bind(root_joint).qpos[:3] = position
if quaternion is not None:
physics.bind(root_joint).qpos[3:] = quaternion
else:
attachment_frame = mjcf.get_attachment_frame(self.mjcf_model)
if attachment_frame is None:
raise RuntimeError(_NO_ATTACHMENT_FRAME)
if position is not None:
physics.bind(attachment_frame).pos = position
if quaternion is not None:
normalised_quaternion = quaternion / np.linalg.norm(quaternion)
physics.bind(attachment_frame).quat = normalised_quaternion
def get_pose(self, physics):
"""Get the position and orientation of this entity relative to its parent.
Note that the semantics differ slightly depending on whether or not the
entity has a free joint:
* If it has a free joint the position and orientation are always given in
global coordinates.
* If the entity is fixed or attached with a different joint type then the
position and orientation are given relative to the parent frame.
For entities that are either attached directly to the worldbody, or to other
entities that are positioned at the global origin (e.g. the arena) the
global and relative poses are equivalent.
Args:
physics: An instance of `mjcf.Physics`.
Returns:
A 2-tuple where the first entry is a (3,) numpy array representing the
position and the second is a (4,) numpy array representing orientation as
a quaternion.
Raises:
RuntimeError: If the entity is not attached.
"""
root_joint = mjcf.get_frame_freejoint(self.mjcf_model)
if root_joint:
position = physics.bind(root_joint).qpos[:3]
quaternion = physics.bind(root_joint).qpos[3:]
else:
attachment_frame = mjcf.get_attachment_frame(self.mjcf_model)
if attachment_frame is None:
raise RuntimeError(_NO_ATTACHMENT_FRAME)
position = physics.bind(attachment_frame).pos
quaternion = physics.bind(attachment_frame).quat
return position, quaternion
def initialize_episode_mjcf(self, random_state):
super().initialize_episode_mjcf(random_state)
self._target_position = self._maze_arena.target_positions[
random_state.randint(0, len(self._maze_arena.target_positions))]
mjcf.get_attachment_frame(
self._target.mjcf_model).pos = self._target_position
def prop_zaxis(physics):
prop_xmat = physics.bind(
mjcf.get_attachment_frame(self._prop.mjcf_model)).xmat
return prop_xmat[[2, 5, 8]]
def root_body(self):
if self.parent:
return mjcf.get_attachment_frame(self.mjcf_model)
else:
return self.mjcf_model.worldbody
