Add support for evaluating arbitrary nested structures of "variation" callables

alimuldal · saran-t · commit 6cb8f775a314 · 2019-11-15T17:20:09.000Z
This introduces a dependency on `dm-tree`.

I also corrected the name of the `dm-env` package in `setup.py` and `requirements.txt` (the module is called `dm_env`, the package is called `dm-env`).

PiperOrigin-RevId: 280667371
Change-Id: I5d23f354482fa9e8cad84e88bba54f20d25cf3d3
diff --git a/dm_control/composer/variation/variation_values.py b/dm_control/composer/variation/variation_values.py
@@ -1,4 +1,4 @@
-# Copyright 2018 The dm_control Authors.
+# Copyright 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.
@@ -13,22 +13,27 @@
 # limitations under the License.
 # ============================================================================
 
-"""Utility function for evaluating callables or constants."""
+"""Utilities for handling nested structures of callables or constants."""
 
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
+import tree
 
-def evaluate(x, *args, **kwargs):
-  """Evaluates a callable or constant value.
+
+def evaluate(structure, *args, **kwargs):
+  """Evaluates a arbitrarily nested structure of callables or constant values.
 
   Args:
-    x: Either a callable or a constant value.
-    *args: Positional arguments passed to `x` if `x` is callable.
-    **kwargs: Keyword arguments passed to `x` if `x` is callable.
+    structure: An arbitrarily nested structure of callables or constant values.
+      By "structures", we mean lists, tuples, namedtuples, or dicts.
+    *args: Positional arguments passed to each callable in `structure`.
+    **kwargs: Keyword arguments passed to each callable in `structure.
 
   Returns:
-    Either the result of calling `x` if `x` is callable or else `x`.
+    The same nested structure, with each callable replaced by the value returned
+    by calling it.
   """
-  return x(*args, **kwargs) if callable(x) else x
+  return tree.map_structure(
+      lambda x: x(*args, **kwargs) if callable(x) else x, structure)
diff --git a/requirements.txt b/requirements.txt
@@ -1,5 +1,6 @@
 absl-py==0.7.0
-dm_env
+dm-env
+dm-tree
 enum34==1.1.6
 future==0.17.1
 futures==3.1.1
diff --git a/setup.py b/setup.py
@@ -172,7 +172,8 @@ def find_data_files(package_dir, patterns):
     install_requires=[
         'absl-py>=0.7.0',
         'enum34; python_version < "3.4"',
-        'dm_env',
+        'dm-env',
+        'dm-tree',
         'future',
         'futures; python_version == "2.7"',
         'glfw',
