Deprecate composer.Environment.step_spec, add reward_spec and discount_spec

alimuldal · alimuldal · commit 4b1392300a31 · 2019-03-22T17:29:32.000Z
PiperOrigin-RevId: 239807451
diff --git a/dm_control/composer/environment.py b/dm_control/composer/environment.py
@@ -331,7 +331,11 @@ def _reset_attempt(self):
         discount=None,
         observation=self._observation_updater.get_observation())
 
+  # TODO(b/129061424): Remove this method.
   def step_spec(self):
+    """DEPRECATED: please use `reward_spec` and `discount_spec` instead."""
+    warnings.warn('`step_spec` is deprecated, please use `reward_spec` and '
+                  '`discount_spec` instead.', DeprecationWarning)
     if (self._task.get_reward_spec() is None or
         self._task.get_discount_spec() is None):
       raise NotImplementedError
@@ -400,6 +404,38 @@ def action_spec(self):
     """Returns the action specification for this environment."""
     return self._task.action_spec(self._physics_proxy)
 
+  def reward_spec(self):
+    """Describes the reward returned by this environment.
+
+    This will be the output of `self.task.reward_spec()` if it is not None,
+    otherwise it will be the default spec returned by
+    `environment.Base.reward_spec()`.
+
+    Returns:
+      An `ArraySpec`, or a nested dict, list or tuple of `ArraySpec`s.
+    """
+    task_reward_spec = self._task.get_reward_spec()
+    if task_reward_spec is not None:
+      return task_reward_spec
+    else:
+      return super(Environment, self).reward_spec()
+
+  def discount_spec(self):
+    """Describes the discount returned by this environment.
+
+    This will be the output of `self.task.discount_spec()` if it is not None,
+    otherwise it will be the default spec returned by
+    `environment.Base.discount_spec()`.
+
+    Returns:
+      An `ArraySpec`, or a nested dict, list or tuple of `ArraySpec`s.
+    """
+    task_discount_spec = self._task.get_discount_spec()
+    if task_discount_spec is not None:
+      return task_discount_spec
+    else:
+      return super(Environment, self).discount_spec()
+
   def observation_spec(self):
     """Returns the observation specification for this environment.
 
diff --git a/dm_control/composer/environment_test.py b/dm_control/composer/environment_test.py
@@ -21,18 +21,26 @@
 
 # Internal dependencies.
 from absl.testing import absltest
+from absl.testing import parameterized
 from dm_control import composer
 from dm_control import mjcf
+import mock
 from six.moves import range
 
 
-class TaskWithResetFailures(composer.NullTask):
+class DummyTask(composer.NullTask):
+
+  def __init__(self):
+    null_entity = composer.ModelWrapperEntity(mjcf.RootElement())
+    super(DummyTask, self).__init__(null_entity)
+
+
+class DummyTaskWithResetFailures(DummyTask):
 
   def __init__(self, num_reset_failures):
+    super(DummyTaskWithResetFailures, self).__init__()
     self.num_reset_failures = num_reset_failures
     self.reset_counter = 0
-    null_entity = composer.ModelWrapperEntity(mjcf.RootElement())
-    super(TaskWithResetFailures, self).__init__(null_entity)
 
   def initialize_episode_mjcf(self, random_state):
     self.reset_counter += 1
@@ -42,19 +50,38 @@ def initialize_episode(self, physics, random_state):
       raise composer.EpisodeInitializationError()
 
 
-class EnvironmentTest(absltest.TestCase):
+class EnvironmentTest(parameterized.TestCase):
 
   def test_failed_resets(self):
     total_reset_failures = 5
     env_reset_attempts = 2
-    task = TaskWithResetFailures(num_reset_failures=total_reset_failures)
+    task = DummyTaskWithResetFailures(num_reset_failures=total_reset_failures)
     env = composer.Environment(task, max_reset_attempts=env_reset_attempts)
     for _ in range(total_reset_failures // env_reset_attempts):
       with self.assertRaises(composer.EpisodeInitializationError):
         env.reset()
     env.reset()  # should not raise an exception
     self.assertEqual(task.reset_counter, total_reset_failures + 1)
 
+  @parameterized.parameters(
+      dict(name='reward_spec', defined_in_task=True),
+      dict(name='reward_spec', defined_in_task=False),
+      dict(name='discount_spec', defined_in_task=True),
+      dict(name='discount_spec', defined_in_task=False))
+  def test_get_spec(self, name, defined_in_task):
+    task = DummyTask()
+    env = composer.Environment(task)
+    with mock.patch.object(task, 'get_' + name) as mock_task_get_spec:
+      if defined_in_task:
+        expected_spec = mock.Mock()
+        mock_task_get_spec.return_value = expected_spec
+      else:
+        expected_spec = getattr(super(composer.Environment, env), name)()
+        mock_task_get_spec.return_value = None
+      spec = getattr(env, name)()
+    mock_task_get_spec.assert_called_once_with()
+    self.assertSameStructure(spec, expected_spec)
+
 
 if __name__ == '__main__':
   absltest.main()
