Add support for multi-array action_specs in viewer.

DeepMind · alimuldal · commit ea433b44a0df · 2019-02-08T16:46:15.000Z
PiperOrigin-RevId: 230801344
diff --git a/dm_control/viewer/runtime.py b/dm_control/viewer/runtime.py
@@ -1,4 +1,4 @@
-# Copyright 2018 The dm_control Authors.
+# Copyright 2018-2019 The dm_control Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -42,11 +42,16 @@ def _get_default_action(action_spec):
   * For unbounded intervals this will be zero.
 
   Args:
-    action_spec: An instance of `BoundedArraySpec`.
+    action_spec: An instance of `BoundedArraySpec` or a list or tuple
+      containing these.
 
   Returns:
-    A numpy array of actions.
+    A numpy array of actions if `action_spec` is a single `BoundedArraySpec`, or
+    a tuple of such arrays if `action_spec` is a list or tuple.
   """
+  if isinstance(action_spec, (list, tuple)):
+    return tuple(_get_default_action(spec) for spec in action_spec)
+
   minimum = np.broadcast_to(action_spec.minimum, action_spec.shape)
   maximum = np.broadcast_to(action_spec.maximum, action_spec.shape)
   left_bounded = np.isfinite(minimum)
@@ -55,7 +60,9 @@ def _get_default_action(action_spec):
       condlist=[left_bounded & right_bounded, left_bounded, right_bounded],
       choicelist=[0.5 * (minimum + maximum), minimum, maximum],
       default=0.)
-  return action.astype(action_spec.dtype, copy=False)
+  action = action.astype(action_spec.dtype, copy=False)
+  action.flags.writeable = False
+  return action
 
 
 class State(enum.Enum):
@@ -233,7 +240,7 @@ def _step(self):
       if self._policy:
         action = self._policy(self._time_step)
       else:
-        action = self._default_action.copy()
+        action = self._default_action
       self._time_step = self._env.step(action)
       self._last_action = action
       finished = self._time_step.last()
diff --git a/dm_control/viewer/runtime_test.py b/dm_control/viewer/runtime_test.py
@@ -1,4 +1,4 @@
-# Copyright 2018 The dm_control Authors.
+# Copyright 2018-2019 The dm_control Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -23,6 +23,7 @@
 from dm_control.viewer import runtime
 import mock
 import numpy as np
+from six.moves import zip
 from dm_control.rl import environment
 from dm_control.rl import specs
 
@@ -288,8 +289,7 @@ def test_step_without_policy(self):
         runtime.__name__ + '._get_default_action') as mock_get_default_action:
       this_runtime = runtime.Runtime(environment=self.env, policy=None)
     this_runtime._step()
-    self.env.step.assert_called_once_with(
-        mock_get_default_action.return_value.copy())
+    self.env.step.assert_called_once_with(mock_get_default_action.return_value)
 
   def test_stepping_paused(self):
     with mock.patch(runtime.__name__ + '.mjlib') as mjlib:
@@ -336,5 +336,48 @@ def raise_exception(*unused_args, **unused_kwargs):
     self.assertTrue(finished)
 
 
+class DefaultActionFromSpecTest(parameterized.TestCase):
+
+  def assertNestedArraysEqual(self, expected, actual):
+    """Asserts that two potentially nested structures of arrays are equal."""
+    if isinstance(expected, (list, tuple)):
+      self.assertIsInstance(actual, (list, tuple))
+      self.assertLen(actual, len(expected))
+      for expected_item, actual_item in zip(expected, actual):
+        self.assertNestedArraysEqual(expected_item, actual_item)
+    else:
+      np.testing.assert_array_equal(expected, actual)
+
+  _SHAPE = (2,)
+  _DTYPE = np.float64
+  _ACTION = np.zeros(_SHAPE)
+  _ACTION_SPEC = specs.BoundedArraySpec(_SHAPE, np.float64, -1, 1)
+
+  @parameterized.named_parameters(
+      ('single_array', _ACTION_SPEC, _ACTION),
+      ('tuple', (_ACTION_SPEC, _ACTION_SPEC), (_ACTION, _ACTION)),
+      ('list', [_ACTION_SPEC, _ACTION_SPEC], (_ACTION, _ACTION)))
+  def test_action_structure(self, action_spec, expected_action):
+    self.assertNestedArraysEqual(expected_action,
+                                 runtime._get_default_action(action_spec))
+
+  @parameterized.named_parameters(
+      ('closed',
+       specs.BoundedArraySpec(_SHAPE, _DTYPE, minimum=1., maximum=2.),
+       np.full(_SHAPE, fill_value=1.5, dtype=_DTYPE)),
+      ('left_open',
+       specs.BoundedArraySpec(_SHAPE, _DTYPE, minimum=-np.inf, maximum=2.),
+       np.full(_SHAPE, fill_value=2., dtype=_DTYPE)),
+      ('right_open',
+       specs.BoundedArraySpec(_SHAPE, _DTYPE, minimum=1., maximum=np.inf),
+       np.full(_SHAPE, fill_value=1., dtype=_DTYPE)),
+      ('unbounded',
+       specs.BoundedArraySpec(_SHAPE, _DTYPE, minimum=-np.inf, maximum=np.inf),
+       np.full(_SHAPE, fill_value=0., dtype=_DTYPE)))
+  def test_action_spec_interval(self, action_spec, expected_action):
+    self.assertNestedArraysEqual(expected_action,
+                                 runtime._get_default_action(action_spec))
+
+
 if __name__ == '__main__':
   absltest.main()
