$ show Farama-Foundation/SuperSuit
SuperSuit is a Python library that provides a collection of lightweight, composable 'microwrappers' for preprocessing reinforcement learning environments. It supports both single-agent Gymnasium environments and multi-agent PettingZoo environments (AEC and Parallel APIs). The wrappers cover common transformations such as color reduction, frame stacking, resizing, dtype conversion, reward clipping, action scaling, and NaN handling. Additionally, it offers vectorization utilities to run multiple environment instances in parallel, including synchronous, asynchronous, and multiprocessing-based vector environments. The library is designed to be used as a preprocessing pipeline before feeding environments into RL algorithms. It is maintained primarily for compatibility with newer versions of PettingZoo and Gymnasium, with the stated intention of eventually merging relevant functionality into PettingZoo's own wrappers.
The LLM identified 4 files worth highlighting for reuse:
Implements a `WrapperChooser` class that automatically selects the correct wrapper implementation based on the environment type (gymnasium.Env, pettingzoo.AECEnv, or pettingzoo.ParallelEnv). This allows a single wrapper function to work across all supported RL environment APIs without the user needing to specify the type. The pattern is self-contained and can be reused in any multi-API RL library to dispatch to appropriate wrapper classes.
Provides pure functions (`stack_obs_space`, `stack_init`, `stack_obs`) for stacking observations over time, supporting both Box and Discrete observation spaces with configurable stack dimension. The logic handles 1D, 2D, and 3D observations and uses numpy tiling for Box spaces and modular arithmetic for Discrete spaces. These functions are cleanly separated from the wrapper classes and can be extracted for use in any RL project requiring frame stacking.
Contains functions to homogenize action and observation spaces across agents in multi-agent environments by padding Box spaces to the maximum dimensions and expanding Discrete spaces to the maximum n. Also provides dehomogenization for actions and homogenization for observations. This is essential for vectorizing heterogeneous multi-agent environments and is implemented as a reusable utility independent of the wrapper infrastructure.
Implements `ProcConcatVec`, a high-performance multiprocessing vector environment that uses shared memory (via `multiprocessing.Array`) and pipes for inter-process communication. It conforms to Gymnasium's `VectorEnv` API and supports reset, step, render, and env_is_wrapped. The implementation handles graceful shutdown, error propagation, and efficient observation/action sharing. This is a non-trivial example of multiprocessing for RL that could be adapted for other vector environment implementations.