Add option to step the physics with mj_step instead of mj_step2->mj_step1.

PiperOrigin-RevId: 514728526
Change-Id: Icfac463419f1f9365b90ce36d4d96b0d5311ec73

Author: kevinzakka

dm_control/composer/environment.py

@@ -169,7 +169,8 @@ def __init__(self, task, time_limit=float('inf'), random_state=None,
                n_sub_steps=None,
                raise_exception_on_physics_error=True,
                strip_singleton_obs_buffer_dim=False,
-               delayed_observation_padding=ObservationPadding.ZERO):
+               delayed_observation_padding=ObservationPadding.ZERO,
+               legacy_step: bool = True):
     """Initializes an instance of `_CommonEnvironment`.
 
     Args:
@@ -194,6 +195,9 @@ def __init__(self, task, time_limit=float('inf'), random_state=None,
         observables. If `ZERO` then the buffer is initially filled with zeroes.
         If `INITIAL_VALUE` then the buffer is initially filled with the first
         observation values.
+      legacy_step: If True, steps the state with up-to-date position and
+        velocity dependent fields. See Page 6 of
+        https://arxiv.org/abs/2006.12983 for more information.
     """
     if not isinstance(delayed_observation_padding, ObservationPadding):
       raise ValueError(
@@ -210,6 +214,7 @@ def __init__(self, task, time_limit=float('inf'), random_state=None,
     self._raise_exception_on_physics_error = raise_exception_on_physics_error
     self._strip_singleton_obs_buffer_dim = strip_singleton_obs_buffer_dim
     self._delayed_observation_padding = delayed_observation_padding
+    self._legacy_step = legacy_step
 
     if n_sub_steps is not None:
       warnings.simplefilter('once', DeprecationWarning)
@@ -248,6 +253,7 @@ def _recompile_physics(self):
       self._physics.free()
     self._physics = mjcf.Physics.from_mjcf_model(
         self._task.root_entity.mjcf_model)
+    self._physics.legacy_step = self._legacy_step
 
   def _make_observation_updater(self):
     pad_with_initial_value = (
@@ -291,7 +297,8 @@ def __init__(self, task, time_limit=float('inf'), random_state=None,
                raise_exception_on_physics_error=True,
                strip_singleton_obs_buffer_dim=False,
                max_reset_attempts=1,
-               delayed_observation_padding=ObservationPadding.ZERO):
+               delayed_observation_padding=ObservationPadding.ZERO,
+               legacy_step: bool = True):
     """Initializes an instance of `Environment`.
 
     Args:
@@ -320,6 +327,8 @@ def __init__(self, task, time_limit=float('inf'), random_state=None,
         observables. If `ZERO` then the buffer is initially filled with zeroes.
         If `INITIAL_VALUE` then the buffer is initially filled with the first
         observation values.
+      legacy_step: If True, steps the state with up-to-date position and
+        velocity dependent fields.
     """
     super().__init__(
         task=task,
@@ -328,7 +337,8 @@ def __init__(self, task, time_limit=float('inf'), random_state=None,
         n_sub_steps=n_sub_steps,
         raise_exception_on_physics_error=raise_exception_on_physics_error,
         strip_singleton_obs_buffer_dim=strip_singleton_obs_buffer_dim,
-        delayed_observation_padding=delayed_observation_padding)
+        delayed_observation_padding=delayed_observation_padding,
+        legacy_step=legacy_step)
     self._max_reset_attempts = max_reset_attempts
     self._reset_next_step = True
 

dm_control/mujoco/engine.py

@@ -144,30 +144,37 @@ def set_control(self, control):
     """
     np.copyto(self.data.ctrl, control)
 
-  def step(self, nstep=1):
-    """Advances physics with up-to-date position and velocity dependent fields.
-
-    Args:
-      nstep: Optional integer, number of steps to take.
-
-    The actuation can be updated by calling the `set_control` function first.
-    """
+  def _step_with_up_to_date_position_velocity(self, nstep: int = 1) -> None:
+    """Physics step with up-to-date position and velocity dependent fields."""
     # In the case of Euler integration we assume mj_step1 has already been
     # called for this state, finish the step with mj_step2 and then update all
     # position and velocity related fields with mj_step1. This ensures that
     # (most of) mjData is in sync with qpos and qvel. In the case of non-Euler
     # integrators (e.g. RK4) an additional mj_step1 must be called after the
     # last mj_step to ensure mjData syncing.
+    if self.model.opt.integrator != mujoco.mjtIntegrator.mjINT_RK4.value:
+      mujoco.mj_step2(self.model.ptr, self.data.ptr)
+      if nstep > 1:
+        mujoco.mj_step(self.model.ptr, self.data.ptr, nstep-1)
+    else:
+      mujoco.mj_step(self.model.ptr, self.data.ptr, nstep)
+
+    mujoco.mj_step1(self.model.ptr, self.data.ptr)
+
+  def step(self, nstep: int = 1) -> None:
+    """Advances the physics state by `nstep`s.
+
+    Args:
+      nstep: Optional integer, number of steps to take.
+
+    The actuation can be updated by calling the `set_control` function first.
+    """
     with self.check_invalid_state():
-      if self.model.opt.integrator != mujoco.mjtIntegrator.mjINT_RK4.value:
-        mujoco.mj_step2(self.model.ptr, self.data.ptr)
-        if nstep > 1:
-          mujoco.mj_step(self.model.ptr, self.data.ptr, nstep-1)
+      if self.legacy_step:
+        self._step_with_up_to_date_position_velocity(nstep)
       else:
         mujoco.mj_step(self.model.ptr, self.data.ptr, nstep)
 
-      mujoco.mj_step1(self.model.ptr, self.data.ptr)
-
   def render(
       self,
       height=240,

dm_control/rl/control.py

@@ -34,7 +34,8 @@ def __init__(self,
                time_limit=float('inf'),
                control_timestep=None,
                n_sub_steps=None,
-               flat_observation=False):
+               flat_observation=False,
+               legacy_step: bool = True):
     """Initializes a new `Environment`.
 
     Args:
@@ -48,12 +49,16 @@ def __init__(self,
         `control_timestep` is not specified.
       flat_observation: If True, observations will be flattened and concatenated
         into a single numpy array.
+      legacy_step: If True, steps the state with up-to-date position and
+        velocity dependent fields. See Page 6 of
+        https://arxiv.org/abs/2006.12983 for more information.
 
     Raises:
       ValueError: If both `n_sub_steps` and `control_timestep` are supplied.
     """
     self._task = task
     self._physics = physics
+    self._physics.legacy_step = legacy_step
     self._flat_observation = flat_observation
 
     if n_sub_steps is not None and control_timestep is not None:
@@ -201,6 +206,8 @@ def _spec_from_observation(observation):
 class Physics(metaclass=abc.ABCMeta):
   """Simulates a physical environment."""
 
+  legacy_step: bool = True
+
   @abc.abstractmethod
   def step(self, n_sub_steps=1):
     """Updates the simulation state.