def initialize_episode(self, physics: 'mjcf.Physics',
random_state: np.random.RandomState):
"""Randomly selects a starting point and set the walker."""
self._get_clip_to_track(random_state)
# Set the walker at the beginning of the clip.
self._set_walker(physics)
self._walker_features = utils.get_features(physics, self._walker)
self._walker_features_prev = utils.get_features(physics, self._walker)
self._walker_joints = np.array(physics.bind(self._walker.mocap_joints).qpos) # pytype: disable=attribute-error
# compute initial error
self._compute_termination_error()
# assert error is 0 at initialization. In particular this will prevent
# a proto/walker mismatch.
if self._termination_error > 1e-2:
raise ValueError(('The termination exceeds 1e-2 at initialization. '
'This is likely due to a proto/walker mismatch.'))
self._update_ghost(physics)
# reset reward channels
self._reset_reward_channels()
def after_step(self, physics: 'mjcf.Physics', random_state):
"""Update the data after step."""
super().after_step(physics, random_state)
self._time_step += 1
# Update the walker's data for this timestep.
self._walker_features = utils.get_features(physics, self._walker)
# features for default error
self._walker_joints = np.array(
physics.bind(self._walker.mocap_joints).qpos) # pytype: disable=attribute-error
self._current_reference_features = {
k: v[self._time_step].copy()
for k, v in self._clip_reference_features.items()
}
# Error.
self._compute_termination_error()
# Terminate based on the error.
self._end_mocap = self._time_step == self._last_step
self._reference_observations.update(
self.get_all_reference_observations(physics))
self._update_ghost(physics)