diff --git a/gymnasium/vector/utils/space_utils.py b/gymnasium/vector/utils/space_utils.py index 8d38f5704..ae317a05a 100644 --- a/gymnasium/vector/utils/space_utils.py +++ b/gymnasium/vector/utils/space_utils.py @@ -308,6 +308,7 @@ def _concatenate_dict( @concatenate.register(Text) @concatenate.register(Sequence) @concatenate.register(Space) +@concatenate.register(OneOf) def _concatenate_custom(space: Space, items: Iterable, out: None) -> tuple[Any, ...]: return tuple(items) @@ -413,6 +414,11 @@ def _create_empty_array_sequence( return tuple(tuple() for _ in range(n)) +@create_empty_array.register(OneOf) +def _create_empty_array_oneof(space: OneOf, n: int = 1, fn=np.zeros): + return tuple(tuple() for _ in range(n)) + + @create_empty_array.register(Space) def _create_empty_array_custom(space, n=1, fn=np.zeros): return None