JetBot

Train a differential-drive robot to turn toward a randomly changing direction and move along it. Two action values directly control the left and right wheels.

Save this complete app as jetbot/app.py:

Show complete codeHide complete codejetbot/app.py

jetbot/app.py

from __future__ import annotations

import math
from typing import Any, Tuple

import simulo

# The JetBot robot — a validated, version-pinned global-catalog asset.
jetbot = simulo.Asset.from_registry("simulo/robot/jetbot:v2")

# A named, durable, writable volume for the trained checkpoint. Metadata handle only
# at packaging time; the runner materialises it at execution (local:
# ~/.simulo/volumes/jetbot-checkpoints/, cloud: a durable volume). The job reads the
# real path via ``vol.path`` (execution-mode only).
vol = simulo.Volume.from_name("jetbot-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("jetbot", 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)


def _quat_to_forward(quat: torch.Tensor) -> torch.Tensor:
    """Rotate the unit X vector ``[1, 0, 0]`` by a per-env quaternion (wxyz).

    Plain typed module-level helper (not itself decorated) — ``torch.jit.script``
    recursively compiles it when it is called from the ``@app.runtime.torch_jit`` reward
    kernel below, and it also runs eagerly when called directly from
    ``JetbotTask.get_observations``.
    """
    w, x, y, z = quat[:, 0], quat[:, 1], quat[:, 2], quat[:, 3]
    forward_x = 1.0 - 2.0 * (y * y + z * z)
    forward_y = 2.0 * (x * y + w * z)
    forward_z = 2.0 * (x * z - w * y)
    return torch.stack([forward_x, forward_y, forward_z], dim=-1)


@app.runtime.torch_jit
def _compute_rewards(
    rew_scale_alignment: float,
    rew_scale_velocity: float,
    root_quat: torch.Tensor,
    root_lin_vel: torch.Tensor,
    commands: torch.Tensor,
) -> torch.Tensor:
    """JIT-compiled alignment + velocity reward kernel.

    ``@app.runtime.torch_jit`` is a no-op marker at submit and real ``torch.jit.script`` on
    the worker, so this lives at module level and is torch-free to *define* during
    discovery (its body never runs at submit). ``_quat_to_forward`` is compiled
    transitively — TorchScript recursively scripts plain typed functions it calls.
    """
    forward = _quat_to_forward(root_quat)
    alignment = (forward * commands).sum(dim=-1)
    velocity_in_cmd_dir = (root_lin_vel[:, :2] * commands[:, :2]).sum(dim=-1)
    reward = rew_scale_alignment * alignment + rew_scale_velocity * velocity_in_cmd_dir
    # ``.view(-1)`` is harmless belt-and-braces here — nothing above ``.squeeze()``s,
    # so ``reward`` is already (num_envs,) even when num_envs == 1 — but it locks in
    # the per-env reward contract explicitly rather than relying on that being true.
    return reward.view(-1)


class JetbotTask(simulo.Task):
    """Drive a two-wheeled Jetbot in a commanded direction, as fast as possible.

    Observation (6-dim): forward direction unit vector (3) + commanded direction unit
    vector (3, XY-plane). Action (2-dim): left / right wheel angular velocity.
    """

    observation_dim = 6
    action_dim = 2

    episode_length_s = 5.0
    velocity_scale = 10.0  # [rad/s] wheel angular velocity scale

    rew_scale_alignment = 1.0
    rew_scale_velocity = 0.5

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

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

    def on_start(self, env: simulo.LearningEnv) -> None:
        left = self.robot.find_joints("left_wheel_joint")
        right = self.robot.find_joints("right_wheel_joint")
        self._wheel_joint_ids = left + right

        self._commands = torch.zeros((self.num_envs, 3), device=self.device)
        self._randomize_commands(torch.arange(self.num_envs, device=self.device))

    def get_observations(self) -> torch.Tensor:
        # 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).
        # pose is [x, y, z, qw, qx, qy, qz]; the quaternion is the last four columns.
        forward = _quat_to_forward(self.robot.state.pose[:, 3:7])
        return torch.cat([forward, self._commands], dim=-1)

    def get_rewards(self) -> torch.Tensor:
        return _compute_rewards(
            self.rew_scale_alignment,
            self.rew_scale_velocity,
            self.robot.state.pose[:, 3:7],
            self.robot.state.linear_velocity,
            self._commands,
        )

    def get_dones(self) -> Tuple[torch.Tensor, torch.Tensor]:
        truncated = self.episode_length_buf >= self.max_episode_length - 1
        terminated = torch.zeros_like(truncated)
        return terminated, truncated

    def apply_actions(self, actions: torch.Tensor) -> None:
        scaled_velocities = actions * self.velocity_scale
        self.robot.set_joint_velocity_target(scaled_velocities, joint_ids=self._wheel_joint_ids)

    def reset_idx(self, env_ids: torch.Tensor) -> None:
        num_resets = len(env_ids)
        if num_resets == 0:
            return
        self.robot.reset(env_ids)
        self._randomize_commands(env_ids)

    def _randomize_commands(self, env_ids: torch.Tensor) -> None:
        """Generate new random XY-plane unit-vector direction commands."""
        n = len(env_ids)
        angles = torch.rand(n, device=self.device) * 2 * math.pi
        self._commands[env_ids, 0] = torch.cos(angles)
        self._commands[env_ids, 1] = torch.sin(angles)
        self._commands[env_ids, 2] = 0.0


# Retries are safe now that resumability exists: ResumableCheckpoint declares periodic
# checkpoints (every 50 iterations) and resume defaults to "auto" — see cartpole's
# identical comment on its own job.
@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,
    callbacks=[simulo.callbacks.ResumableCheckpoint(every=50)],
)
def train_jetbot(num_envs: int = 512, max_iterations: int = 700) -> dict[str, Any]:
    """Train the Jetbot direction-following policy with PPO and save the checkpoint.

    Args:
        num_envs: Number of parallel environments to simulate. The default (512) is
            sized for a ~2-3 minute demo run on a single RTX 3090 — see the module
            docstring's calibration note.
        max_iterations: Number of PPO policy-update iterations. The default (700) sees
            reward rise off the random-policy floor and converge (this task is simple
            enough to plateau well within the demo window — see the calibration note).

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

    stats = trainer.train(max_iterations=max_iterations)

    checkpoint = f"{vol.path}/jetbot_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 jetbot/app.py --num-envs 512 --max-iterations 700

The reward rises as the robot points and drives in the commanded direction. The completed result names the saved checkpoint. Use --num-envs 64 --max-iterations 2 for a launch smoke test.

What to reuse

• Express heading as a unit vector so the observation stays compact and has no angle wrap-around.
• Use one action per wheel when the actuator layout already matches the control problem.
• Resample the command on reset to train one policy across many directions.

See Volumes & Assets for the pinned JetBot asset and Tasks & the RL Loop for the training structure.