Pass random_state to observable attributes that are Variation objects.

sbohez · alimuldal · commit 9c0b199d1e0e · 2020-06-19T22:23:32.000+01:00
PiperOrigin-RevId: 317105039
Change-Id: I82c551e1ad99f726c56f96bb9ec189c87c106f8e
diff --git a/dm_control/composer/observation/updater.py b/dm_control/composer/observation/updater.py
@@ -20,8 +20,10 @@
 from __future__ import print_function
 
 import collections
+import functools
 from absl import logging
 
+from dm_control.composer import variation
 from dm_control.composer.observation import obs_buffer
 from dm_env import specs
 import numpy as np
@@ -37,6 +39,7 @@ class _EnabledObservable(object):
   """Encapsulates an enabled observable, its buffer, and its update schedule."""
 
   __slots__ = ('observable', 'observation_callable',
+               'update_interval', 'delay', 'buffer_size',
                'buffer', 'update_schedule')
 
   def __init__(self, observable, physics, random_state,
@@ -45,6 +48,16 @@ def __init__(self, observable, physics, random_state,
     self.observation_callable = (
         observable.observation_callable(physics, random_state))
 
+    self._bind_attribute_from_observable('update_interval',
+                                         DEFAULT_UPDATE_INTERVAL,
+                                         random_state)
+    self._bind_attribute_from_observable('delay',
+                                         DEFAULT_DELAY,
+                                         random_state)
+    self._bind_attribute_from_observable('buffer_size',
+                                         DEFAULT_BUFFER_SIZE,
+                                         random_state)
+
     obs_spec = self.observable.array_spec
     if obs_spec is None:
       # We take an observation to determine the shape and dtype of the array.
@@ -58,11 +71,22 @@ def __init__(self, observable, physics, random_state,
       obs_array = np.asarray(obs_array)
       obs_spec = specs.Array(shape=obs_array.shape, dtype=obs_array.dtype)
     self.buffer = obs_buffer.Buffer(
-        buffer_size=(observable.buffer_size or DEFAULT_BUFFER_SIZE),
+        buffer_size=self.buffer_size,
         shape=obs_spec.shape, dtype=obs_spec.dtype,
         strip_singleton_buffer_dim=strip_singleton_buffer_dim)
     self.update_schedule = collections.deque()
 
+  def _bind_attribute_from_observable(self, attr, default_value, random_state):
+    obs_attr = getattr(self.observable, attr)
+    if obs_attr:
+      if isinstance(obs_attr, variation.Variation):
+        setattr(self, attr,
+                functools.partial(obs_attr, random_state=random_state))
+      else:
+        setattr(self, attr, obs_attr)
+    else:
+      setattr(self, attr, default_value)
+
 
 def _call_if_callable(arg):
   if callable(arg):
@@ -137,7 +161,7 @@ def make_buffers_dict(observables):
 
     self._step_counter = 0
     for enabled in self._enabled_list:
-      first_delay = _call_if_callable(enabled.observable.delay or DEFAULT_DELAY)
+      first_delay = _call_if_callable(enabled.delay)
       enabled.buffer.insert(
           0, first_delay,
           enabled.observation_callable())
@@ -225,14 +249,11 @@ def prepare_for_next_control_step(self):
     if self._enabled_structure is None:
       raise RuntimeError('`reset` must be called before `before_step`.')
     for enabled in self._enabled_list:
-      update_interval = (
-          enabled.observable.update_interval or DEFAULT_UPDATE_INTERVAL)
-      delay = enabled.observable.delay or DEFAULT_DELAY
-      buffer_size = enabled.observable.buffer_size or DEFAULT_BUFFER_SIZE
-
-      if (update_interval == DEFAULT_UPDATE_INTERVAL and delay == DEFAULT_DELAY
-          and buffer_size < self._physics_steps_per_control_step):
-        for i in reversed(range(buffer_size)):
+
+      if (enabled.update_interval == DEFAULT_UPDATE_INTERVAL
+          and enabled.delay == DEFAULT_DELAY
+          and enabled.buffer_size < self._physics_steps_per_control_step):
+        for i in reversed(range(enabled.buffer_size)):
           next_step = (
               self._step_counter + self._physics_steps_per_control_step - i)
           next_delay = DEFAULT_DELAY
@@ -244,9 +265,9 @@ def prepare_for_next_control_step(self):
           last_scheduled_step = self._step_counter
         max_step = self._step_counter + 2 * self._physics_steps_per_control_step
         while last_scheduled_step < max_step:
-          next_update_interval = _call_if_callable(update_interval)
+          next_update_interval = _call_if_callable(enabled.update_interval)
           next_step = last_scheduled_step + next_update_interval
-          next_delay = _call_if_callable(delay)
+          next_delay = _call_if_callable(enabled.delay)
           enabled.update_schedule.append((next_step, next_delay))
           last_scheduled_step = next_step
         # Optimize the schedule by planning ahead and dropping unseen entries.
diff --git a/setup.py b/setup.py
@@ -166,7 +166,7 @@ def find_data_files(package_dir, patterns):
 
 setup(
     name='dm_control',
-    version='0.0.316448566',
+    version='0.0.317105039',
     description='Continuous control environments and MuJoCo Python bindings.',
     author='DeepMind',
     license='Apache License, Version 2.0',
