Available Ocean environments

from pufferlib.ocean.cartpole import Cartpole

# Create environment
env = Cartpole(
    num_envs=4096,
    continuous=True,  # or False for discrete
    cart_mass=1.0,
    pole_mass=0.1,
    gravity=9.8,
)

obs, info = env.reset()

for _ in range(1000):
    actions = policy(obs)  # Your policy here
    obs, rewards, terms, truncs, info = env.step(actions)

env.close()

from pufferlib.ocean.snake import Snake

# Create 16 environments with 256 snakes each
env = Snake(
    num_envs=16,
    num_snakes=256,  # Per environment
    width=640,
    height=360,
    vision=5,
    reward_food=0.1,
    reward_death=-1.0,
)

obs, info = env.reset()

# obs.shape = (4096, 11, 11)  # 16 * 256 agents
# Each agent sees local 11x11 grid around itself

for _ in range(1000):
    actions = policy(obs)
    obs, rewards, terms, truncs, info = env.step(actions)

env.close()

from pufferlib.ocean.battle import Battle

env = Battle(
    num_envs=8,
    num_agents=1024,  # Per environment
    num_armies=4,
    width=1920,
    height=1080,
    render_mode='human',
)

obs, info = env.reset()
# obs.shape = (8192, obs_dim)  # 8 * 1024 agents

for _ in range(1000):
    actions = policy(obs)  # Continuous actions
    obs, rewards, terms, truncs, info = env.step(actions)

env.render()  # Visualize first environment
env.close()

import pufferlib.ocean

# Get environment class
env_class = pufferlib.ocean.env_creator('puffer_snake')

# Create environment
env = env_class(num_envs=16, num_snakes=256)