[Feature] Dict specs in vmas (pytorch#2415)

Co-authored-by: Vincent Moens <vmoens@meta.com>
ai4co · Sep 4, 2024 · df4fa78 · df4fa78
1 parent 49d7f74
commit df4fa78
Showing 1 changed file with 10 additions and 0 deletions.
diff --git a/torchrl/envs/libs/vmas.py b/torchrl/envs/libs/vmas.py
@@ -94,6 +94,16 @@ def _vmas_to_torchrl_spec_transform(
                 device=device,
             )
         )
+    elif isinstance(spec, gym_spaces.Dict):
+        spec_out = {}
+        for key in spec.keys():
+            spec_out[key] = _vmas_to_torchrl_spec_transform(
+                spec[key],
+                device=device,
+                categorical_action_encoding=categorical_action_encoding,
+            )
+        # the batch-size must be set later
+        return Composite(spec_out, device=device)
     else:
         raise NotImplementedError(
             f"spec of type {type(spec).__name__} is currently unaccounted for vmas"