Class Balance

dm_control/suite/cartpole.py:151–225

A Cartpole `Task` to balance the pole. State is initialized either close to the target configuration or at a random configuration.

class Balance(base.Task):
  """A Cartpole `Task` to balance the pole.

  State is initialized either close to the target configuration or at a random
  configuration.
  """
  _CART_RANGE = (-.25, .25)
  _ANGLE_COSINE_RANGE = (.995, 1)

  def __init__(self, swing_up, sparse, random=None):
    """Initializes an instance of `Balance`.

    Args:
      swing_up: A `bool`, which if `True` sets the cart to the middle of the
        slider and the pole pointing towards the ground. Otherwise, sets the
        cart to a random position on the slider and the pole to a random
        near-vertical position.
      sparse: A `bool`, whether to return a sparse or a smooth reward.
      random: Optional, either a `numpy.random.RandomState` instance, an
        integer seed for creating a new `RandomState`, or None to select a seed
        automatically (default).
    """
    self._sparse = sparse
    self._swing_up = swing_up
    super().__init__(random=random)

  def initialize_episode(self, physics):
    """Sets the state of the environment at the start of each episode.

    Initializes the cart and pole according to `swing_up`, and in both cases
    adds a small random initial velocity to break symmetry.

    Args:
      physics: An instance of `Physics`.
    """
    nv = physics.model.nv
    if self._swing_up:
      physics.named.data.qpos['slider'] = .01*self.random.randn()
      physics.named.data.qpos['hinge_1'] = np.pi + .01*self.random.randn()
      physics.named.data.qpos[2:] = .1*self.random.randn(nv - 2)
    else:
      physics.named.data.qpos['slider'] = self.random.uniform(-.1, .1)
      physics.named.data.qpos[1:] = self.random.uniform(-.034, .034, nv - 1)
    physics.named.data.qvel[:] = 0.01 * self.random.randn(physics.model.nv)
    super().initialize_episode(physics)

  def get_observation(self, physics):
    """Returns an observation of the (bounded) physics state."""
    obs = collections.OrderedDict()
    obs['position'] = physics.bounded_position()
    obs['velocity'] = physics.velocity()
    return obs

  def _get_reward(self, physics, sparse):
    if sparse:
      cart_in_bounds = rewards.tolerance(physics.cart_position(),
                                         self._CART_RANGE)
      angle_in_bounds = rewards.tolerance(physics.pole_angle_cosine(),