Implement repr and equality for all stateless variations in dm_control.

Equality is useful for tests, and the repr implementations are useful for debugging, especially for binary operations, as at the moment you get opaque strings.

PiperOrigin-RevId: 737684596
Change-Id: I2a3a466200efdfc97f60ff0e4b825333a3b44a00

Author: nimrod-gileadi

dm_control/composer/variation/base.py

@@ -91,9 +91,21 @@ def __init__(self, op, variation):
     self._op = op
     self._variation = variation
 
+  def __eq__(self, other):
+    if not isinstance(other, _UnaryOperation):
+      return False
+    return self._op == other._op and self._variation == other._variation
+
+  def __str__(self):
+    return f"{self._op.__name__}({self._variation})"
+
+  def __repr__(self):
+    return f"UnaryOperation({self._op.__name__}({self._variation}))"
+
   def __call__(self, initial_value=None, current_value=None, random_state=None):
     value = variation_values.evaluate(
-        self._variation, initial_value, current_value, random_state)
+        self._variation, initial_value, current_value, random_state
+    )
     return self._op(value)
 
 
@@ -105,21 +117,56 @@ def __init__(self, op, first, second):
     self._second = second
     self._op = op
 
+  def __eq__(self, other):
+    if not isinstance(other, _BinaryOperation):
+      return False
+    return (
+        self._op == other._op
+        and self._first == other._first
+        and self._second == other._second
+    )
+
+  def __str__(self):
+    return f"{self._op.__name__}({self._first}, {self._second})"
+
+  def __repr__(self):
+    return (
+        f"BinaryOperation({self._op.__name__}({self._first!r},"
+        f" {self._second!r}))"
+    )
+
   def __call__(self, initial_value=None, current_value=None, random_state=None):
     first_value = variation_values.evaluate(
-        self._first, initial_value, current_value, random_state)
+        self._first, initial_value, current_value, random_state
+    )
     second_value = variation_values.evaluate(
-        self._second, initial_value, current_value, random_state)
+        self._second, initial_value, current_value, random_state
+    )
     return self._op(first_value, second_value)
 
 
 class _GetItemOperation(Variation):
+  """Returns a single element from the output of a Variation."""
 
   def __init__(self, variation, index):
     self._variation = variation
     self._index = index
 
+  def __eq__(self, other):
+    if not isinstance(other, _GetItemOperation):
+      return False
+    return self._variation == other._variation and self._index == other._index
+
+  def __str__(self):
+    return f"{self._variation}[{self._index}]"
+
+  def __repr__(self):
+    return (
+        f"GetItemOperation({self._variation!r}[{self._index}])"
+    )
+
   def __call__(self, initial_value=None, current_value=None, random_state=None):
     value = variation_values.evaluate(
-        self._variation, initial_value, current_value, random_state)
+        self._variation, initial_value, current_value, random_state
+    )
     return np.asarray(value)[self._index]

dm_control/composer/variation/colors.py

@@ -43,6 +43,16 @@ def __call__(self, initial_value=None, current_value=None, random_state=None):
         variation_values.evaluate([self._r, self._g, self._b, self._alpha],
                                   initial_value, current_value, random_state))
 
+  def __eq__(self, other):
+    if not isinstance(other, RgbVariation):
+      return False
+    return (
+        self._r == other._r
+        and self._g == other._g
+        and self._b == other._b
+        and self._alpha == other._alpha
+    )
+
 
 class HsvVariation(base.Variation):
   """Represents a variation in the HSV color space.
@@ -62,6 +72,22 @@ def __call__(self, initial_value=None, current_value=None, random_state=None):
         random_state)
     return np.asarray(list(colorsys.hsv_to_rgb(h, s, v)) + [alpha])
 
+  def __eq__(self, other):
+    if not isinstance(other, HsvVariation):
+      return False
+    return (
+        self._h == other._h
+        and self._s == other._s
+        and self._v == other._v
+        and self._alpha == other._alpha
+    )
+
+  def __repr__(self):
+    return (
+        f"HsvVariation(h={self._h}, s={self._s}, v={self._v}, "
+        f"alpha={self._alpha})"
+    )
+
 
 class GrayVariation(HsvVariation):
   """Represents a variation in gray level.
@@ -73,3 +99,8 @@ class GrayVariation(HsvVariation):
 
   def __init__(self, gray_level, alpha=1.0):
     super().__init__(h=0.0, s=0.0, v=gray_level, alpha=alpha)
+
+  def __repr__(self):
+    return (
+        f"GrayVariation(gray_level={self._v}, alpha={self._alpha})"
+    )

dm_control/composer/variation/deterministic.py

@@ -33,6 +33,17 @@ def __init__(self, value):
   def __call__(self, initial_value=None, current_value=None, random_state=None):
     return self._value
 
+  def __eq__(self, other):
+    if not isinstance(other, Constant):
+      return False
+    return self._value == other._value
+
+  def __str__(self):
+    return f"{self._value}"
+
+  def __repr__(self):
+    return f"Constant({self._value!r})"
+
 
 class Sequence(base.Variation):
   """Variation representing a fixed sequence of values."""
@@ -52,6 +63,8 @@ def __call__(self, initial_value=None, current_value=None, random_state=None):
 
 
 class Identity(base.Variation):
-
   def __call__(self, initial_value=None, current_value=None, random_state=None):
     return current_value
+
+  def __eq__(self, other):
+    return isinstance(other, Identity)

dm_control/composer/variation/distributions.py

@@ -71,6 +71,22 @@ def __getattr__(self, name):
   def _callable(self, random_state):
     raise NotImplementedError
 
+  def __eq__(self, other):
+    if not isinstance(other, type(self)):
+      return False
+    return (
+        self._args == other._args
+        and self._kwargs == other._kwargs
+        and self._single_sample == other._single_sample
+    )
+
+  def __repr__(self):
+    return '{}(args={}, kwargs={}, single_sample={})'.format(
+        type(self).__name__,
+        self._args,
+        self._kwargs,
+        self._single_sample)
+
 
 class Uniform(Distribution):
   __slots__ = ()
@@ -118,6 +134,17 @@ def __call__(self, initial_value=None, current_value=None, random_state=None):
     axis /= np.linalg.norm(axis, axis=-1, keepdims=True)
     return axis
 
+  def __eq__(self, other):
+    if not isinstance(other, UniformPointOnSphere):
+      return False
+    return self._single_sample == other._single_sample
+
+  def __repr__(self):
+    return '{}(single_sample={})'.format(
+        type(self).__name__,
+        self._single_sample,
+    )
+
 
 class Normal(Distribution):
   __slots__ = ()
@@ -215,3 +242,17 @@ def __call__(self, initial_value=None, current_value=None, random_state=None):
         self._retain * self._value +
         random_state.normal(loc=0.0, scale=self._scale))
     return self._value
+
+  def __eq__(self, other):
+    # __eq__ shouldn't be used for this one, because it's stateful.
+    return id(self) == id(other)
+
+  def __repr__(self):
+    # include id(self), to make sure that two instances with the same parameters
+    # don't appear equal in logs.
+    return '{}(id={}, scale={}, retain={}, value={})'.format(
+        type(self).__name__,
+        id(self),
+        self._scale,
+        self._retain,
+        self._value)

dm_control/composer/variation/distributions_test.py

@@ -13,8 +13,6 @@
 # limitations under the License.
 # ============================================================================
 
-"""Tests for distributions."""
-
 from absl.testing import absltest
 from absl.testing import parameterized
 from dm_control.composer.variation import distributions
@@ -43,6 +41,10 @@ def testUniform(self):
           variation(random_state=self._variation_random_state),
           self._np_random_state.uniform(lower, upper))
 
+    self.assertEqual(variation, distributions.Uniform(low=lower, high=upper))
+    self.assertNotEqual(variation, distributions.Uniform(low=upper, high=upper))
+    self.assertIn('[2, 3, 4]', repr(variation))
+
   def testUniformChoice(self):
     choices = ['apple', 'banana', 'cherry']
     variation = distributions.UniformChoice(choices)
@@ -51,6 +53,8 @@ def testUniformChoice(self):
           variation(random_state=self._variation_random_state),
           self._np_random_state.choice(choices))
 
+    self.assertIn('banana', repr(variation))
+
   def testUniformPointOnSphere(self):
     variation = distributions.UniformPointOnSphere()
     samples = []
@@ -60,8 +64,11 @@ def testUniformPointOnSphere(self):
       np.testing.assert_approx_equal(np.linalg.norm(sample), 1.0)
       samples.append(sample)
     # Make sure that none of the samples are the same.
-    self.assertLen(
-        set(np.reshape(np.asarray(samples), -1)), 3 * NUM_ITERATIONS)
+    self.assertLen(set(np.reshape(np.asarray(samples), -1)), 3 * NUM_ITERATIONS)
+    self.assertEqual(variation, distributions.UniformPointOnSphere())
+    self.assertNotEqual(
+        variation, distributions.UniformPointOnSphere(single_sample=True)
+    )
 
   def testNormal(self):
     loc, scale = 1, 2
@@ -70,6 +77,14 @@ def testNormal(self):
       self.assertEqual(
           variation(random_state=self._variation_random_state),
           self._np_random_state.normal(loc, scale))
+    self.assertEqual(variation, distributions.Normal(loc=loc, scale=scale))
+    self.assertNotEqual(
+        variation, distributions.Normal(loc=loc*2, scale=scale)
+    )
+    self.assertEqual(
+        "Normal(args=(), kwargs={'loc': 1, 'scale': 2}, single_sample=False)",
+        repr(variation),
+    )
 
   def testExponential(self):
     scale = 3
@@ -78,6 +93,14 @@ def testExponential(self):
       self.assertEqual(
           variation(random_state=self._variation_random_state),
           self._np_random_state.exponential(scale))
+    self.assertEqual(variation, distributions.Exponential(scale=scale))
+    self.assertNotEqual(
+        variation, distributions.Exponential(scale=scale*2)
+    )
+    self.assertEqual(
+        "Exponential(args=(), kwargs={'scale': 3}, single_sample=False)",
+        repr(variation),
+    )
 
   def testPoisson(self):
     lam = 4
@@ -86,6 +109,14 @@ def testPoisson(self):
       self.assertEqual(
           variation(random_state=self._variation_random_state),
           self._np_random_state.poisson(lam))
+    self.assertEqual(variation, distributions.Poisson(lam=lam))
+    self.assertNotEqual(
+        variation, distributions.Poisson(lam=lam*2)
+    )
+    self.assertEqual(
+        "Poisson(args=(), kwargs={'lam': 4}, single_sample=False)",
+        repr(variation),
+    )
 
   @parameterized.parameters(0, 10)
   def testBiasedRandomWalk(self, timescale):

dm_control/composer/variation/math.py

@@ -26,10 +26,10 @@
 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.
+  Subclasses need to implement `_callable`, 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):
@@ -54,6 +54,21 @@ def __call__(self, initial_value=None, current_value=None, random_state=None):
   def _callable(self):
     pass
 
+  def __eq__(self, other):
+    if not isinstance(other, type(self)):
+      return False
+    return (
+        self._args == other._args
+        and self._kwargs == other._kwargs
+    )
+
+  def __repr__(self):
+    return '{}(args={}, kwargs={})'.format(
+        type(self).__name__,
+        self._args,
+        self._kwargs,
+    )
+
 
 class Log(MathOp):
 

dm_control/composer/variation/noises.py

@@ -39,6 +39,19 @@ def __call__(self, initial_value=None, current_value=None, random_state=None):
         variation_values.evaluate(self._variation, initial_value, current_value,
                                   random_state))
 
+  def __eq__(self, other):
+    if not isinstance(other, Additive):
+      return False
+    return (
+        self._variation == other._variation
+        and self._cumulative == other._cumulative
+    )
+
+  def __repr__(self):
+    return (
+        f"Additive(variation={self._variation}, cumulative={self._cumulative})"
+    )
+
 
 class Multiplicative(base.Variation):
   """A variation that multiplies to an existing value.
@@ -58,3 +71,17 @@ def __call__(self, initial_value=None, current_value=None, random_state=None):
     return base_value * (
         variation_values.evaluate(self._variation, initial_value, current_value,
                                   random_state))
+
+  def __eq__(self, other):
+    if not isinstance(other, Multiplicative):
+      return False
+    return (
+        self._variation == other._variation
+        and self._cumulative == other._cumulative
+    )
+
+  def __repr__(self):
+    return (
+        f"Multiplicative(variation={self._variation}, "
+        f"cumulative={self._cumulative})"
+    )