[Proposal] AsyncVectorEnv with Graph observation spaces

[ramondalmau]

Proposal

Add support for environments with Graph observation spaces in AsyncVectorEnv in Gymnasium. Currently, AsyncVectorEnv assumes observations can be stacked in a typical array-like format. However, environments that return graph-based observations (e.g., adjacency matrices, node/edge features) are incompatible with the default vectorized operations. This proposal seeks to introduce native support for vectorizing such environments, allowing the parallelization of environments that utilize Graph-like spaces for their observations.

Motivation

The current limitation of AsyncVectorEnv is that it can only handle environments with simple observation spaces. For graphs (e.g., adjacency matrices or node/edge data), the default handling of observation stacking fails, leading to errors and the inability to utilize vectorized environments.

Pitch

Modify AsyncVectorEnv to natively handle observation spaces that are graph-structured, by allowing users to define custom stacking or returning observations as lists of graphs instead of arrays. This could be done via:

• Adding a specific check for Graph-type spaces in the observation stacking method.
• Providing an option for environments with complex observations to return observations as lists, or by defining a custom merging mechanism for such spaces.

Alternatives

No response

Additional context

Many modern reinforcement learning tasks involve graph-structured data, particularly in domains such as networked systems, biology, and operations research. A robust way to handle graph-based environments in a vectorized setup would greatly enhance Gymnasium's flexibility and application scope.

Additionally, this proposal would align with Gymnasium's mission to support a wide range of environment types, particularly as more graph-based tasks emerge in reinforcement learning.

Checklist

• [X] I have checked that there is no similar issue in the repo

[pseudo-rnd-thoughts]

With AsyncVectorEnv(..., shared_memory=True) then it is required that the observations have a fixed shape. There are two options,

1. disable shared memory - turning off shared memory will allow dynamic observation shapes
2. fixed shape graph instances - we could add a feature such that observations have a fixed number of nodes (and thus edges)