Cartpole

Train a cart to keep its pole upright by applying horizontal force. This is the core Simulo learning shape: build a task, create parallel environments, train a policy, and save the checkpoint to a durable volume.

Save this complete app as cartpole/app.py:

Show complete codeHide complete codecartpole/app.py

cartpole/app.py

from __future__ import annotations

import math
from typing import Any, Tuple

import simulo

# The cartpole robot — a validated, version-pinned global-catalog asset.
# Capturing it at module level lets submit resolve this exact version before the
# job runs; the worker mounts it and the Task below consumes the same handle.
cartpole = simulo.Asset.from_registry("simulo/robot/cartpole:v1")

# A named, durable, writable volume for the trained checkpoint. This line is just a
# metadata handle — it creates nothing at packaging time. At *execution* the runner
# materialises it: under local execution as ~/.simulo/volumes/cartpole-checkpoints/,
# in the cloud as a durable volume. The job reads the real path via ``vol.path``
# (execution-mode only), never the ``/out`` mount point declared below.
vol = simulo.Volume.from_name("cartpole-checkpoints", create_if_missing=True)

# Advanced: pick a different Simulo runtime with
# App("name", runtime=simulo.Runtime.from_registry("simulo/gpu-rl:2026.06"));
# see the Runtimes docs.
app = simulo.App("cartpole", mounts={"/out": vol})

# The ONE module-level heavy import — deferred under the runtime guard so discovery
# records it as a remote import instead of resolving it.
with app.runtime.imports():
    import torch  # noqa: F401  (resolved only in execution mode, on the worker)


@app.runtime.torch_jit
def _compute_rewards(
    rew_scale_alive: float,
    rew_scale_terminated: float,
    rew_scale_pole_pos: float,
    rew_scale_cart_vel: float,
    rew_scale_pole_vel: float,
    pole_pos: torch.Tensor,
    pole_vel: torch.Tensor,
    cart_pos: torch.Tensor,
    cart_vel: torch.Tensor,
    reset_terminated: torch.Tensor,
) -> torch.Tensor:
    """JIT-compiled reward kernel (the classic cartpole balance reward).

    ``@app.runtime.torch_jit`` is a no-op marker at submit (no ``torch`` locally) and
    real ``torch.jit.script`` on the worker — so this lives at module level and is
    still torch-free to *define* during discovery (its body never runs at submit).
    """
    pole_pos = pole_pos.squeeze()
    pole_vel = pole_vel.squeeze()
    cart_pos = cart_pos.squeeze()
    cart_vel = cart_vel.squeeze()
    reset_terminated = reset_terminated.squeeze()

    rew_alive = rew_scale_alive * (1.0 - reset_terminated.float())
    rew_termination = rew_scale_terminated * reset_terminated.float()
    rew_pole_pos = rew_scale_pole_pos * torch.square(pole_pos)
    rew_cart_vel = rew_scale_cart_vel * torch.abs(cart_vel)
    rew_pole_vel = rew_scale_pole_vel * torch.abs(pole_vel)

    reward: torch.Tensor = rew_alive + rew_termination + rew_pole_pos + rew_cart_vel + rew_pole_vel
    # Keep the (num_envs,) per-env reward contract even when num_envs == 1: the
    # squeezes above collapse a single-env batch to a 0-d scalar, which breaks
    # per-env consumers (e.g. the MCAP recorder's per-env /reward indexing).
    return reward.view(-1)


class CartpoleTask(simulo.Task):
    """Balance a pole on a cart by applying horizontal forces to the cart.

    Observation (4-dim): pole angle, pole angular velocity, cart position, cart
    velocity. Action (1-dim): scaled horizontal force on the cart.

    Defined at module level: ``simulo.Task`` is a torch-free contract stand-in at
    submit and the real ``simulo.core.Task`` on the worker, so the same class
    authors lean and trains heavy.
    """

    observation_dim = 4
    action_dim = 1

    episode_length_s = 5.0
    action_scale = 100.0  # [N]

    max_cart_pos = 3.0  # [m]
    initial_pole_angle_range = (-0.25, 0.25)  # fraction of pi [rad]

    rew_scale_alive = 1.0
    rew_scale_terminated = -2.0
    rew_scale_pole_pos = -1.0
    rew_scale_cart_vel = -0.01
    rew_scale_pole_vel = -0.005

    # Framework-injected at runtime by ``simulo.core.Task`` / ``LearningEnv``
    # (declared here only so the type checker sees the names the methods read; the
    # annotations are PEP 563 strings and never shadow the inherited values).
    device: str
    max_episode_length: int
    episode_length_buf: torch.Tensor
    reset_terminated: torch.Tensor

    def build(self, scene: simulo.Scene) -> None:
        scene.add(simulo.Terrain.plane(name="ground"), at="/", per_environment=False)
        scene.add(
            simulo.Light.dome(name="light", intensity=2000.0, color=(0.75, 0.75, 0.75)),
            at="/",
            per_environment=False,
        )
        self.robot = simulo.Robot(asset=cartpole, initial_pose=simulo.Pose.identity())
        scene.add(self.robot, at="/World/Robot")

    def on_start(self, env: simulo.LearningEnv) -> None:
        self._cart_dof_idx = self.robot.find_joints("slider_to_cart")
        self._pole_dof_idx = self.robot.find_joints("cart_to_pole")
        # robot.state is the supported, typed way to read live state (robot.internals
        # is the unstable engine escape hatch — see the Scene, Robot & World docs).
        self._joint_pos = self.robot.state.joint_positions
        self._joint_vel = self.robot.state.joint_velocities

    def get_observations(self) -> torch.Tensor:
        pole_idx = self._pole_dof_idx[0]
        cart_idx = self._cart_dof_idx[0]
        pole_pos = self._joint_pos[:, pole_idx].view(-1, 1)
        pole_vel = self._joint_vel[:, pole_idx].view(-1, 1)
        cart_pos = self._joint_pos[:, cart_idx].view(-1, 1)
        cart_vel = self._joint_vel[:, cart_idx].view(-1, 1)
        return torch.cat((pole_pos, pole_vel, cart_pos, cart_vel), dim=-1)

    def get_rewards(self) -> torch.Tensor:
        return _compute_rewards(
            self.rew_scale_alive,
            self.rew_scale_terminated,
            self.rew_scale_pole_pos,
            self.rew_scale_cart_vel,
            self.rew_scale_pole_vel,
            self._joint_pos[:, self._pole_dof_idx[0]],
            self._joint_vel[:, self._pole_dof_idx[0]],
            self._joint_pos[:, self._cart_dof_idx[0]],
            self._joint_vel[:, self._cart_dof_idx[0]],
            self.reset_terminated,
        )

    def get_dones(self) -> Tuple[torch.Tensor, torch.Tensor]:
        self._joint_pos = self.robot.state.joint_positions
        self._joint_vel = self.robot.state.joint_velocities
        pole_idx = self._pole_dof_idx[0]
        cart_idx = self._cart_dof_idx[0]
        truncated = self.episode_length_buf >= self.max_episode_length - 1
        cart_out = torch.abs(self._joint_pos[:, cart_idx]) > self.max_cart_pos
        pole_fallen = torch.abs(self._joint_pos[:, pole_idx]) > math.pi / 2
        terminated = cart_out | pole_fallen
        return terminated, truncated

    def apply_actions(self, actions: torch.Tensor) -> None:
        self.robot.set_joint_effort_target(self.action_scale * actions, joint_ids=self._cart_dof_idx)

    def reset_idx(self, env_ids: torch.Tensor) -> None:
        num_resets = len(env_ids)
        if num_resets == 0:
            return
        self.robot.reset(env_ids)
        pole_idx = self._pole_dof_idx[0]
        # robot.state has no default-joint-value equivalent, so this stays on the
        # internals escape hatch (there is nothing unstable about reading it here,
        # just no supported, typed name for it yet).
        joint_pos = self.robot.internals.default_joint_pos[env_ids].clone()
        random_angles = torch.empty(num_resets, device=self.device).uniform_(
            self.initial_pole_angle_range[0] * math.pi,
            self.initial_pole_angle_range[1] * math.pi,
        )
        joint_pos[:, pole_idx] += random_angles
        # set_joint_state writes both positions and velocities through the engine's
        # own command path (Articulation.write_joint_state_to_sim ->
        # write_joint_{position,velocity}_to_sim), which updates robot.state's
        # backing buffers in place AND pushes to the physics view in the same call.
        # self._joint_pos / self._joint_vel are the SAME objects as those buffers
        # (the hold-safety contract on core/robot.py), so they are already current
        # after this call — no separate write into either tensor is needed, and
        # robot.state's contract is never write into a member's tensor directly.
        joint_vel = self.robot.internals.default_joint_vel[env_ids]
        self.robot.set_joint_state(joint_pos, velocities=joint_vel, env_ids=env_ids)


# Retries are safe now that resumability exists: ResumableCheckpoint declares
# periodic checkpoints (every 50 iterations) and resume defaults to "auto", so a
# retried / preempted run resumes from the latest checkpoint instead of restarting
# — the runner exports the checkpoint envelope and trainer.train() picks it up;
# the job body below stays unchanged.
@app.job(
    # Tier 1: T4 GPU, 16 GB VRAM. Run `simulo systems` for the full four-tier catalog.
    system=simulo.SystemType.TIER_1,
    timeout=8 * 60 * 60,
    retries=2,
    # Writes update one latest.pt/latest.json in place; pass keep_last=N to retain numbered iter_<n>.pt copies.
    callbacks=[simulo.callbacks.ResumableCheckpoint(every=50)],
)
def train_cartpole(num_envs: int = 4096, max_iterations: int = 200) -> dict[str, Any]:
    """Train a cartpole-balancing policy with PPO and save the checkpoint.

    All heavy work happens here, on the worker: ``simulo.LearningEnv`` and
    ``simulo.RLTrainer`` resolve to the real ``simulo.core`` training types in
    execution mode, and ``CartpoleTask`` (a module-level ``simulo.Task`` subclass)
    is instantiated against the GPU runtime.

    Args:
        num_envs: Number of parallel environments to simulate.
        max_iterations: Number of PPO policy-update iterations.

    Returns:
        A JSON-serialisable dict: the saved ``checkpoint`` path plus training
        ``stats`` (``iterations``, ``num_envs``).
    """
    env = simulo.LearningEnv(
        task=CartpoleTask(),
        num_envs=num_envs,
        device="cuda",
        dt=1.0 / 120.0,
        physics_steps_per_action=2,
        env_spacing=4.0,
        headless=True,
        seed=42,
    )
    trainer = simulo.RLTrainer(env=env, algorithm="PPO", device="cuda", seed=42)

    stats = trainer.train(max_iterations=max_iterations)

    # Persist the trained policy to the durable volume (resolved in execution mode).
    checkpoint = f"{vol.path}/cartpole_final.pt"
    trainer.save(checkpoint)

    # Close the trainer before the environment so skrl releases its resources first.
    trainer.close()
    env.close()

    return {"checkpoint": checkpoint, "num_envs": num_envs, **stats}

Run the app

simulo run cartpole/app.py --num-envs 4096 --max-iterations 200

Training logs will show reward progress. The completed result names the saved checkpoint, and simulo models lists the downloadable model files. For a short smoke run, use --num-envs 64 --max-iterations 2.

What to reuse

• Task keeps scene construction, observations, rewards, actions, and resets in one testable unit.
• app.runtime.imports() defers torch until the worker executes the app, so submission stays lightweight.
• ResumableCheckpoint gives retries and preempted runs a recent recovery point.

See Tasks & the RL Loop for the lifecycle behind these methods and Callbacks for checkpoint behavior.