Method initialize_episode

dm_control/suite/stacker.py:139–178

Sets the state of the environment at the start of each episode.

(self, physics)

    super().__init__(random=random)

  def initialize_episode(self, physics):
    """Sets the state of the environment at the start of each episode."""
    # Local aliases
    randint = self.random.randint
    uniform = self.random.uniform
    model = physics.named.model
    data = physics.named.data

    # Find a collision-free random initial configuration.
    penetrating = True
    while penetrating:

      # Randomise angles of arm joints.
      is_limited = model.jnt_limited[_ARM_JOINTS].astype(bool)
      joint_range = model.jnt_range[_ARM_JOINTS]
      lower_limits = np.where(is_limited, joint_range[:, 0], -np.pi)
      upper_limits = np.where(is_limited, joint_range[:, 1], np.pi)
      angles = uniform(lower_limits, upper_limits)
      data.qpos[_ARM_JOINTS] = angles

      # Symmetrize hand.
      data.qpos['finger'] = data.qpos['thumb']

      # Randomise target location.
      target_height = 2*randint(self._n_boxes) + 1
      box_size = model.geom_size['target', 0]
      model.body_pos['target', 'z'] = box_size * target_height
      model.body_pos['target', 'x'] = uniform(-.37, .37)

      # Randomise box locations.
      for name in self._box_names:
        data.qpos[name + '_x'] = uniform(.1, .3)
        data.qpos[name + '_z'] = uniform(0, .7)
        data.qpos[name + '_y'] = uniform(0, 2*np.pi)

      # Check for collisions.
      physics.after_reset()
      penetrating = physics.data.ncon > 0

    super().initialize_episode(physics)

  def get_observation(self, physics):
    """Returns either features or only sensors (to be used with pixels)."""