dm_env
dm_env copied to clipboard
Add TreeSpec to support dataclass-based nested specs (Issue #13)
This PR adds a minimal TreeSpec implementation to dm_env to support structured action and observation specifications using Python dataclasses. This directly addresses Issue #13 ("Tree Spec").
✔ What is added
- A new
TreeSpecclass in specs.py. - A recursive
_generate_tree_value()helper to generate valid test values. - Support for nested structures:
- dataclasses
- dicts
- tuples
- lists
- existing dm_env specs (Array, BoundedArray, DiscreteArray, etc.)
- A new test:
test_tree_spec_generate_value.
✔ Why this is useful
Many JAX-based RL environments define structured I/O using dataclasses. While dm_env already supports tuple/dict specs, dataclass-based specs were not recognized, causing type checking and clarity issues. TreeSpec provides an explicit and type-safe way to represent these structures.
✔ Compatibility
This feature is fully backward-compatible and does not modify any existing spec behavior. All existing tests pass (166 passed).