Add math module to variation package and make sequence variation evaluate items.

PiperOrigin-RevId: 495758194
Change-Id: I8fae6c2709e4a0d37c985863cde47ee0dea61057

Author: sbohez

dm_control/composer/variation/deterministic.py

@@ -17,6 +17,7 @@
 
 
 from dm_control.composer.variation import base
+from dm_control.composer.variation.variation_values import evaluate
 
 
 class Constant(base.Variation):
@@ -42,10 +43,12 @@ def __init__(self, values):
 
   def __call__(self, initial_value=None, current_value=None, random_state=None):
     try:
-      return next(self._iterator)
+      return evaluate(next(self._iterator), initial_value=initial_value,
+                      current_value=current_value, random_state=random_state)
     except StopIteration:
       self._iterator = iter(self._values)
-      return next(self._iterator)
+      return evaluate(next(self._iterator), initial_value=initial_value,
+                      current_value=current_value, random_state=random_state)
 
 
 class Identity(base.Variation):

dm_control/composer/variation/distributions.py

@@ -17,8 +17,8 @@
 import abc
 import functools
 
-from dm_control.composer import variation
 from dm_control.composer.variation import base
+from dm_control.composer.variation.variation_values import evaluate
 import numpy as np
 
 
@@ -45,12 +45,12 @@ def __call__(self, initial_value=None, current_value=None, random_state=None):
     size = (
         None if self._single_sample or initial_value is None  # pylint: disable=g-long-ternary
         else np.shape(initial_value))
-    local_args = variation.evaluate(
+    local_args = evaluate(
         self._args,
         initial_value=initial_value,
         current_value=current_value,
         random_state=random_state)
-    local_kwargs = variation.evaluate(
+    local_kwargs = evaluate(
         self._kwargs,
         initial_value=initial_value,
         current_value=current_value,

dm_control/composer/variation/math.py

@@ -0,0 +1,83 @@
+# Copyright 2018 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.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or  implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+
+"""Math operations on variation objects."""
+
+import abc
+
+from dm_control.composer.variation import base
+from dm_control.composer.variation.variation_values import evaluate
+
+import numpy as np
+
+
+class MathOp(base.Variation):
+  """Base MathOp class for applying math operations on variation objects.
+
+  Subclasses need to implement `_op`, which takes in a single value and applies
+  the desired math operation. This operation gets applied to the result of the
+  evaluated base variation object passed at construction. Structured variation
+  objects are automatically traversed.
+  """
+
+  def __init__(self, *args, **kwargs):
+    self._args = args
+    self._kwargs = kwargs
+
+  def __call__(self, initial_value=None, current_value=None, random_state=None):
+    local_args = evaluate(
+        self._args,
+        initial_value=initial_value,
+        current_value=current_value,
+        random_state=random_state)
+    local_kwargs = evaluate(
+        self._kwargs,
+        initial_value=initial_value,
+        current_value=current_value,
+        random_state=random_state)
+    return self._callable(*local_args, **local_kwargs)
+
+  @property
+  @abc.abstractmethod
+  def _callable(self):
+    pass
+
+
+class Log(MathOp):
+
+  @property
+  def _callable(self):
+    return np.log
+
+
+class Max(MathOp):
+
+  @property
+  def _callable(self):
+    return np.max
+
+
+class Min(MathOp):
+
+  @property
+  def _callable(self):
+    return np.min
+
+
+class Norm(MathOp):
+
+  @property
+  def _callable(self):
+    return np.linalg.norm