Tree Spec

[joeryjoery]

Hi, I was wondering if there were plans to by default support tree like structures in dm_env.specs.

This would be extremely relevant to jax-based environments, or generic environments utilizing deepmind/tree, to specify any arbitrary shape of agent-environment IO.

[joeryjoery]

Forgot to clarify. I'm aware that tuple, dicts, etc. are already supported as valid specs. The reason why I'm suggesting an explicit Tree spec is for my type-checker to OK when I return an explicit dataclass as a spec.

For example:

@dataclass
class MyAction
    a: jax.Array
    b: jax.Array

def action_spec(self) -> specs.Spec:
    return MyAction(a=specs.Array(...), b=specs.Array(...))  # mismatching types, specs.Array != jax.Array

I would find it confusing and more complicated to read if I had to add a | specs.Array type for all my dataclass type-hints.