Fixes Gymnasium spaces issues due to Hydra/OmegaConf limitations (#1306)

# Description

Fixed issues with defining Gymnasium spaces in Direct workflows due to
Hydra/OmegaConf limitations with non-primitive types (see
#1264 (reply in thread))

```
omegaconf.errors.UnsupportedValueType: Value 'XXXXX' is not a supported primitive type
```


## Type of change

<!-- As you go through the list, delete the ones that are not
applicable. -->

- Bug fix (non-breaking change which fixes an issue)

## Checklist

- [x] I have run the [`pre-commit` checks](https://pre-commit.com/) with
`./isaaclab.sh --format`
- [ ] I have made corresponding changes to the documentation
- [x] My changes generate no new warnings
- [x] I have added tests that prove my fix is effective or that my
feature works
- [x] I have updated the changelog and the corresponding version in the
extension's `config/extension.toml` file
- [x] I have added my name to the `CONTRIBUTORS.md` or my name already
exists there

<!--
As you go through the checklist above, you can mark something as done by
putting an x character in it

For example,
- [x] I have done this task
- [ ] I have not done this task
-->

Author: Toni-SM

source/extensions/omni.isaac.lab/config/extension.toml

@@ -1,7 +1,7 @@
 [package]
 
 # Note: Semantic Versioning is used: https://semver.org/
-version = "0.27.4"
+version = "0.27.5"
 
 # Description
 title = "Isaac Lab framework for Robot Learning"

source/extensions/omni.isaac.lab/docs/CHANGELOG.rst

@@ -1,6 +1,15 @@
 Changelog
 ---------
 
+0.27.5 (2024-10-25)
+~~~~~~~~~~~~~~~~~~~
+
+Added
+^^^^^
+
+* Added utilities for serializing/deserializing Gymnasium spaces.
+
+
 0.27.4 (2024-10-18)
 ~~~~~~~~~~~~~~~~~~~
 

source/extensions/omni.isaac.lab/omni/isaac/lab/envs/utils/spaces.py

@@ -4,6 +4,7 @@
 # SPDX-License-Identifier: BSD-3-Clause
 
 import gymnasium as gym
+import json
 import numpy as np
 import torch
 from typing import Any
@@ -90,3 +91,131 @@ def tensorize(s, x):
 
     sample = (gym.vector.utils.batch_space(space, batch_size) if batch_size > 0 else space).sample()
     return tensorize(space, sample)
+
+
+def serialize_space(space: SpaceType) -> str:
+    """Serialize a space specification as JSON.
+
+    Args:
+        space: Space specification.
+
+    Returns:
+        Serialized JSON representation.
+    """
+    # Gymnasium spaces
+    if isinstance(space, gym.spaces.Discrete):
+        return json.dumps({"type": "gymnasium", "space": "Discrete", "n": int(space.n)})
+    elif isinstance(space, gym.spaces.Box):
+        return json.dumps({
+            "type": "gymnasium",
+            "space": "Box",
+            "low": space.low.tolist(),
+            "high": space.high.tolist(),
+            "shape": space.shape,
+        })
+    elif isinstance(space, gym.spaces.MultiDiscrete):
+        return json.dumps({"type": "gymnasium", "space": "MultiDiscrete", "nvec": space.nvec.tolist()})
+    elif isinstance(space, gym.spaces.Tuple):
+        return json.dumps({"type": "gymnasium", "space": "Tuple", "spaces": tuple(map(serialize_space, space.spaces))})
+    elif isinstance(space, gym.spaces.Dict):
+        return json.dumps(
+            {"type": "gymnasium", "space": "Dict", "spaces": {k: serialize_space(v) for k, v in space.spaces.items()}}
+        )
+    # Python data types
+    # Box
+    elif isinstance(space, int) or (isinstance(space, list) and all(isinstance(x, int) for x in space)):
+        return json.dumps({"type": "python", "space": "Box", "value": space})
+    # Discrete
+    elif isinstance(space, set) and len(space) == 1:
+        return json.dumps({"type": "python", "space": "Discrete", "value": next(iter(space))})
+    # MultiDiscrete
+    elif isinstance(space, list) and all(isinstance(x, set) and len(x) == 1 for x in space):
+        return json.dumps({"type": "python", "space": "MultiDiscrete", "value": [next(iter(x)) for x in space]})
+    # composite spaces
+    # Tuple
+    elif isinstance(space, tuple):
+        return json.dumps({"type": "python", "space": "Tuple", "value": [serialize_space(x) for x in space]})
+    # Dict
+    elif isinstance(space, dict):
+        return json.dumps(
+            {"type": "python", "space": "Dict", "value": {k: serialize_space(v) for k, v in space.items()}}
+        )
+    raise ValueError(f"Unsupported space ({space})")
+
+
+def deserialize_space(string: str) -> gym.spaces.Space:
+    """Deserialize a space specification encoded as JSON.
+
+    Args:
+        string: Serialized JSON representation.
+
+    Returns:
+        Space specification.
+    """
+    obj = json.loads(string)
+    # Gymnasium spaces
+    if obj["type"] == "gymnasium":
+        if obj["space"] == "Discrete":
+            return gym.spaces.Discrete(n=obj["n"])
+        elif obj["space"] == "Box":
+            return gym.spaces.Box(low=np.array(obj["low"]), high=np.array(obj["high"]), shape=obj["shape"])
+        elif obj["space"] == "MultiDiscrete":
+            return gym.spaces.MultiDiscrete(nvec=np.array(obj["nvec"]))
+        elif obj["space"] == "Tuple":
+            return gym.spaces.Tuple(spaces=tuple(map(deserialize_space, obj["spaces"])))
+        elif obj["space"] == "Dict":
+            return gym.spaces.Dict(spaces={k: deserialize_space(v) for k, v in obj["spaces"].items()})
+        else:
+            raise ValueError(f"Unsupported space ({obj['spaces']})")
+    # Python data types
+    elif obj["type"] == "python":
+        if obj["space"] == "Discrete":
+            return {obj["value"]}
+        elif obj["space"] == "Box":
+            return obj["value"]
+        elif obj["space"] == "MultiDiscrete":
+            return [{x} for x in obj["value"]]
+        elif obj["space"] == "Tuple":
+            return tuple(map(deserialize_space, obj["value"]))
+        elif obj["space"] == "Dict":
+            return {k: deserialize_space(v) for k, v in obj["value"].items()}
+        else:
+            raise ValueError(f"Unsupported space ({obj['spaces']})")
+    else:
+        raise ValueError(f"Unsupported type ({obj['type']})")
+
+
+def replace_env_cfg_spaces_with_strings(env_cfg: object) -> object:
+    """Replace spaces objects with their serialized JSON representations in an environment config.
+
+    Args:
+        env_cfg: Environment config instance.
+
+    Returns:
+        Environment config instance with spaces replaced if any.
+    """
+    for attr in ["observation_space", "action_space", "state_space"]:
+        if hasattr(env_cfg, attr):
+            setattr(env_cfg, attr, serialize_space(getattr(env_cfg, attr)))
+    for attr in ["observation_spaces", "action_spaces"]:
+        if hasattr(env_cfg, attr):
+            setattr(env_cfg, attr, {k: serialize_space(v) for k, v in getattr(env_cfg, attr).items()})
+    return env_cfg
+
+
+def replace_strings_with_env_cfg_spaces(env_cfg: object) -> object:
+    """Replace spaces objects with their serialized JSON representations in an environment config.
+
+    Args:
+        env_cfg: Environment config instance.
+
+    Returns:
+        Environment config instance with spaces replaced if any.
+    """
+    for attr in ["observation_space", "action_space", "state_space"]:
+        if hasattr(env_cfg, attr):
+            setattr(env_cfg, attr, deserialize_space(getattr(env_cfg, attr)))
+    for attr in ["observation_spaces", "action_spaces"]:
+        if hasattr(env_cfg, attr):
+            setattr(env_cfg, attr, {k: deserialize_space(v) for k, v in getattr(env_cfg, attr).items()})
+    return env_cfg

source/extensions/omni.isaac.lab/test/envs/test_spaces_utils.py

@@ -26,7 +26,7 @@
 import unittest
 from gymnasium.spaces import Box, Dict, Discrete, MultiDiscrete, Tuple
 
-from omni.isaac.lab.envs.utils.spaces import sample_space, spec_to_gym_space
+from omni.isaac.lab.envs.utils.spaces import deserialize_space, sample_space, serialize_space, spec_to_gym_space
 
 
 class TestSpacesUtils(unittest.TestCase):
@@ -104,6 +104,59 @@ def test_sample_space(self):
         self.assertIsInstance(sample, dict)
         self._check_tensorized(sample, batch_size=5)
 
+    def test_space_serialization_deserialization(self):
+        # fundamental spaces
+        # Box
+        space = 1
+        output = deserialize_space(serialize_space(space))
+        self.assertEqual(space, output)
+        space = [1, 2, 3, 4, 5]
+        output = deserialize_space(serialize_space(space))
+        self.assertEqual(space, output)
+        space = Box(low=-1.0, high=1.0, shape=(1, 2))
+        output = deserialize_space(serialize_space(space))
+        self.assertIsInstance(output, Box)
+        self.assertTrue((space.low == output.low).all())
+        self.assertTrue((space.high == output.high).all())
+        self.assertEqual(space.shape, output.shape)
+        # Discrete
+        space = {2}
+        output = deserialize_space(serialize_space(space))
+        self.assertEqual(space, output)
+        space = Discrete(2)
+        output = deserialize_space(serialize_space(space))
+        self.assertIsInstance(output, Discrete)
+        self.assertEqual(space.n, output.n)
+        # MultiDiscrete
+        space = [{1}, {2}, {3}]
+        output = deserialize_space(serialize_space(space))
+        self.assertEqual(space, output)
+        space = MultiDiscrete(np.array([1, 2, 3]))
+        output = deserialize_space(serialize_space(space))
+        self.assertIsInstance(output, MultiDiscrete)
+        self.assertTrue((space.nvec == output.nvec).all())
+        # composite spaces
+        # Tuple
+        space = ([1, 2, 3, 4, 5], {2}, [{1}, {2}, {3}])
+        output = deserialize_space(serialize_space(space))
+        self.assertEqual(space, output)
+        space = Tuple((Box(-1, 1, shape=(1,)), Discrete(2)))
+        output = deserialize_space(serialize_space(space))
+        self.assertIsInstance(output, Tuple)
+        self.assertEqual(len(output), 2)
+        self.assertIsInstance(output[0], Box)
+        self.assertIsInstance(output[1], Discrete)
+        # Dict
+        space = {"box": [1, 2, 3, 4, 5], "discrete": {2}, "multi_discrete": [{1}, {2}, {3}]}
+        output = deserialize_space(serialize_space(space))
+        self.assertEqual(space, output)
+        space = Dict({"box": Box(-1, 1, shape=(1,)), "discrete": Discrete(2)})
+        output = deserialize_space(serialize_space(space))
+        self.assertIsInstance(output, Dict)
+        self.assertEqual(len(output), 2)
+        self.assertIsInstance(output["box"], Box)
+        self.assertIsInstance(output["discrete"], Discrete)
+
     """
     Helper functions.
     """

source/extensions/omni.isaac.lab_tasks/config/extension.toml

@@ -1,7 +1,7 @@
 [package]
 
 # Note: Semantic Versioning is used: https://semver.org/
-version = "0.10.9"
+version = "0.10.10"
 
 # Description
 title = "Isaac Lab Environments"

source/extensions/omni.isaac.lab_tasks/docs/CHANGELOG.rst

@@ -1,6 +1,15 @@
 Changelog
 ---------
 
+0.10.10 (2024-10-25)
+~~~~~~~~~~~~~~~~~~~~
+
+Fixed
+^^^^^
+
+* Fixed issues with defining Gymnasium spaces in Direct workflows due to Hydra/OmegaConf limitations with non-primitive types.
+
+
 0.10.9 (2024-10-22)
 ~~~~~~~~~~~~~~~~~~~
 

source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/utils/hydra.py

@@ -17,6 +17,7 @@
     raise ImportError("Hydra is not installed. Please install it by running 'pip install hydra-core'.")
 
 from omni.isaac.lab.envs import DirectRLEnvCfg, ManagerBasedRLEnvCfg
+from omni.isaac.lab.envs.utils.spaces import replace_env_cfg_spaces_with_strings, replace_strings_with_env_cfg_spaces
 from omni.isaac.lab.utils import replace_slices_with_strings, replace_strings_with_slices
 
 from omni.isaac.lab_tasks.utils.parse_cfg import load_cfg_from_registry
@@ -40,6 +41,9 @@ def register_task_to_hydra(
     # load the configurations
     env_cfg = load_cfg_from_registry(task_name, "env_cfg_entry_point")
     agent_cfg = load_cfg_from_registry(task_name, agent_cfg_entry_point)
+    # replace gymnasium spaces with strings because OmegaConf does not support them.
+    # this must be done before converting the env configs to dictionary to avoid internal reinterpretations
+    replace_env_cfg_spaces_with_strings(env_cfg)
     # convert the configs to dictionary
     env_cfg_dict = env_cfg.to_dict()
     if isinstance(agent_cfg, dict):
@@ -83,6 +87,10 @@ def hydra_main(hydra_env_cfg: DictConfig, env_cfg=env_cfg, agent_cfg=agent_cfg):
                 hydra_env_cfg = replace_strings_with_slices(hydra_env_cfg)
                 # update the configs with the Hydra command line arguments
                 env_cfg.from_dict(hydra_env_cfg["env"])
+                # replace strings that represent gymnasium spaces because OmegaConf does not support them.
+                # this must be done after converting the env configs from dictionary to avoid internal reinterpretations
+                replace_strings_with_env_cfg_spaces(env_cfg)
+                # get agent configs
                 if isinstance(agent_cfg, dict):
                     agent_cfg = hydra_env_cfg["agent"]
                 else: