env-ssl-wrapper

One line turns any simulator's environment — mujoco warp, isaac sim, pybullet, gymnasium, pufferlib, dm_control, robosuite — into the same torch-native interface.

Install

pip install env-ssl-wrapper

Standardize

import torch
from env_ssl_wrapper import StandardizeEnvWrapper

# one master wrapper to standardize any simulator environment

env = StandardizeEnvWrapper(any_env)

obs, info = env.reset() # torch.float32, batched

while not env.all_done:
    actions = torch.randint(0, 2, (1,))
    obs, reward, terminated, truncated, info = env.step(actions)

Works identically for every simulator.

Or compose individual wrappers piecemeal:

from env_ssl_wrapper import compose_env

env = compose_env(
    any_env,
    ('tensor', dict(device = 'cpu')),
    'done_tracker',
)

Research Wrappers

Memory Trace

A cheap memory that keeps track of an exponential moving average (EMA) of a state, or subset (Eberhard et al., 2025).

import gymnasium as gym
from env_ssl_wrapper import StandardizeEnvWrapper, MemoryTraceWrapper

env = StandardizeEnvWrapper(gym.make('LunarLander-v3'))
env = MemoryTraceWrapper(env, lambdas = (0.9, 0.99))

obs, info = env.reset()
# obs['obs']        -> (1, 8)
# obs['trace_0.9']  -> (1, 8)
# obs['trace_0.99'] -> (1, 8)

Or on a subset of dictionary observations:

env = MemoryTraceWrapper(env, lambdas = (0.9, 0.99), keys = 'proprio')
# obs['proprio_trace_0.9'] -> (1, 2)

Run the PPO benchmark on POMDP LunarLander:

uv run test_memory_trace.py

Wrappers

Pass wrappers as strings (default config) or (name, dict) tuples (custom config), in any order.

Every env emits the same contract: obs torch.float32, rewards torch.float32, terminated/truncated torch.bool. env.seed(n) works on every sim.

Terminated envs are uniformly padded (zeros / False obs; rewards zeroed only after the terminating step, so the terminal transition's reward is never lost), and info['final_observation'] — the true terminal obs, frozen per env and re-emitted while the env stays done — is always present once any env has terminated, with info['_final_observation'] masking which envs it applies to. env.is_done always reflects the per-env done mask.

Conforming Janky Simulators

Say you are handed a legacy or custom simulator with completely idiosyncratic signatures — non-standard method names, costs instead of rewards, inverted booleans, and custom rendering:

class JankyRoboticsSim:
    def boot(self):
        return {'sensors': [0.1, -0.5, 1.2]}

    def actuate(self, torque):
        # returns sensor_dict, cost, is_alive
        return {'sensors': [0.2, -0.4, 1.1]}, 0.05, True

    def grab_pixels(self, w, h):
        return np.zeros((h, w, 3), dtype = np.uint8)

Wrap all that idiosyncrasy into an adapter in a few lines:

import numpy as np
from env_ssl_wrapper import BaseEnvAdapter, register_adapter, compose_env

class JankySimAdapter(BaseEnvAdapter):
    @classmethod
    def matches(cls, env):
        return isinstance(env, JankyRoboticsSim)

    def reset(self, **kwargs):
        obs = self.env.boot()
        return {'sensors': np.asarray(obs['sensors'])}, {}

    def step(self, action):
        data, cost, is_alive = self.env.actuate(action)
        return {'sensors': np.asarray(data['sensors'])}, -cost, not is_alive, False, {}

    def render(self, height, width, camera = None):
        return self.env.grab_pixels(width, height)

register_adapter(JankySimAdapter)

Now it behaves like every first-class simulator in the ecosystem:

env = compose_env(
    JankyRoboticsSim(),
    ('image', dict(image_size = (64, 64))),
    'auto_batch',
    'tensor',
    'done_tracker'
)

obs, info = env.reset()
# obs['image']   -> torch.Size([1, 3, 64, 64])
# obs['sensors'] -> torch.Size([1, 3])

obs, reward, terminated, truncated, info = env.step(torch.randn(1, 1))

Multiprocessing

Parallelize any single environment or factory into an autoresetting vector env:

from env_ssl_wrapper import MultiprocessingVecEnv, compose_env

with MultiprocessingVecEnv('CartPole-v1', num_envs = 8) as env:
    env = compose_env(env, 'tensor', 'done_tracker')
    obs, info = env.reset()

Tests

uv sync --extra test
uv run pytest

Citations

@article{eberhard2025partially,
    title   = {Partially Observable Reinforcement Learning with Memory Traces},
    author  = {Onno Eberhard and Michael Muehlebach and Claire Vernade},
    journal = {arXiv preprint arXiv:2503.15200},
    year    = {2025}
}