Skip to content
Simulo

Cartpole

Train a cart to keep its pole upright by pushing the cart left and right. This is the core Simulo training shape: a task that builds the scene and scores each step, and a job that trains it with PPO across many parallel environments.

Save these files in a folder called cartpole:

Show complete codeHide complete codecartpole/task.py
cartpole/task.py
from __future__ import annotations


import math
from typing import 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")


# 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")


# 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 __init__(self, with_camera: bool = False):
        super().__init__()
        self._with_camera = with_camera


    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")


        # Playback adds one world-frame camera. Training keeps the default off,
        # so its scene and render cost remain unchanged.
        if self._with_camera:
            scene.add(
                simulo.Camera(
                    width=640,
                    height=480,
                    data_types=["rgb"],
                    update_period=1.0 / 30.0,
                    offset=simulo.SensorOffset.look_at(
                        pos=(0.0, 0.0, 5.0), target=(0.0, 0.0, 0.0)
                    ),
                ),
                at="/World/overhead_cam",
                per_environment=False,
            )


    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")


    def get_observations(self) -> torch.Tensor:
        pole_idx = self._pole_dof_idx[0]
        cart_idx = self._cart_dof_idx[0]
        joint_pos = self.robot.state.joint_positions
        joint_vel = self.robot.state.joint_velocities
        pole_pos = joint_pos[:, pole_idx].view(-1, 1)
        pole_vel = joint_vel[:, pole_idx].view(-1, 1)
        cart_pos = joint_pos[:, cart_idx].view(-1, 1)
        cart_vel = 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:
        joint_pos = self.robot.state.joint_positions
        joint_vel = self.robot.state.joint_velocities
        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,
            joint_pos[:, self._pole_dof_idx[0]],
            joint_vel[:, self._pole_dof_idx[0]],
            joint_pos[:, self._cart_dof_idx[0]],
            joint_vel[:, self._cart_dof_idx[0]],
            self.reset_terminated,
        )


    def get_dones(self) -> Tuple[dict[str, torch.Tensor], torch.Tensor]:
        pole_idx = self._pole_dof_idx[0]
        cart_idx = self._cart_dof_idx[0]
        joint_pos = self.robot.state.joint_positions
        truncated = self.episode_length_buf >= self.max_episode_length - 1
        cart_out = torch.abs(joint_pos[:, cart_idx]) > self.max_cart_pos
        pole_fallen = torch.abs(joint_pos[:, pole_idx]) > math.pi / 2
        terminated = {"cart left the track": cart_out, "pole fell": 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]
        # Start from the default joint positions for the selected environments.
        joint_pos = self.robot.default_joint_positions[env_ids]
        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
        # Apply the randomized starting positions and default velocities together.
        joint_vel = self.robot.default_joint_velocities[env_ids]
        self.robot.set_joint_state(joint_pos, velocities=joint_vel, env_ids=env_ids)
Show complete codeHide complete codecartpole/train.py
cartpole/train.py
from __future__ import annotations


from typing import Any


import simulo


from task import CartpoleTask, app




@app.job(
    type="train",
    system=simulo.SystemType.TIER_1,
    timeout=8 * 60 * 60,
    retries=2,
)
def train(num_envs: int = 4096, max_iterations: int = 200) -> dict[str, Any]:
    """Train Cartpole with PPO."""
    env = simulo.LearningEnv(
        task=CartpoleTask(),
        num_envs=num_envs,
        device="cuda",
        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)
    trainer.close()
    env.close()
    return {"num_envs": num_envs, **stats}
Show complete codeHide complete codecartpole/eval.py
cartpole/eval.py
from __future__ import annotations


import math


import simulo


from task import CartpoleTask, app




@app.job(type="eval", system=simulo.SystemType.TIER_1, timeout=8 * 60 * 60)
def evaluate(
    policy: simulo.PolicyCheckpoint,
    episodes: int = 100,
    num_envs: int = 100,
    actions: str = "best",
):
    return simulo.evaluate(
        CartpoleTask(), policy, episodes=episodes, num_envs=num_envs, actions=actions
    )




@app.success
def balanced(task):
    """Count an episode when the pole remains upright until the time limit."""
    pole_angle = task.robot.state.joint_positions[
        :, task.robot.find_joints("cart_to_pole")[0]
    ]
    lasted = task.check("lasted to the time limit", task.survived())
    upright = task.check(
        "pole within 10 degrees of upright", pole_angle.abs() < math.radians(10)
    )
    return lasted & upright
Show complete codeHide complete codecartpole/play.py
cartpole/play.py
from __future__ import annotations


from task import CartpoleTask, app


import simulo




@app.job(type="play", system=simulo.SystemType.TIER_1, timeout=15 * 60)
def play(policy: simulo.PolicyCheckpoint, episodes: int = 3) -> None:
    """Play a policy: simulo run play.py --policy policy_x:best."""
    env = simulo.LearningEnv(
        task=CartpoleTask(with_camera=True),
        num_envs=1,
        device="cuda",
        dt=1.0 / 120.0,
        physics_steps_per_action=2,
        env_spacing=4.0,
        headless=True,
        seed=42,
        enable_cameras=True,
    )
    player = simulo.RLPlayer(env=env, policy=policy, device="cuda")
    try:
        player.play(
            num_episodes=episodes, record=simulo.RecordConfig(include_video=True)
        )
    finally:
        player.close()
        env.close()

The task observes 4 values (pole angle, pole angular velocity, cart position, cart velocity) and emits 1 action, the horizontal force on the cart. An episode ends when the pole tips past 90 degrees, when the cart leaves the track, or after 5 seconds.

Run it

Section titled “Run it”
Terminal window
simulo run cartpole/train.py --num-envs 512 --max-iterations 600 --tag baseline

The flags override the job defaults, --num-envs 4096 and --max-iterations 200. --tag baseline labels the job so you can find it later, and you can repeat --tag to add more labels. The job asks for a Tier 1 GPU, has an 8-hour timeout, and is retried up to 2 times after a failure. All attempts share the one timeout.

While it trains, the logs show Checkpoint written at iteration N/600 as training progresses and Best checkpoint saved at iteration N (reward=…) each time the mean episode reward improves. When it finishes, simulo run prints the policy’s checkpoints and the next commands to try:

Checkpoints:  best (iteration <n>)  latest (iteration 600)
Next: simulo policy get policy_brisk-heron-4f8k2m:best
      simulo run cartpole/eval.py --policy policy_brisk-heron-4f8k2m:best
      simulo run cartpole/train.py --max-iterations 1200 --from policy_brisk-heron-4f8k2m:best
      simulo export policy_brisk-heron-4f8k2m:best

simulo result prints what the job returned, including best_reward and iterations. simulo policy list lists the policy with its best and latest checkpoints.

To keep training the same policy, start a new job from one of its checkpoints. --max-iterations is the total, so this run trains the 300 iterations the checkpoint doesn’t have yet:

Terminal window
simulo run cartpole/train.py --num-envs 512 --max-iterations 900 \
  --from policy_brisk-heron-4f8k2m:latest

Evaluate it

Section titled “Evaluate it”

The success rule counts an episode only when the pole lasts to the time limit and stays within 10 degrees of upright. Run simulo run cartpole/eval.py --policy <policy-id>:best; see Evaluate a trained policy.

Play it back

Section titled “Play it back”

play.py runs a checkpoint back through the same task and records what happens. It builds CartpoleTask(with_camera=True), which adds one overhead camera on top of the training scene:

Terminal window
simulo run cartpole/play.py --policy policy_brisk-heron-4f8k2m:best
PLAY  job_swift-otter-8h4q2t
Policy:  policy_brisk-heron-4f8k2m:best (iteration 600)
...
[RLPlayer] Completed: 3 episodes, <steps> steps
Next: simulo recordings job_swift-otter-8h4q2t

Download the recording the command names. It holds the policy’s actions and the camera’s video:

Terminal window
simulo recordings job_swift-otter-8h4q2t

See Play a policy and record it for what the checkpoint’s task-compatibility checks do and what to do when they fail.

What to reuse

Section titled “What to reuse”
  • Task keeps scene construction, observations, rewards, dones, actions, and resets in one class you can test on its own.
  • with app.runtime.imports(): defers import torch until the job runs on the GPU, so submitting stays fast and needs no GPU libraries on your machine.
  • A training job gets its best and latest checkpoints automatically. The job body only builds the environment, trains, and returns its statistics.
  • The camera in play.py only exists in playback; training stays headless and unchanged.

See Tasks and RL training for the lifecycle behind the task methods, Continue training for starting a job from a checkpoint, and Policies and checkpoints for what a checkpoint holds.