Reinforcement Learning

The following is a simple example of training a reinforcement learning agent on the Lwmr environment using Stable Baselines3’s PPO implementation.

from dataclasses import dataclass

import gymnasium as gym
import lwmr  # noqa: F401  # registers "lwmr/Lwmr-v0"
import tyro
from stable_baselines3 import PPO
from stable_baselines3.common.env_util import make_vec_env
from stable_baselines3.common.vec_env import SubprocVecEnv


@dataclass
class Args:
    """Example using Lwmr environment."""

    quiet: bool = True
    seed: int = 47
    device: str = "cpu"

    n_envs: int = 8
    n_steps: int = 512
    total_timesteps: int = 1_000_000

    model_path: str = "ppo_lwmr"
    eval_steps: int = 100


# SubprocVecEnv requires the script to use __main__
if __name__ == "__main__":
    args = tyro.cli(Args)

    verbose = 0 if args.quiet else 1

    env_kwargs = {"quiet": args.quiet, "render_mode": "none"}
    env = make_vec_env("lwmr/Lwmr-v0", n_envs=args.n_envs, vec_env_cls=SubprocVecEnv, env_kwargs=env_kwargs)
    model = PPO("MlpPolicy", env, n_steps=args.n_steps, verbose=verbose, device=args.device)
    model.learn(total_timesteps=args.total_timesteps, progress_bar=True)
    model.save(args.model_path)

    env = gym.make("lwmr/Lwmr-v0", quiet=True, render_mode="viser")
    model = PPO.load(args.model_path, device=args.device)

    obs = env.reset()
    obs = obs[0]
    for _ in range(args.eval_steps):
        action, _ = model.predict(obs, deterministic=True)
        obs, reward, dones, info, blah = env.step(action)
        env.render()

    env.close()

!python ./basic_ppo-sb3.py --help

usage: ./basic_ppo-sb3.py [-h] [OPTIONS]



Example using Lwmr environment.



╭─ options ──────────────────────────────────────────────╮

│ -h, --help             show this help message and exit │

│ --quiet, --no-quiet    (default: True)                 │

│ --seed INT             (default: 47)                   │

│ --device STR           (default: cpu)                  │

│ --n-envs INT           (default: 8)                    │

│ --n-steps INT          (default: 512)                  │

│ --total-timesteps INT  (default: 1000000)              │

│ --model-path STR       (default: ppo_lwmr)             │

│ --eval-steps INT       (default: 100)                  │

╰────────────────────────────────────────────────────────╯

!python ./basic_ppo-sb3.py --quiet --total-timesteps 10000

2026-10-09 11:46:47.825 | WARNING  | lwmr.envs.plane:__init__:68 - CUDA device specified but not available, falling back to CPU

   ⋮ (7 identical messages omitted)


   0% ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 0/10,000  [ 0:00:00 < -:--:-- , ? it/s ]

   ⋮ (151 progress updates omitted)

 100% ━━━━━━━━━━━━━━━━━━━━━━━━━━ 12,288/10,000  [ 0:00:12 < 0:00:00 , 782 it/s ]


2026-10-09 11:47:12.926 | WARNING  | lwmr.envs.plane:__init__:68 - CUDA device specified but not available, falling back to CPU


  0%|                                                   | 0/100 [00:00<?, ?it/s]

   ⋮ (7 progress updates omitted)

100%|████████████████████████████████████████| 100/100 [00:00<00:00, 129.93it/s]

(viser) Server stopped

import gymnasium as gym
import lwmr  # noqa: F401
from stable_baselines3 import PPO
from tqdm.auto import trange

model_path = "ppo_lwmr.zip"
device = "cpu"
eval_steps = 200

env = gym.make("lwmr/Lwmr-v0", quiet=True, render_mode="viser")
model = PPO.load(model_path, device=device)

obs = env.reset()
obs = obs[0]
for _ in trange(eval_steps):
    action, _ = model.predict(obs, deterministic=True)
    obs, reward, dones, info, blah = env.step(action)
    env.render()
env.close()

2026-10-09 11:47:16.656 | WARNING  | lwmr.envs.plane:__init__:68 - CUDA device specified but not available, falling back to CPU

  0%|          | 0/200 [00:00<?, ?it/s]
   ⋮ (13 progress updates omitted)
100%|██████████| 200/200 [00:01<00:00, 127.28it/s]

(viser) Server stopped

# !python -m http.server 8047 1>/dev/null 2>&1
# http://localhost:8047/?playbackPath=./recordings/lwmr_plane.viser