Add TreeSpec to support dataclass-based nested specs (Issue #13)

dm_env/__init__.py

@@ -27,3 +27,4 @@
 termination = _environment.termination
 transition = _environment.transition
 truncation = _environment.truncation
+from dm_env import specs

dm_env/specs.py

@@ -256,8 +256,12 @@ def validate(self, value):
     return value
 
   def generate_value(self):
-    return (np.ones(shape=self.shape, dtype=self.dtype) *
-            self.dtype.type(self.minimum))
+    """Generate a random value within [minimum, maximum] that matches this spec."""
+    return np.random.uniform(
+        low=self.minimum,
+        high=self.maximum,
+        size=self.shape
+    ).astype(self.dtype)
 
   def __reduce__(self):
     return BoundedArray, (self._shape, self._dtype, self._minimum,
@@ -402,3 +406,46 @@ def __repr__(self):
 
   def __reduce__(self):
     return type(self), (self.shape, self.string_type, self.name)
+
+from dataclasses import is_dataclass, fields
+
+class TreeSpec:
+  """A container for nested spec-like structures, including dataclasses."""
+
+  def __init__(self, structure):
+    self.structure = structure
+
+  def generate_value(self):
+    return _generate_tree_value(self.structure)
+
+  def __repr__(self):
+    return f"TreeSpec({self.structure!r})"
+
+
+def _generate_tree_value(structure):
+  """Recursively generate test values for nested spec structures."""
+
+  # Case 1 — Array, BoundedArray, DiscreteArray, StringArray
+  if isinstance(structure, Array):
+    return structure.generate_value()
+
+  # Case 2 — dataclass
+  if is_dataclass(structure):
+    return type(structure)(**{
+        f.name: _generate_tree_value(getattr(structure, f.name))
+        for f in fields(structure)
+    })
+
+  # Case 3 — dict
+  if isinstance(structure, dict):
+    return {k: _generate_tree_value(v) for k, v in structure.items()}
+
+  # Case 4 — tuple
+  if isinstance(structure, tuple):
+    return tuple(_generate_tree_value(v) for v in structure)
+
+  # Case 5 — list
+  if isinstance(structure, list):
+    return [_generate_tree_value(v) for v in structure]
+
+  raise TypeError(f"Unsupported element in TreeSpec: {type(structure)}")

dm_env/tests/test_bounded_array.py

@@ -0,0 +1,16 @@
+import numpy as np
+from dm_env import specs
+
+def test_generate_value_within_bounds():
+    spec = specs.BoundedArray(
+        shape=(2, 2),
+        dtype=np.float32,
+        minimum=0.0,
+        maximum=5.0
+    )
+    value = spec.generate_value()
+
+    assert value.shape == (2, 2)
+    assert value.dtype == np.float32
+    assert np.all(value >= 0.0)
+    assert np.all(value <= 5.0)

dm_env/tests/test_tree_spec.py

@@ -0,0 +1,24 @@
+from dataclasses import dataclass
+from dm_env import specs
+
+@dataclass
+class MyAction:
+    a: specs.Array
+    b: specs.Array
+
+def test_tree_spec_generate_value():
+    spec = specs.TreeSpec(
+        MyAction(
+            a=specs.Array(shape=(2,), dtype=float),
+            b=specs.Array(shape=(3,), dtype=float)
+        )
+    )
+
+    value = spec.generate_value()
+
+    # type check
+    assert isinstance(value, MyAction)
+
+    # shape checks
+    assert value.a.shape == (2,)
+    assert value.b.shape == (3,)