Allow custom spec types to be passed through in composer obs updater.

PiperOrigin-RevId: 530928000
Change-Id: If155fb61320163c4f256300771dcc836bc7a6aae

Author: DeepMind

dm_control/composer/observation/updater.py

@@ -195,6 +195,16 @@ def make_observation_spec_dict(enabled_dict):
       out_dict = type(enabled_dict)()
       for name, enabled in enabled_dict.items():
 
+        if (enabled.observable.aggregator is None
+            and enabled.observable.array_spec is not None):
+          # If possible, keep the original array spec, just updating the name
+          # and modifying the dimension for buffering. Doing this allows for
+          # custom spec types to be exposed by the environment where possible.
+          out_dict[name] = enabled.observable.array_spec.replace(
+              name=name, shape=enabled.buffer.shape
+          )
+          continue
+
         if isinstance(enabled.observable.array_spec, specs.BoundedArray):
           bounds = (enabled.observable.array_spec.minimum,
                     enabled.observable.array_spec.maximum)

dm_control/composer/observation/updater_test.py

@@ -54,6 +54,17 @@ def array_spec(self):
                               maximum=self._bounds[1])
 
 
+class MyArraySpec(specs.Array):
+  pass
+
+
+class GenericObservableWithMyArraySpec(observable.Generic):
+  @property
+  def array_spec(self):
+    datum = np.array(self(None, None))
+    return MyArraySpec(shape=datum.shape, dtype=datum.dtype)
+
+
 class UpdaterTest(parameterized.TestCase):
 
   @parameterized.parameters(list, tuple)
@@ -163,6 +174,23 @@ def testObservationSpecInference(self):
     self.assertCorrectSpec(spec['matrix'], (4, 2, 3), int, 'matrix')
     self.assertCorrectSpec(spec['sqrt'], (3,), float, 'sqrt')
 
+  def testCustomSpecTypePassedThrough(self):
+    physics = fake_physics.FakePhysics()
+    physics.observables['two_twos'] = GenericObservableWithMyArraySpec(
+        lambda _: [2.0, 2.0], buffer_size=3
+    )
+
+    physics.observables['two_twos'].enabled = True
+
+    observation_updater = updater.Updater(physics.observables)
+    observation_updater.reset(physics=physics, random_state=None)
+
+    spec = observation_updater.observation_spec()
+    self.assertIsInstance(spec['two_twos'], MyArraySpec)
+    self.assertEqual(spec['two_twos'].shape, (3, 2))
+    self.assertEqual(spec['two_twos'].dtype, float)
+    self.assertEqual(spec['two_twos'].name, 'two_twos')
+
   @parameterized.parameters(True, False)
   def testObservation(self, pad_with_initial_value):
     physics = fake_physics.FakePhysics()