Environment setup + train loop (#5)

Mujoco environment setup (vectorized on GPU) + training loop + simulate script

.gitignore

@@ -1,3 +1,7 @@
+# Model files
+artifacts/*
+runs/*
+
 # Python-generated files
 __pycache__/
 *.py[oc]
@@ -515,4 +519,3 @@ Icon
 Network Trash Folder
 Temporary Items
 .apdisk
-

README.md

@@ -0,0 +1,17 @@
+# Brittle Star
+
+## Usage
+
+### UV
+
+To set up the UV module, you can run the following command:
+
+```bash
+uv sync --frozen
+```
+
+example command:
+
+```bash
+uv run src/train.py --model_name my_model --epochs 50 --batch_size 32
+```

docs/api/environment.md

@@ -0,0 +1,24 @@
+# Brittle star environment
+
+## Creation
+The environment package contains a factory class `BrittleStarEnvFactory` 
+that creates instances of the environment/morphologies/... It uses the
+configuration classes defined in `env_config.py` to create the instances.
+
+## Configuration
+The data classes in `env_config` have default values as stated in the tutorials.
+* MorphologyConfig: configuration for the morphology of the brittle star. Contains
+number of arms, number of segments per arm, and control mode.
+* ArenaConfig: configuration for the arena. Sets the size of the arena, whether to 
+set the ground floor to sand, attach a target and sizes of the walls.
+* EnvConfig: configuration for the environment. These set shared settings
+such as camera locations, simulation time and the task.
+
+## Backend and Task enums
+The Backend enum specifies either an MJC or MJX backend.
+* MJC: runs on CPU
+* MJX: uses jax on the gpu
+
+The Task enum specifies which task to use. 2 items are present:
+* DIRECTED_LOCOMOTION: move to a target location
+* LIGHT_ESCAPE: situation where the robot must move to a darker location

docs/api/train_simulate.md

@@ -0,0 +1,30 @@
+# Training and Simulation for Brittle Star Models
+
+## Training a model
+
+To train a model, you can use the `train.py` script. This script allows to pass some parameters to customize the training process:
+
+- `--out`: The output path where the trained model will be saved.
+- `--model_type`: The type of model to train (e.g., `random`, ...)
+- `--task`: The task to train on (e.g., `directed_locomotion`, ...)
+- `--seed`: The random seed for reproducibility.
+- `--epochs`: The number of epochs to train for.
+
+This will then train the specified model on the specified task for the given number of epochs and save the trained model to the specified output path.
+
+```bash
+python train.py --out artifacts/my_model --model-type random --task directed_locomotion --seed 0 --epochs 50 
+```
+
+## Simulating a model
+
+In order to simulate and view the behavior of a trained model, you can use the `simulate.py` script. This script allows you to specify the path to a trained model and will launch a simulation using that model. This script has the following parameters:
+
+- `--model`: The path to the trained model artifact to simulate.
+- `--model-type`: The type of model to simulate (e.g., `random`, ...)
+- `--task`: The task to simulate (e.g., `directed_locomotion`, ...)
+- `--seed`: The random seed for reproducibility.
+
+```bash
+python simulate.py --model artifacts/my_model --model-type random --task directed_locomotion --seed 0
+```

experiments/simulate.py

@@ -0,0 +1,100 @@
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+
+from brittle_star_project import (
+    ArenaConfig,
+    Backend,
+    BrittleStarEnv,
+    BrittleStarEnvFactory,
+    EnvConfig,
+    MorphologyConfig,
+    Task,
+)
+from brittle_star_project.rl import RLModel  # imports concrete models via rl.__init__
+from brittle_star_project.rl.base import get_rl_model_registry
+from brittle_star_project.renderer import SimulationConfig, simulate_policy
+
+MODEL_BY_NAME = get_rl_model_registry()
+MODEL_OPTIONS = sorted(MODEL_BY_NAME)
+
+
+def parse_args() -> argparse.Namespace:
+    p = argparse.ArgumentParser(description="Simulate a trained policy in the MuJoCo viewer.")
+    p.add_argument(
+        "--model",
+        type=str,
+        default=None,
+        help="Path to a saved model artifact. If omitted, a model is created from --model-type.",
+    )
+    p.add_argument(
+        "--model-type",
+        choices=MODEL_OPTIONS,
+        default="random",
+        help="Which model class to instantiate when --model is omitted.",
+    )
+    p.add_argument(
+        "--task",
+        choices=[t.value for t in Task],
+        default=Task.DIRECTED_LOCOMOTION.value,
+    )
+    p.add_argument("--seed", type=int, default=None)
+    return p.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+
+    # ======= ENVIRONMENT SETUP =======
+
+    backend = Backend.MJC
+    task = Task(args.task)
+
+    morphology_cfg = MorphologyConfig()
+    arena_cfg = ArenaConfig(attach_target=(task == Task.DIRECTED_LOCOMOTION))
+    env_cfg = EnvConfig(task=task)
+
+    factory = BrittleStarEnvFactory()
+    raw_env = factory.create_environment(backend, morphology_cfg, arena_cfg, env_cfg)
+    env = BrittleStarEnv(raw_env, backend=backend, config=env_cfg)
+
+    seed_for_env = int(args.seed) if args.seed is not None else 0
+    state = env.reset(seed=seed_for_env)
+
+    # ======= MODEL SETUP =======
+
+    # Extract the number of actuators (nu) from the environment's model, so we can pass it to the
+    # policy/model.
+    nu = int(state.mj_model.nu)
+
+    if args.model is not None:
+        model_path = Path(args.model)
+        policy = RLModel.load(model_path)
+        if hasattr(policy, "nu"):
+            policy.nu = nu
+    else:
+        model_cls = MODEL_BY_NAME[str(args.model_type)]
+        policy = model_cls(seed=seed_for_env)
+        if hasattr(policy, "nu"):
+            policy.nu = nu
+
+    # If the policy/model has a `seed` attribute, use the provided seed (or default) to reset it.
+    default_seed = int(getattr(policy, "seed", seed_for_env))
+    if args.seed is not None and hasattr(policy, "reset"):
+        policy.reset(int(args.seed))
+
+    # ======= SIMULATION =======
+
+    rollout_cfg = SimulationConfig(
+        realtime=True,
+        seed=int(args.seed) if args.seed is not None else default_seed,
+    )
+
+    simulate_policy(policy, rollout_cfg, state)
+
+    env.close()
+
+
+if __name__ == "__main__":
+    main()

pyproject.toml

@@ -3,16 +3,21 @@ name = "2026sel3-project"
 version = "0.1.0"
 description = "Add your description here"
 readme = "README.md"
-requires-python = ">=3.12, <3.13"
+requires-python = ">= 3.12, < 3.13"
 dependencies = [
     "biorobot==0.4.2",
+    "cleanrl>=0.4.8",
     "evosax==0.2.0",
+    "flax>=0.12.2",
+    "gymnasium>=1.2.3",
     "ipykernel==7.2.0",
     "jax==0.9.0.1",
     "matplotlib==3.10.8",
     "mediapy==1.2.6",
+    "optax>=0.2.6",
     "pyopengl>=3.1.10",
     "pyopengl-accelerate>=3.1.10",
+    "tyro>=1.0.10",
     "wandb==0.24.2",
 ]
 

ruff.toml

@@ -1,3 +1,5 @@
+line-length = 100
+
 [lint]
 extend-select = [
   # "PLC0103",  # invalid-name

src/brittle_star_project/__init__.py

@@ -0,0 +1,14 @@
+from .environment.env_types import Backend, Task
+from .environment.env_config import ArenaConfig, EnvConfig, MorphologyConfig
+from .environment.factory import BrittleStarEnvFactory
+from .environment.env_wrapper import BrittleStarEnv
+
+__all__ = [
+    "ArenaConfig",
+    "Backend",
+    "BrittleStarEnv",
+    "BrittleStarEnvFactory",
+    "EnvConfig",
+    "MorphologyConfig",
+    "Task",
+]

src/brittle_star_project/dataclasses/EpisodeStatistics.py

@@ -0,0 +1,10 @@
+import flax.struct
+import jax.numpy as jnp
+
+
+@flax.struct.dataclass
+class EpisodeStatistics:
+    episode_returns: jnp.array
+    episode_lengths: jnp.array
+    returned_episode_returns: jnp.array
+    returned_episode_lengths: jnp.array

src/brittle_star_project/dataclasses/PPOArgs.py

@@ -0,0 +1,103 @@
+from dataclasses import dataclass
+
+
+@dataclass
+class PPOArgs:
+    """
+    source: https://github.com/vwxyzjn/cleanrl/blob/master/cleanrl/ppo_atari_envpool_xla_jax_scan.py
+    """
+
+    # the name of this experiment
+    exp_name: str = "brittle_star_ppo"
+
+    # seed of the experiment
+    seed: int = 1
+
+    # if toggled, `torch.backends.cudnn.deterministic=False`
+    torch_deterministic: bool = True
+
+    # if toggled, cuda will be enabled by default
+    cuda: bool = True
+
+    # if toggled, this experiment will be tracked with Weights and Biases
+    track: bool = False
+
+    # the wandb's project name
+    wandb_project_name: str = "PPO-Modularity"
+
+    # the entity (team) of wandb's project
+    wandb_entity: str | None = None
+
+    # whether to capture videos of the agent performances (check out `videos` folder)
+    capture_video: bool = False
+
+    # whether to save model into the `runs/{run_name}` folder
+    save_model: bool = True
+
+    # whether to upload the saved model to huggingface
+    upload_model: bool = False
+
+    # the user or org name of the model repository from the Hugging Face Hub
+    hf_entity: str = ""
+
+    # ==== Algorithm specific dataclasses ====
+    # the id of the environment
+    env_id: str = ""  # todo
+
+    # total timesteps of the experiments
+    total_timesteps: int = 10000000
+
+    # the learning rate of the optimizer
+    learning_rate: float = 2.5e-4
+
+    # the number of parallel game environments
+    num_envs: int = 16
+
+    # the number of steps to run in each environment per policy rollout
+    num_steps: int = 128
+
+    # Toggle learning rate annealing for policy and value networks
+    anneal_lr: bool = True
+
+    # the discount factor gamma
+    gamma: float = 0.99
+
+    # the lambda for the general advantage estimation
+    gae_lambda: float = 0.95
+
+    # the number of mini-batches
+    num_minibatches: int = 4
+
+    # the K epochs to update the policy
+    update_epochs: int = 4
+
+    # Toggles advantages normalization
+    norm_adv: bool = True
+
+    # the surrogate clipping coefficient
+    clip_coef: float = 0.1
+
+    # Toggles whether or not to use a clipped loss for the value function, as per the paper.
+    clip_vloss: bool = True
+
+    # coefficient of the entropy
+    ent_coef: float = 0.01
+
+    # coefficient of the value function
+    vf_coef: float = 0.5
+
+    # the maximum norm for the gradient clipping
+    max_grad_norm: float = 0.5
+
+    # the target KL divergence threshold
+    target_kl: float | None = None
+
+    # ==== to be filled in runtime ====
+    # the batch size (computed in runtime)
+    batch_size: int = 0
+
+    # the mini-batch size (computed in runtime)
+    minibatch_size: int = 0
+
+    # the number of iterations (computed in runtime)
+    num_iterations: int = 0

src/brittle_star_project/dataclasses/__init__.py

@@ -0,0 +1,8 @@
+from .PPOArgs import PPOArgs
+from .EpisodeStatistics import EpisodeStatistics
+
+
+__all__ = [
+    "PPOArgs",
+    "EpisodeStatistics",
+]

src/brittle_star_project/environment/BrittleStarJaxEnvWrapper.py

@@ -0,0 +1,78 @@
+import jax
+import jax.numpy as jnp
+
+from brittle_star_project import (
+    EnvConfig,
+    BrittleStarEnvFactory,
+    MorphologyConfig,
+    ArenaConfig,
+    Backend,
+)
+
+
+class BrittleStarJaxEnvWrapper:
+    def __init__(
+        self,
+        morphology: MorphologyConfig,
+        arena: ArenaConfig,
+        env_config: EnvConfig,
+        num_envs: int,
+        backend: Backend = Backend.MJX,
+    ):
+        self._morphology = morphology
+        self._arena = arena
+        self._env_config = env_config
+        self._backend = backend
+        self._num_envs = num_envs
+        self._env = BrittleStarEnvFactory.create_environment(
+            self._backend, self._morphology, self._arena, self._env_config
+        )
+
+        self._vectorized_reset = jax.jit(jax.vmap(self._env.reset))
+        self._vectorized_step = jax.jit(jax.vmap(self._env.step))
+        self._vectorized_action_sample = jax.jit(jax.vmap(self._env.action_space.sample))
+
+        self._action_rng = None
+
+    @property
+    def backend(self):
+        return self._backend
+
+    @property
+    def raw(self):
+        return self._env
+
+    @property
+    def single_action_space(self):
+        return self._env.action_space
+
+    @property
+    def single_observation_space(self):
+        return self._env.observation_space
+
+    def reset(self, seed: int = 0):
+        self._action_rng, env_rng = jax.random.split(jax.random.PRNGKey(seed), 2)
+        env_rngs = jnp.array(jax.random.split(env_rng, self._num_envs))
+        return self._vectorized_reset(rng=env_rngs)
+
+    def sample_actions(self):
+        assert self._action_rng is not None, "Call reset() before sample_actions()"
+        self._action_rng, *sub_rngs = jnp.array(
+            jax.random.split(self._action_rng, self._num_envs + 1)
+        )
+        return self._vectorized_action_sample(rng=jnp.array(sub_rngs))
+
+    def step(self, state, action):
+        return self._vectorized_step(state=state, action=action)
+
+    def close(self):
+        self._env.close()
+
+    @staticmethod
+    def default(num_envs: int, backend: Backend = Backend.MJX) -> "BrittleStarJaxEnvWrapper":
+        morphology = MorphologyConfig()
+        arena = ArenaConfig()
+        env_config = EnvConfig()
+        return BrittleStarJaxEnvWrapper(
+            morphology, arena, env_config, num_envs=num_envs, backend=backend
+        )

src/brittle_star_project/environment/__init__.py

@@ -0,0 +1,15 @@
+from .env_config import ArenaConfig, EnvConfig, MorphologyConfig
+from .env_types import Backend, Task
+from .env_wrapper import BrittleStarEnv, StepResult
+from .factory import BrittleStarEnvFactory
+
+__all__ = [
+    "ArenaConfig",
+    "EnvConfig",
+    "MorphologyConfig",
+    "Backend",
+    "Task",
+    "BrittleStarEnv",
+    "StepResult",
+    "BrittleStarEnvFactory",
+]

src/brittle_star_project/environment/env_config.py

@@ -0,0 +1,53 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+
+from .env_types import Task
+
+
+@dataclass(frozen=True, slots=True)
+class MorphologyConfig:
+    num_arms: int = 5
+    num_segments_per_arm: int = 4
+    use_p_control: bool = True
+    use_torque_control: bool = False
+
+
+@dataclass(frozen=True, slots=True)
+class ArenaConfig:
+    size: tuple[float, float] = (10.0, 5.0)
+    sand_ground_color: bool = True
+    attach_target: bool = True
+    wall_height: float = 1.5
+    wall_thickness: float = 0.1
+
+
+@dataclass(frozen=True, slots=True)
+class EnvConfig:
+    """Shared environment settings.
+
+    Note: Some tasks have additional parameters (see fields below).
+    """
+
+    task: Task = Task.DIRECTED_LOCOMOTION
+
+    simulation_time: float = 5.0
+    num_physics_steps_per_control_step: int = 10
+    time_scale: int = 2
+
+    camera_ids: list[int] = field(default_factory=lambda: [0, 1])
+    # (height, width)
+    render_size: tuple[int, int] = (480, 640)
+
+    joint_randomization_noise_scale: float = 0.0
+
+    # Directed locomotion
+    target_distance: float = 3.0
+
+    # Light escape
+    # Per docs in upstream env config: integer factors of 200.
+    light_perlin_noise_scale: int = 0
+
+    @staticmethod
+    def from_json(path: str) -> EnvConfig:
+        pass

src/brittle_star_project/environment/env_types.py

@@ -0,0 +1,21 @@
+from __future__ import annotations
+
+from enum import Enum
+
+
+class Backend(str, Enum):
+    """Physics backend.
+
+    - MJC: MuJoCo C engine
+    - MJX: MuJoCo XLA (JAX) engine
+    """
+
+    MJC = "MJC"
+    MJX = "MJX"
+
+
+class Task(str, Enum):
+    """Which brittle-star task/environment to instantiate."""
+
+    DIRECTED_LOCOMOTION = "directed_locomotion"
+    LIGHT_ESCAPE = "light_escape"

src/brittle_star_project/environment/env_wrapper.py

@@ -0,0 +1,87 @@
+from __future__ import annotations
+
+import inspect
+from dataclasses import dataclass
+from typing import Any
+
+import numpy as np
+
+from .env_config import EnvConfig
+from .env_types import Backend
+
+
+@dataclass(slots=True)
+class StepResult:
+    state: Any
+    reward: float | None = None
+    terminated: bool | None = None
+    truncated: bool | None = None
+    info: dict[str, Any] | None = None
+
+
+class BrittleStarEnv:
+    """Thin wrapper around the underlying DualMuJoCoEnvironment.
+
+    Goal: hide backend-specific RNG setup and provide a stable place to plug in RL.
+    """
+
+    def __init__(self, env: Any, *, backend: Backend, config: EnvConfig) -> None:
+        self._env = env
+        self._backend = backend
+        self._config = config
+
+    @property
+    def raw(self) -> Any:
+        return self._env
+
+    @property
+    def backend(self) -> Backend:
+        return self._backend
+
+    @property
+    def config(self) -> EnvConfig:
+        return self._config
+
+    def make_rng(self, seed: int):
+        if self._backend == Backend.MJC:
+            return np.random.RandomState(seed)
+
+        import jax
+
+        return jax.random.PRNGKey(seed)
+
+    def reset(self, *, seed: int = 0):
+        rng = self.make_rng(seed)
+        state = self._env.reset(rng=rng)
+        return state
+
+    def render(self, *, state: Any):
+        return self._env.render(state=state)
+
+    def close(self) -> None:
+        self._env.close()
+
+    def step(self, *, state: Any, action: Any, rng: Any | None = None) -> StepResult:
+        """Best-effort step wrapper.
+
+        Different env libraries return different tuples; we normalize common cases.
+        """
+
+        if not hasattr(self._env, "step"):
+            raise AttributeError("Underlying env has no step() method")
+
+        step_fn = self._env.step
+        sig = inspect.signature(step_fn)
+        params = list(sig.parameters)
+
+        # Common patterns:
+        # - step(state, action)
+        # - step(state, action, rng)
+        # - step(state, action, key)
+        # We pass rng only if the callable accepts a 3rd arg.
+        if len(params) >= 3 and rng is not None:
+            out = step_fn(state, action, rng)
+        else:
+            out = step_fn(state, action)
+
+        return out

src/brittle_star_project/environment/factory.py

@@ -0,0 +1,106 @@
+from __future__ import annotations
+
+from dataclasses import asdict
+
+from moojoco.environment.dual import DualMuJoCoEnvironment
+
+from .env_config import ArenaConfig, EnvConfig, MorphologyConfig
+from .env_types import Backend, Task
+
+
+class BrittleStarEnvFactory:
+    """Creates brittle-star morphology, arena, and task environment instances."""
+
+    @staticmethod
+    def create_morphology(config: MorphologyConfig):
+        from biorobot.brittle_star.mjcf.morphology.morphology import (
+            MJCFBrittleStarMorphology,
+        )
+        from biorobot.brittle_star.mjcf.morphology.specification.default import (
+            default_brittle_star_morphology_specification,
+        )
+
+        spec = default_brittle_star_morphology_specification(
+            num_arms=config.num_arms,
+            num_segments_per_arm=config.num_segments_per_arm,
+            use_p_control=config.use_p_control,
+            use_torque_control=config.use_torque_control,
+        )
+        return MJCFBrittleStarMorphology(specification=spec)
+
+    @staticmethod
+    def create_arena(config: ArenaConfig):
+        from biorobot.brittle_star.mjcf.arena.aquarium import (
+            AquariumArenaConfiguration,
+            MJCFAquariumArena,
+        )
+
+        arena_config = AquariumArenaConfiguration(**asdict(config))
+        return MJCFAquariumArena(configuration=arena_config)
+
+    @staticmethod
+    def create_environment_configuration(config: EnvConfig):
+        # Import locally so the project can still be imported without these deps.
+        from biorobot.brittle_star.environment.directed_locomotion.shared import (
+            BrittleStarDirectedLocomotionEnvironmentConfiguration,
+        )
+        from biorobot.brittle_star.environment.light_escape.shared import (
+            BrittleStarLightEscapeEnvironmentConfiguration,
+        )
+
+        common = dict(
+            joint_randomization_noise_scale=config.joint_randomization_noise_scale,
+            render_mode="human",
+            simulation_time=config.simulation_time,
+            num_physics_steps_per_control_step=config.num_physics_steps_per_control_step,
+            time_scale=config.time_scale,
+            camera_ids=config.camera_ids,
+            render_size=config.render_size,
+        )
+
+        match config.task:
+            case Task.DIRECTED_LOCOMOTION:
+                return BrittleStarDirectedLocomotionEnvironmentConfiguration(
+                    target_distance=config.target_distance,
+                    **common,
+                )
+            case Task.LIGHT_ESCAPE:
+                return BrittleStarLightEscapeEnvironmentConfiguration(
+                    light_perlin_noise_scale=config.light_perlin_noise_scale,
+                    **common,
+                )
+            case _:
+                raise ValueError(f"Unsupported task: {config.task}")
+
+    @staticmethod
+    def create_environment(
+        backend: Backend,
+        morphology_config: MorphologyConfig,
+        arena_config: ArenaConfig,
+        env_config: EnvConfig,
+    ) -> DualMuJoCoEnvironment:
+        from biorobot.brittle_star.environment.directed_locomotion.dual import (
+            BrittleStarDirectedLocomotionEnvironment,
+        )
+        from biorobot.brittle_star.environment.light_escape.dual import (
+            BrittleStarLightEscapeEnvironment,
+        )
+
+        morphology = BrittleStarEnvFactory.create_morphology(morphology_config)
+        arena = BrittleStarEnvFactory.create_arena(arena_config)
+        env_configuration = BrittleStarEnvFactory.create_environment_configuration(env_config)
+
+        match env_config.task:
+            case Task.DIRECTED_LOCOMOTION:
+                env_class = BrittleStarDirectedLocomotionEnvironment
+            case Task.LIGHT_ESCAPE:
+                env_class = BrittleStarLightEscapeEnvironment
+            case _:
+                raise ValueError(f"Unsupported task: {env_config.task}")
+
+        return env_class.from_morphology_and_arena(
+            morphology=morphology,
+            arena=arena,
+            configuration=env_configuration,
+            backend=backend.value,
+        )

src/brittle_star_project/render/renderer.py

@@ -0,0 +1,75 @@
+from __future__ import annotations
+
+import time
+from dataclasses import dataclass
+from typing import Any, Protocol
+
+import numpy as np
+
+
+@dataclass
+class SimulationConfig:
+    realtime: bool = True
+    seed: int = 0
+
+
+class ControlPolicy(Protocol):
+    def act(self, *, obs: np.ndarray | None = None, t: float = 0.0) -> np.ndarray: ...
+
+
+def _default_observations(data: Any) -> np.ndarray:
+    qpos = np.asarray(data.qpos, dtype=np.float32).ravel()
+    qvel = np.asarray(data.qvel, dtype=np.float32).ravel()
+    return np.concatenate([qpos, qvel], axis=0)
+
+
+def simulate_policy(
+    policy: ControlPolicy,
+    config: SimulationConfig,
+    state: Any | None = None,
+) -> None:
+    """Open MuJoCo's native viewer and step using actions from a policy.
+
+    This path drives MuJoCo physics directly (mj_step) and uses the policy output
+    as `data.ctrl`.
+    """
+
+    import mujoco.viewer
+
+    model = state.mj_model
+    data = state.mj_data
+
+    start = time.time()
+    with mujoco.viewer.launch_passive(model, data) as viewer:
+        while viewer.is_running():
+            step_start = time.time()
+
+            t = time.time() - start
+
+            # Input vector for the policy
+            # TODO: custom input
+            obs = _default_observations(data)
+
+            # Policy action
+            ctrl = policy.act(obs=obs, t=t)
+
+            # Check if the policy output vector give an input for each actuator (nu)
+            # TODO: what if model trained on full morphology but we want to test on a damaged one?
+            # (nu mismatch)
+            if model.nu > 0:
+                ctrl = np.asarray(ctrl, dtype=np.float32).ravel()
+                if ctrl.shape != (model.nu,):
+                    raise ValueError(
+                        f"Policy returned ctrl shape {ctrl.shape}, expected ({model.nu},)"
+                    )
+                data.ctrl[:] = ctrl
+
+            # Step the simulation and update the viewer
+            mujoco.mj_step(model, data)
+            viewer.sync()
+
+            # If we're running in realtime mode, sleep to maintain real-time pacing.
+            if config.realtime:
+                remaining = model.opt.timestep - (time.time() - step_start)
+                if remaining > 0:
+                    time.sleep(remaining)

src/brittle_star_project/rl/DummyAgent.py

@@ -0,0 +1,71 @@
+from dataclasses import dataclass, fields
+
+import flax
+import flax.linen as nn
+import jax.numpy as jnp
+import jax.tree_util
+import numpy as np
+from flax.linen.initializers import constant, orthogonal
+
+
+class Network(nn.Module):
+    """
+    Dummy model only used for testing purposes
+
+    inspired by: https://github.com/vwxyzjn/cleanrl/blob/master/cleanrl/ppo_atari_envpool_xla_jax_scan.py
+    """
+
+    hidden_dim: int = 195
+
+    @nn.compact
+    def __call__(self, x):
+        x = nn.Dense(self.hidden_dim, kernel_init=orthogonal(np.sqrt(2)), bias_init=constant(0.0))(
+            x
+        )
+        x = nn.relu(x)
+        x = nn.Dense(self.hidden_dim, kernel_init=orthogonal(np.sqrt(2)), bias_init=constant(0.0))(
+            x
+        )
+        x = nn.relu(x)
+        return x
+
+
+class Critic(nn.Module):
+    @nn.compact
+    def __call__(self, x):
+        return nn.Dense(1, kernel_init=orthogonal(1), bias_init=constant(0.0))(x)
+
+
+class Actor(nn.Module):
+    action_dim: int
+
+    @nn.compact
+    def __call__(self, x):
+        mean = nn.Dense(self.action_dim, kernel_init=orthogonal(0.01), bias_init=constant(0.0))(x)
+        log_std = self.param("log_std", nn.initializers.zeros, (self.action_dim,))
+        return mean, log_std
+
+
+@jax.tree_util.register_dataclass
+@dataclass
+class AgentParams:
+    network_params: flax.core.FrozenDict
+    actor_params: flax.core.FrozenDict
+    critic_params: flax.core.FrozenDict
+
+
+@jax.tree_util.register_dataclass
+@dataclass
+class Storage:
+    obs: jnp.array
+    actions: jnp.array
+    logprobs: jnp.array
+    dones: jnp.array
+    values: jnp.array
+    advantages: jnp.array
+    returns: jnp.array
+    rewards: jnp.array
+
+    def replace(self, **kwargs) -> "Storage":
+        fs = fields(self)
+        return Storage(**{f.name: kwargs.get(f.name, getattr(self, f.name)) for f in fs})

src/brittle_star_project/rl/__init__.py

@@ -0,0 +1,25 @@
+from .DummyAgent import Network, Critic, Actor, AgentParams, Storage
+from .base import (
+    RLAlgorithm,
+    RLModel,
+    Transition,
+    create_model,
+    register_rl_model,
+    registered_model_types,
+)
+from .random_policy_model import RandomPolicyModel
+
+__all__ = [
+    "RLAlgorithm",
+    "RLModel",
+    "RandomPolicyModel",
+    "Transition",
+    "create_model",
+    "register_rl_model",
+    "registered_model_types",
+    "Network",
+    "Critic",
+    "Actor",
+    "AgentParams",
+    "Storage",
+]

src/brittle_star_project/rl/base.py

@@ -0,0 +1,162 @@
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+import json
+from pathlib import Path
+from typing import Any
+
+
+@dataclass(frozen=True, slots=True)
+class Transition:
+    """A minimal transition container for RL.
+
+    This is intentionally generic because the underlying env state type may be a
+    JAX pytree, a numpy struct, or something library-specific.
+    """
+
+    obs: Any
+    action: Any
+    reward: float
+    next_obs: Any
+    terminated: bool
+    truncated: bool
+    info: dict[str, Any] | None = None
+
+
+class RLAlgorithm(ABC):
+    """Insertable RL algorithm interface."""
+
+    @abstractmethod
+    def select_action(self, *, obs: Any, rng: Any | None = None) -> Any:
+        raise NotImplementedError
+
+    def observe(self, transition: Transition) -> None:
+        """Optional hook to store transitions."""
+
+    def update(self, *, rng: Any | None = None) -> dict[str, float]:
+        """Optional hook to run one training update."""
+
+        return {}
+
+    def save(self, path: str) -> None:
+        raise NotImplementedError("Save not implemented")
+
+    def load(self, path: str) -> None:
+        raise NotImplementedError("Load not implemented")
+
+
+_RL_MODEL_REGISTRY: dict[str, type["RLModel"]] = {}
+
+
+def registered_model_types() -> list[str]:
+    return sorted(_RL_MODEL_REGISTRY)
+
+
+def create_model(type_name: str, *, payload: dict[str, Any]) -> "RLModel":
+    model_cls = _RL_MODEL_REGISTRY.get(type_name)
+    if model_cls is None:
+        known = ", ".join(sorted(_RL_MODEL_REGISTRY)) or "<none>"
+        raise ValueError(f"Unknown RLModel type '{type_name}'. Known: {known}")
+    return model_cls.from_payload(payload)
+
+
+def get_rl_model_registry() -> dict[str, type["RLModel"]]:
+    """Return a copy of the current RLModel registry.
+
+    The registry is populated by importing concrete model modules that use the
+    `@register_rl_model(...)` decorator.
+    """
+
+    return dict(_RL_MODEL_REGISTRY)
+
+
+def register_rl_model(*type_names: str):
+    """Decorator to register an `RLModel` for generic loading.
+
+    Concrete model modules should apply this decorator, so `base.py` never needs
+    to import concrete models (avoids circular imports).
+    """
+
+    if not type_names:
+        raise TypeError("register_rl_model() requires at least one type name")
+
+    primary = type_names[0]
+
+    def _decorator(cls: type[RLModel]):
+        for name in type_names:
+            _RL_MODEL_REGISTRY[name] = cls
+        cls.type_name = primary
+        return cls
+
+    return _decorator
+
+
+class RLModel(ABC):
+    """Serializable policy/model interface.
+
+    This is the artifact that `train.py` writes and `simulate.py` loads.
+    """
+
+    # Overwritten by the `@register_rl_model(...)` decorator.
+    type_name: str = "RLModel"
+
+    def reset(self, seed: int | None = None) -> None:
+        """Optional hook for RNG/stateful models."""
+
+    @abstractmethod
+    def act(self, *, obs: Any | None = None, t: float = 0.0) -> Any:
+        raise NotImplementedError
+
+    def train(self, *, env: Any, num_epochs: int = 1) -> None:
+        """Optional training hook.
+
+        Many models won't learn; for those this can be a no-op.
+        """
+
+        _ = (env, num_epochs)
+
+    def to_payload(self) -> dict[str, Any]:
+        """Return JSON-serializable model parameters."""
+
+        return {}
+
+    @classmethod
+    def from_payload(cls, payload: dict[str, Any]) -> "RLModel":
+        """Reconstruct a model from `to_payload()` output."""
+
+        return cls(**payload)  # type: ignore[arg-type]
+
+    def save(self, path: str | Path) -> Path:
+        out = Path(path)
+        out.parent.mkdir(parents=True, exist_ok=True)
+        doc = {
+            "type": self.type_name,
+            "version": 1,
+            "payload": self.to_payload(),
+        }
+        out.write_text(json.dumps(doc, indent=2, sort_keys=True) + "\n")
+        return out
+
+    @classmethod
+    def load(cls, path: str | Path) -> "RLModel":
+        p = Path(path)
+        doc = json.loads(p.read_text())
+
+        type_name = doc.get("type")
+        if not isinstance(type_name, str):
+            raise ValueError("Model artifact missing string field 'type'")
+
+        model_cls = _RL_MODEL_REGISTRY.get(type_name)
+        if model_cls is None:
+            known = ", ".join(sorted(_RL_MODEL_REGISTRY)) or "<none>"
+            raise ValueError(f"Unknown RLModel type '{type_name}'. Known: {known}")
+
+        payload = doc.get("payload")
+        # Backward compatibility: older artifacts stored fields at top-level.
+        if payload is None:
+            payload = {k: v for k, v in doc.items() if k not in ("type", "version")}
+        if not isinstance(payload, dict):
+            raise ValueError("Model artifact field 'payload' must be an object")
+
+        return model_cls.from_payload(payload)

src/brittle_star_project/rl/random_policy_model.py

@@ -0,0 +1,52 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+
+import numpy as np
+
+from .base import RLModel, register_rl_model
+
+
+@register_rl_model("random")
+@dataclass(slots=True)
+class RandomPolicyModel(RLModel):
+    """A minimal, serializable policy model that outputs random controls.
+
+    This is intentionally *not* a learning algorithm yet. It exists so we can:
+    - produce a stable model artifact from `train.py`
+    - load that artifact in `simulate.py`
+    - drive the MuJoCo viewer with the model's actions
+    """
+
+    nu: int = 0
+    seed: int = 0
+    ctrl_noise_scale: float = 0.5
+
+    _rng: np.random.RandomState = field(init=False, repr=False)
+
+    def __post_init__(self) -> None:
+        self.reset(self.seed)
+
+    def reset(self, seed: int | None = None) -> None:
+        if seed is not None:
+            self.seed = int(seed)
+        self._rng = np.random.RandomState(self.seed)
+
+    def act(self, *, obs: np.ndarray | None = None, t: float = 0.0) -> np.ndarray:
+        if self.nu <= 0:
+            return np.zeros((0,), dtype=np.float32)
+        ctrl = self.ctrl_noise_scale * self._rng.randn(self.nu)
+        return ctrl.astype(np.float32)
+
+    def to_payload(self) -> dict[str, object]:
+        return {
+            "seed": int(self.seed),
+            "ctrl_noise_scale": float(self.ctrl_noise_scale),
+        }
+
+    @classmethod
+    def from_payload(cls, payload: dict[str, object]) -> RandomPolicyModel:
+        return cls(
+            seed=int(payload.get("seed", 0)),
+            ctrl_noise_scale=float(payload.get("ctrl_noise_scale", 0.5)),
+        )

src/train.py

@@ -0,0 +1,377 @@
+import random
+import time
+from dataclasses import asdict
+from functools import partial
+from typing import Callable
+
+import flax
+import jax
+import jax.numpy as jnp
+import numpy as np
+import optax
+import torch
+import tqdm
+import tyro
+from flax.training.train_state import TrainState
+from torch.utils.tensorboard import SummaryWriter
+
+from brittle_star_project.dataclasses import PPOArgs
+from brittle_star_project.dataclasses.EpisodeStatistics import EpisodeStatistics
+from brittle_star_project.environment.BrittleStarJaxEnvWrapper import BrittleStarJaxEnvWrapper
+from brittle_star_project.rl import Network, Actor, Critic, AgentParams, Storage
+
+
+def convert_obs_dict_to_array(obs_dict: dict) -> jnp.ndarray:
+    return jax.vmap(lambda o: jnp.concatenate([v.flatten() for v in o.values() if v.size > 0]))(
+        obs_dict
+    )
+
+
+def make_env(num_envs: int) -> Callable:
+    def thunk():
+        return BrittleStarJaxEnvWrapper.default(num_envs=num_envs)
+
+    return thunk
+
+
+def train(args: PPOArgs):
+    args.batch_size = args.num_envs * args.num_steps
+    args.minibatch_size = args.batch_size // args.num_minibatches
+    args.num_iterations = args.total_timesteps // args.batch_size
+    run_name = f"{args.exp_name}__seed_{args.seed}__{int(time.time())}"
+    print(f"running name: {run_name}")
+
+    if args.track:
+        import wandb
+
+        wandb.init(
+            project=args.wandb_project_name,
+            entity=args.wandb_entity,
+            sync_tensorboard=True,
+            config=vars(args),
+            name=run_name,
+            save_code=True,
+        )
+
+    writer = SummaryWriter(f"runs/{run_name}")
+    writer.add_text(
+        "hyperparameters",
+        "|param|value|\n|---|---|\n" + "\n".join(f"|{k}|{v}|" for k, v in vars(args).items()),
+    )
+
+    random.seed(args.seed)
+    np.random.seed(args.seed)
+    key = jax.random.PRNGKey(args.seed)
+    key, network_key, actor_key, critic_key = jax.random.split(key, 4)
+
+    torch.backends.cudnn.deterministic = args.torch_deterministic
+    device = torch.device("cuda" if torch.cuda.is_available() and args.cuda else "cpu")
+    print(f"Running on device: {device}")
+
+    print("Creating the environment...")
+    env = make_env(num_envs=args.num_envs)()
+
+    episode_stats = EpisodeStatistics(
+        episode_returns=jnp.zeros(args.num_envs, dtype=jnp.float32),
+        episode_lengths=jnp.zeros(args.num_envs, dtype=jnp.int32),
+        returned_episode_returns=jnp.zeros(args.num_envs, jnp.float32),
+        returned_episode_lengths=jnp.zeros(args.num_envs, dtype=jnp.int32),
+    )
+
+    def step_env_wrapped(episode_stats: EpisodeStatistics, env_state, action):
+        next_env_state = env.step(env_state, action)
+
+        # Extract per-environment signals from the state object
+        reward = next_env_state.reward  # (num_envs,)
+        terminated = next_env_state.terminated  # (num_envs,)
+        truncated = next_env_state.truncated  # (num_envs,)
+        done = terminated | truncated  # (num_envs,)
+
+        new_episode_return = episode_stats.episode_returns + reward
+        new_episode_length = episode_stats.episode_lengths + 1
+
+        episode_stats = episode_stats.replace(
+            episode_returns=new_episode_return * (1 - done),
+            episode_lengths=new_episode_length * (1 - done),
+            returned_episode_returns=jnp.where(
+                done, new_episode_return, episode_stats.returned_episode_returns
+            ),
+            returned_episode_lengths=jnp.where(
+                done, new_episode_length, episode_stats.returned_episode_lengths
+            ),
+        )
+        return (
+            episode_stats,
+            next_env_state,
+            (convert_obs_dict_to_array(next_env_state.observations), reward, done),
+        )
+
+    def linear_schedule(count):
+        frac = 1.0 - (count // (args.num_minibatches * args.update_epochs)) / args.num_iterations
+        return args.learning_rate * frac
+
+    print("Initializing the models...")
+    network = Network()
+    actor = Actor(action_dim=env.single_action_space.shape[0])  # continuous actions for MJX
+    critic = Critic()
+
+    sample_obs = jnp.concatenate(
+        [
+            v.flatten()
+            for v in env.single_observation_space.sample(rng=jax.random.PRNGKey(0)).values()
+            if v.size > 0
+        ]
+    )
+    network_params = network.init(network_key, sample_obs)
+    actor_params = actor.init(actor_key, network.apply(network_params, sample_obs))
+    critic_params = critic.init(critic_key, network.apply(network_params, sample_obs))
+
+    agent_state = TrainState.create(
+        apply_fn=None,
+        params=asdict(AgentParams(network_params, actor_params, critic_params)),
+        tx=optax.chain(
+            optax.clip_by_global_norm(args.max_grad_norm),
+            optax.inject_hyperparams(optax.adam)(
+                learning_rate=linear_schedule if args.anneal_lr else args.learning_rate, eps=1e-5
+            ),
+        ),
+    )
+
+    network.apply = jax.jit(network.apply)
+    actor.apply = jax.jit(actor.apply)
+    critic.apply = jax.jit(critic.apply)
+
+    @jax.jit
+    def get_action_and_value_noise(
+        agent_state: TrainState,
+        next_obs: jnp.ndarray,
+        key: jax.random.PRNGKey,
+    ):
+        hidden = network.apply(agent_state.params["network_params"], next_obs)
+        # Continuous actions: sample from a Gaussian parameterized by the actor
+        mean, log_std = actor.apply(agent_state.params["actor_params"], hidden)
+        key, subkey = jax.random.split(key)
+        noise = jax.random.normal(subkey, shape=mean.shape)
+        std = jnp.exp(log_std)
+        action = mean + noise * std
+        logprob = -0.5 * (((action - mean) / std) ** 2 + 2 * log_std + jnp.log(2 * jnp.pi)).sum(-1)
+        value = critic.apply(agent_state.params["critic_params"], hidden)
+        return action, logprob, value.squeeze(-1), key
+
+    @jax.jit
+    def get_action_and_value(
+        params: flax.core.FrozenDict,
+        x: jnp.ndarray,
+        action: np.ndarray,
+    ):
+        hidden = network.apply(params["network_params"], x)
+        mean, log_std = actor.apply(params["actor_params"], hidden)
+        std = jnp.exp(log_std)
+        logprob = -0.5 * (((action - mean) / std) ** 2 + 2 * log_std + jnp.log(2 * jnp.pi)).sum(-1)
+        entropy = (0.5 + 0.5 * jnp.log(2 * jnp.pi) + log_std).sum(-1)
+        value = critic.apply(params["critic_params"], hidden).squeeze(-1)
+        return logprob, entropy, value
+
+    @jax.jit
+    def compute_gae_once(carry, inp, gamma, gae_lambda):
+        advantages = carry
+        nextdone, nextvalues, curvalues, reward = inp
+        nextnonterminal = 1.0 - nextdone
+        delta = reward + gamma * nextvalues * nextnonterminal - curvalues
+        advantages = delta + gamma * gae_lambda * nextnonterminal * advantages
+        return advantages, advantages
+
+    @jax.jit
+    def compute_gae(agent_state, next_obs, next_done, storage):
+        next_value = critic.apply(
+            agent_state.params["critic_params"],
+            network.apply(agent_state.params["network_params"], next_obs),
+        ).squeeze(-1)
+
+        advantages = jnp.zeros((args.num_envs,))
+        dones = jnp.concatenate([storage.dones, next_done[None, :]], axis=0)
+        values = jnp.concatenate([storage.values, next_value[None, :]], axis=0)
+        _, advantages = jax.lax.scan(
+            partial(compute_gae_once, gamma=args.gamma, gae_lambda=args.gae_lambda),
+            advantages,
+            (dones[1:], values[1:], values[:-1], storage.rewards),
+            reverse=True,
+        )
+        return storage.replace(advantages=advantages, returns=advantages + storage.values)
+
+    def ppo_loss(params, x, a, logp, mb_advantages, mb_returns):
+        newlogprob, entropy, newvalue = get_action_and_value(params, x, a)
+        logratio = newlogprob - logp
+        ratio = jnp.exp(logratio)
+        approx_kl = ((ratio - 1) - logratio).mean()
+
+        if args.norm_adv:
+            mb_advantages = (mb_advantages - mb_advantages.mean()) / (mb_advantages.std() + 1e-8)
+
+        pg_loss1 = -mb_advantages * ratio
+        pg_loss2 = -mb_advantages * jnp.clip(ratio, 1 - args.clip_coef, 1 + args.clip_coef)
+        pg_loss = jnp.maximum(pg_loss1, pg_loss2).mean()
+        v_loss = 0.5 * ((newvalue - mb_returns) ** 2).mean()
+        entropy_loss = entropy.mean()
+        loss = pg_loss - args.ent_coef * entropy_loss + v_loss * args.vf_coef
+        return loss, (pg_loss, v_loss, entropy_loss, jax.lax.stop_gradient(approx_kl))
+
+    ppo_loss_grad_fn = jax.value_and_grad(ppo_loss, has_aux=True)
+
+    @jax.jit
+    def update_ppo(agent_state, storage, key):
+        def update_epoch(carry, _):
+            agent_state, key = carry
+            key, subkey = jax.random.split(key)
+
+            def flatten(x):
+                return x.reshape((-1,) + x.shape[2:])
+
+            def convert_data(x):
+                x = jax.random.permutation(subkey, x)
+                return jnp.reshape(x, (args.num_minibatches, -1) + x.shape[1:])
+
+            flatten_storage = jax.tree.map(flatten, storage)
+            shuffled_storage = jax.tree.map(convert_data, flatten_storage)
+
+            def update_minibatch(agent_state, minibatch):
+                (loss, (pg_loss, v_loss, entropy_loss, approx_kl)), grads = ppo_loss_grad_fn(
+                    agent_state.params,
+                    minibatch.obs,
+                    minibatch.actions,
+                    minibatch.logprobs,
+                    minibatch.advantages,
+                    minibatch.returns,
+                )
+                agent_state = agent_state.apply_gradients(grads=grads)
+                return agent_state, (loss, pg_loss, v_loss, entropy_loss, approx_kl, grads)
+
+            agent_state, metrics = jax.lax.scan(update_minibatch, agent_state, shuffled_storage)
+            return (agent_state, key), metrics
+
+        (agent_state, key), (loss, pg_loss, v_loss, entropy_loss, approx_kl, grads) = jax.lax.scan(
+            update_epoch, (agent_state, key), (), length=args.update_epochs
+        )
+        return agent_state, loss, pg_loss, v_loss, entropy_loss, approx_kl, key
+
+    # --- Main training loop ---
+    global_step = 0
+    start_time = time.time()
+
+    # Reset once to get initial state
+    print("Resetting the environment...")
+    next_env_state = env.reset(seed=args.seed)
+    next_obs = convert_obs_dict_to_array(next_env_state.observations)
+    next_done = jnp.zeros(args.num_envs, dtype=jnp.bool_)
+
+    def step_once(carry, _, env_step_fn):
+        agent_state, episode_stats, obs, done, key, env_state = carry
+        action, logprob, value, key = get_action_and_value_noise(agent_state, obs, key)
+
+        episode_stats, env_state, (next_obs, reward, next_done) = env_step_fn(
+            episode_stats, env_state, action
+        )
+
+        storage = Storage(
+            obs=obs,
+            actions=action,
+            logprobs=logprob,
+            dones=done,
+            values=value,
+            rewards=reward,
+            returns=jnp.zeros_like(reward),
+            advantages=jnp.zeros_like(reward),
+        )
+        return (agent_state, episode_stats, next_obs, next_done, key, env_state), storage
+
+    def rollout(
+        agent_state, episode_stats, next_obs, next_done, key, env_state, step_once_fn, max_steps
+    ):
+        (agent_state, episode_stats, next_obs, next_done, key, env_state), storage = jax.lax.scan(
+            step_once_fn,
+            (agent_state, episode_stats, next_obs, next_done, key, env_state),
+            (),
+            max_steps,
+        )
+        return agent_state, episode_stats, next_obs, next_done, storage, key, env_state
+
+    rollout = partial(
+        rollout,
+        step_once_fn=partial(step_once, env_step_fn=step_env_wrapped),
+        max_steps=args.num_steps,
+    )
+
+    print("Starting training...")
+    iters_bar = tqdm.tqdm(range(1, args.num_iterations + 1))
+    for _ in iters_bar:
+        iteration_time_start = time.time()
+
+        agent_state, episode_stats, next_obs, next_done, storage, key, next_env_state = rollout(
+            agent_state, episode_stats, next_obs, next_done, key, next_env_state
+        )
+
+        global_step += args.num_steps * args.num_envs
+        storage = compute_gae(agent_state, next_obs, next_done, storage)
+        agent_state, loss, pg_loss, v_loss, entropy_loss, approx_kl, key = update_ppo(
+            agent_state, storage, key
+        )
+
+        avg_episodic_return = np.mean(jax.device_get(episode_stats.returned_episode_returns))
+        iters_bar.set_postfix_str(
+            f"global_step={global_step}, avg_episodic_return={avg_episodic_return}"
+        )
+
+        writer.add_scalar("charts/avg_episodic_return", avg_episodic_return, global_step)
+        writer.add_scalar(
+            "charts/avg_episodic_length",
+            np.mean(jax.device_get(episode_stats.returned_episode_lengths)),
+            global_step,
+        )
+        writer.add_scalar(
+            "charts/learning_rate",
+            agent_state.opt_state[1].hyperparams["learning_rate"].item(),
+            global_step,
+        )
+        writer.add_scalar("losses/value_loss", v_loss[-1, -1].item(), global_step)
+        writer.add_scalar("losses/policy_loss", pg_loss[-1, -1].item(), global_step)
+        writer.add_scalar("losses/entropy", entropy_loss[-1, -1].item(), global_step)
+        writer.add_scalar("losses/approx_kl", approx_kl[-1, -1].item(), global_step)
+        writer.add_scalar("losses/loss", loss[-1, -1].item(), global_step)
+
+        # iters_bar.set_postfix_str(f"SPS: {int(global_step / (time.time() - start_time))}")
+
+        writer.add_scalar("charts/SPS", int(global_step / (time.time() - start_time)), global_step)
+        writer.add_scalar(
+            "charts/SPS_update",
+            int(args.num_envs * args.num_steps / (time.time() - iteration_time_start)),
+            global_step,
+        )
+
+    if args.save_model:
+        model_path = f"runs/{run_name}/{args.exp_name}.cleanrl_model"
+        with open(model_path, "wb") as f:
+            f.write(
+                flax.serialization.to_bytes(
+                    [
+                        vars(args),
+                        [
+                            agent_state.params["network_params"],
+                            agent_state.params["actor_params"],
+                            agent_state.params["critic_params"],
+                        ],
+                    ]
+                )
+            )
+        print(f"model saved to {model_path}")
+
+    env.close()
+    writer.close()
+
+
+def main() -> None:
+    args = tyro.cli(PPOArgs)
+    train(args)
+
+
+if __name__ == "__main__":
+    main()