dm_env icon indicating copy to clipboard operation
dm_env copied to clipboard

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

Open Pranav-0440 opened this issue 1 month ago • 0 comments

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 TreeSpec class 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).

Pranav-0440 avatar Nov 22 '25 19:11 Pranav-0440