Factor out the creation of observation updater from the _recompile_physics method.

saran-t · alimuldal · commit c4e60f45c5d7 · 2018-11-12T18:02:04.000Z
PiperOrigin-RevId: 220705572
diff --git a/dm_control/composer/environment.py b/dm_control/composer/environment.py
@@ -224,9 +224,7 @@ def _recompile_physics(self):
     self._physics_proxy = weakref.proxy(self._physics)
     self._hooks.refresh_entity_hooks()
 
-    self._observation_updater = observation.Updater(
-        self._task.observables, self._task.physics_steps_per_control_step,
-        self._strip_singleton_obs_buffer_dim)
+    self._observation_updater = self._make_observation_updater()
     self._observation_updater.reset(self._physics_proxy, self._random_state)
 
     # NB: `physics.after_reset()` is called after `task.after_compile()` so
@@ -240,6 +238,11 @@ def _recompile_physics(self):
   def _make_physics(self):
     return mjcf.Physics.from_mjcf_model(self._task.root_entity.mjcf_model)
 
+  def _make_observation_updater(self):
+    return observation.Updater(
+        self._task.observables, self._task.physics_steps_per_control_step,
+        self._strip_singleton_obs_buffer_dim)
+
   @property
   def physics(self):
     """Returns a `weakref.ProxyType` pointing to the current `mjcf.Physics`.
