Gymnasium integration

import gymnasium
import pufferlib.emulation

# Create a Gymnasium environment
env = gymnasium.make('CartPole-v1')

# Wrap it with PufferLib
env = pufferlib.emulation.GymnasiumPufferEnv(env=env)

# Use the vectorized interface
obs, info = env.reset()
actions = env.action_space.sample()
obs, reward, terminal, truncation, info = env.step(actions)

print(obs.shape)      # (1, obs_size) - batch dimension added
print(reward.shape)   # (1,) - single reward per environment

GymnasiumPufferEnv(
    env=None,              # Gymnasium environment instance
    env_creator=None,      # Or callable that creates environment
    env_args=[],           # Args for env_creator
    env_kwargs={},         # Kwargs for env_creator  
    buf=None,              # Pre-allocated buffers (optional)
    seed=0                 # Random seed
)

import gymnasium
import pufferlib.emulation

class MyEnv(gymnasium.Env):
    def __init__(self):
        self.observation_space = gymnasium.spaces.Box(0, 1, (4,))
        # Complex action space
        self.action_space = gymnasium.spaces.Tuple((
            gymnasium.spaces.Discrete(4),
            gymnasium.spaces.Discrete(2),
        ))

env = MyEnv()
puffer_env = pufferlib.emulation.GymnasiumPufferEnv(env=env)

# Action space converted to MultiDiscrete
print(puffer_env.action_space)
# MultiDiscrete([4 2])
print(puffer_env.is_atn_emulated)
# True

# Sample from emulated action space
action = puffer_env.action_space.sample()  # [2, 1]

# Wrapper unpacks to original format: (2, 1)
obs, reward, terminal, truncation, info = puffer_env.step(action)

import gymnasium as gym
import pufferlib
import pufferlib.emulation
from ale_py import AtariEnv

def make(name='breakout', obs_type='grayscale', frameskip=4,
         full_action_space=False, framestack=1, buf=None):
    # Create base Atari environment
    env = AtariEnv(
        name, 
        obs_type=obs_type,
        frameskip=frameskip,
        full_action_space=full_action_space,
        render_mode='rgb_array'
    )
    
    # Fast downscaling wrapper
    env = pufferlib.ResizeObservation(env, downscale=2)
    
    # Optional frame stacking
    if framestack > 1:
        env = gym.wrappers.FrameStack(env, framestack)
    
    # Track episode statistics
    env = pufferlib.EpisodeStats(env)
    
    # Convert to PufferEnv interface
    env = pufferlib.emulation.GymnasiumPufferEnv(env=env, buf=buf)
    return env

from pufferlib.environments import atari

# Create Breakout environment
env = atari.make('breakout', framestack=4)

obs, info = env.reset()
for _ in range(100):
    action = env.action_space.sample()
    obs, reward, terminal, truncation, info = env.step(action)
    
    if terminal or truncation:
        print(f"Episode return: {info['episode_return']}")
        print(f"Episode length: {info['episode_length']}")
        obs, info = env.reset()

import gym
import gymnasium
import shimmy
import pufferlib
import pufferlib.emulation
import procgen

def make(name='bigfish', num_envs=1, num_levels=0, 
         start_level=0, distribution_mode='easy', buf=None):
    # Procgen provides vectorized envs by default
    envs = procgen.ProcgenEnv(
        env_name=name,
        num_envs=num_envs,
        num_levels=num_levels,
        start_level=start_level,
        distribution_mode=distribution_mode,
    )
    
    # Extract RGB observations
    envs = gym.wrappers.TransformObservation(envs, lambda obs: obs["rgb"])
    envs.single_action_space = envs.action_space
    envs.single_observation_space = envs.observation_space["rgb"]
    
    # Apply preprocessing
    envs = gym.wrappers.RecordEpisodeStatistics(envs)
    envs = gym.wrappers.NormalizeReward(envs)
    
    # Convert Gym to Gymnasium API
    envs = shimmy.GymV21CompatibilityV0(env=envs)
    envs = pufferlib.EpisodeStats(envs)
    
    # Convert to PufferEnv
    return pufferlib.emulation.GymnasiumPufferEnv(env=envs, buf=buf)

import gymnasium
import pufferlib
import pufferlib.emulation

class MiniGridWrapper:
    """Remove mission string from observations"""
    def __init__(self, env):
        self.env = env
        # Filter out 'mission' key from observation space
        self.observation_space = gymnasium.spaces.Dict({
            k: v for k, v in env.observation_space.items() 
            if k != 'mission'
        })
        self.action_space = env.action_space
    
    def reset(self, seed=None, options=None):
        obs, info = self.env.reset(seed=seed)
        del obs['mission']
        return obs, info
    
    def step(self, action):
        obs, reward, done, truncated, info = self.env.step(action)
        del obs['mission']
        return obs, reward, done, truncated, info

def make(name='MiniGrid-LavaGapS7-v0', buf=None):
    env = gymnasium.make(name, render_mode='rgb_array')
    env = MiniGridWrapper(env)  # Remove mission
    env = pufferlib.EpisodeStats(env)
    return pufferlib.emulation.GymnasiumPufferEnv(env=env, buf=buf)

from pufferlib.environments import mujoco

env = mujoco.make('Ant-v4')
obs, info = env.reset()

for _ in range(1000):
    # Sample continuous actions
    action = env.action_space.sample()
    obs, reward, terminal, truncation, info = env.step(action)
    
    if terminal or truncation:
        obs, info = env.reset()

import pufferlib
import pufferlib.emulation
from old_package import OldGymEnv  # Returns (obs, reward, done, info)

old_env = OldGymEnv()

# Convert to Gymnasium API
env = pufferlib.GymToGymnasium(old_env)

# Now wrap with PufferLib
env = pufferlib.emulation.GymnasiumPufferEnv(env=env)