Implement __getitem__ for Composer variations.

saran-t · copybara-github · commit b2be68d56dee · 2020-04-22T14:48:01.000-07:00
PiperOrigin-RevId: 307904417
Change-Id: I44f9f2ec70878d95a5f6341dac8dc4d03ffe6d0f
diff --git a/dm_control/composer/variation/base.py b/dm_control/composer/variation/base.py
@@ -23,6 +23,7 @@
 import operator
 
 from dm_control.composer.variation import variation_values
+import numpy as np
 import six
 
 
@@ -82,6 +83,9 @@ def __pow__(self, other):
   def __rpow__(self, other):
     return _BinaryOperation(operator.pow, other, self)
 
+  def __getitem__(self, index):
+    return _GetItemOperation(self, index)
+
 
 class _BinaryOperation(Variation):
   """Represents the result of applying a binary operator to two Variations."""
@@ -97,3 +101,15 @@ def __call__(self, initial_value=None, current_value=None, random_state=None):
     second_value = variation_values.evaluate(
         self._second, initial_value, current_value, random_state)
     return self._op(first_value, second_value)
+
+
+class _GetItemOperation(Variation):
+
+  def __init__(self, variation, index):
+    self._variation = variation
+    self._index = 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)
+    return np.asarray(value)[self._index]
diff --git a/dm_control/composer/variation/variation_test.py b/dm_control/composer/variation/variation_test.py
@@ -24,11 +24,13 @@
 from absl.testing import parameterized
 from dm_control.composer import variation
 from dm_control.composer.variation import deterministic
+import numpy as np
 
 
 class VariationTest(parameterized.TestCase):
 
   def setUp(self):
+    super(VariationTest, self).setUp()
     self.value_1 = 3
     self.variation_1 = deterministic.Constant(self.value_1)
     self.value_2 = 5
@@ -47,6 +49,11 @@ def test_operator(self, name):
         variation.evaluate(func(self.variation_1, self.variation_2)),
         func(self.value_1, self.value_2))
 
+  def test_getitem(self):
+    value = deterministic.Constant(np.array([4, 5, 6, 7, 8]))
+    np.testing.assert_array_equal(
+        variation.evaluate(value[[3, 1]]),
+        [7, 5])
 
 if __name__ == '__main__':
   absltest.main()
diff --git a/setup.py b/setup.py
@@ -166,7 +166,7 @@ def find_data_files(package_dir, patterns):
 
 setup(
     name='dm_control',
-    version='0.0.307814181',
+    version='0.0.307904417',
     description='Continuous control environments and MuJoCo Python bindings.',
     author='DeepMind',
     license='Apache License, Version 2.0',
