How to code wrapper for envs returning a dictionary of tensors of the observation?

[ErcBunny]

For example my env step returns (obs_dict, reward, terminated, truncated, info) as type Tuple[Dict[str, Tensor], Tensor, Tensor, Tensor, Dict[str, Tensor]], and I see the wrapper base that the step should return Tuple[Tensor, Tensor, Tensor, Tensor, Any]. How should I code the wrapper in this case?

Thanks for reviewing and helping.

[Toni-SM]

Hi @ErcBunny

You can take a look at the Isaac Lab env wrapper that also return a dictionary as observation:

skrl/skrl/envs/wrappers/torch/isaaclab_envs.py

Lines 31 to 32 in 636936f

	self._obs_dict, reward, terminated, truncated, info = self._env.step(actions)
	return self._obs_dict["policy"], reward.view(-1, 1), terminated.view(-1, 1), truncated.view(-1, 1), info

[ErcBunny]

Thank you for directing me to the isaac lab wrapper, but is _obs_dict["policy"] just a single tensor? I would like to have multiple tensors as observation.

[Toni-SM]

Hi @ErcBunny

Since memory only support tensors (not collection of tensors: dict, tuples, etc.) it is necessary to convert the observation dict (assuming it is a gymnasium.space.Dict space) as flattened tensor in the wrapper. Then, in the model, you can map the flattened tensor to the gymnasium space again using the Model.tensor_to_space method.

For converting the space to flattened tensor see the gymnasium wrapper _observation_to_tensor method used in

skrl/skrl/envs/wrappers/torch/gymnasium_envs.py

Line 129 in 636936f

observation = self._observation_to_tensor(observation)

Also, see the DeepMind Manipulation: Reach Site Vision dm_manipulation_stack_sac.py example that show how to use the tensor_to_space method.