Method set_state

dm_control/mujoco/engine.py:252–285

Sets the physics state. Args: physics_state: NumPy array containing the full physics simulation state. sig: Optional integer, if specified then physics_state is passed directly to `mj_setState` with `sig`. Raises: ValueError: If `physics_state` has invalid size.

(self, physics_state, sig=None)

      return retval

  def set_state(self, physics_state, sig=None):
    """Sets the physics state.

    Args:
      physics_state: NumPy array containing the full physics simulation state.
      sig: Optional integer, if specified then physics_state is passed directly
        to `mj_setState` with `sig`.

    Raises:
      ValueError: If `physics_state` has invalid size.
    """
    if sig is None:
      state_items = self._physics_state_items()

      expected_shape = (sum(item.size for item in state_items),)
      if expected_shape != physics_state.shape:
        raise ValueError(
            'Input physics state has shape {}. Expected {}.'.format(
                physics_state.shape, expected_shape
            )
        )

      start = 0
      for state_item in state_items:
        size = state_item.size
        np.copyto(state_item, physics_state[start:start + size])
        start += size
    else:
      mujoco.mj_setState(
          self.model.ptr,
          self.data.ptr,
          np.asarray(physics_state, np.float64),
          sig,
      )

  def copy(self, share_model=False):
    """Creates a copy of this `Physics` instance.