From df4fa7808e81f4b95ee2c22a4bb768370a669048 Mon Sep 17 00:00:00 2001 From: Matteo Bettini <55539777+matteobettini@users.noreply.github.com> Date: Wed, 4 Sep 2024 17:44:55 +0200 Subject: [PATCH] [Feature] Dict specs in vmas (#2415) Co-authored-by: Vincent Moens --- torchrl/envs/libs/vmas.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/torchrl/envs/libs/vmas.py b/torchrl/envs/libs/vmas.py index 5811580826d..22f9835303b 100644 --- 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"