Skip to content
Simulo

JetBot

Train a two-wheel robot to turn toward a commanded direction and drive along it. The direction changes at the start of every episode, so one policy learns to follow any heading. Two action values drive the left and right wheels directly.

Save this file as jetbot/train.py. It holds its own task, so it needs no task.py:

Show complete codeHide complete codejetbot/train.py
jetbot/train.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: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("jetbot")


# 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:
        # 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[dict[str, torch.Tensor], torch.Tensor]:
        truncated = self.episode_length_buf >= self.max_episode_length - 1
        terminated = {}
        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




@app.job(
    type="train",
    # 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,
)
def train_jetbot(num_envs: int = 16, max_iterations: int = 700) -> dict[str, Any]:
    """Train the Jetbot direction-following policy with PPO.


    Args:
        num_envs: Number of parallel environments to simulate. More environments give
            the trainer more experience per iteration and use more memory. The default
            of 16 is what was measured for this app on tier 1, where 64 and 512 are
            stopped for exceeding the job's memory budget before training starts. That
            is this workload's measured ceiling on that tier and not a platform-wide
            limit; ``simulo systems`` records no ceiling for the other tiers.
        max_iterations: Number of PPO policy-update iterations.


    Returns:
        A JSON-serialisable dict with training ``stats``.
    """
    env = simulo.LearningEnv(
        task=JetbotTask(),
        num_envs=num_envs,
        device="cuda",
        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)


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


    return {"num_envs": num_envs, **stats}
Show complete codeHide complete codejetbot/eval.py
jetbot/eval.py
from __future__ import annotations


import simulo


from train import JetbotTask, app as _training_app


app = simulo.App(
    _training_app.name, runtime=_training_app.runtime, mounts=_training_app.mounts
)
del _training_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 = 16,
    actions: str = "best",
):
    return simulo.evaluate(
        JetbotTask(), policy, episodes=episodes, num_envs=num_envs, actions=actions
    )




@app.success
def kept_moving(task):
    """Count an episode when the robot moves in its commanded direction."""
    command = task.commands
    return task.check(
        "moved at least 0.1 m/s in the commanded direction",
        (task.robot.state.linear_velocity[:, :2] * command[:, :2]).sum(dim=-1) >= 0.1,
    )

The task observes 6 values: the robot’s forward direction and the commanded direction, each as a unit vector. It emits 2 actions, the left and right wheel velocities. The reward pays for facing the commanded direction and for speed along it. Nothing ends an episode early; every episode runs to its time limit.

Run it

Section titled “Run it”
Terminal window
simulo run jetbot/train.py

The defaults are --num-envs 16 and --max-iterations 700. The job asks for a Tier 1 GPU, has an 8-hour timeout, and is retried up to 2 times after a failure. For a short launch check, add --max-iterations 2.

The 16-environment default is this app’s measured limit on Tier 1: 16 and 32 environments completed there, while 64 and 512 ran out of memory before training started. That limit belongs to this app. Other apps run far more environments on the same hardware. simulo systems shows the Tier 1 figure and lists the other tiers as not yet measured, so on a larger tier, find the right --num-envs by trying it.

On a local RTX 3090 with --num-envs 512, the best reward rose within the first 150 iterations and then held steady as the robot learned to point and drive in the commanded direction. The default collects far less experience per iteration, so expect that rise to come later. No reward has been measured at the default.

When the run finishes, simulo run prints the policy’s best and latest checkpoints and the next commands to try.

Evaluate it

Section titled “Evaluate it”

The success rule checks that the robot moves at least 0.1 m/s in its commanded direction. Run simulo run jetbot/eval.py --policy <policy-id>:best; see Evaluate a trained policy.

What to reuse

Section titled “What to reuse”
  • Give the heading as a unit vector so the observation stays small and never wraps around at 360 degrees.
  • Use one action per wheel when the robot’s actuators already match the control problem.
  • Pick a new command for each environment when it resets, so one policy learns every direction.

See Asset catalogs & version pinning for how the app pins its robot with simulo.Asset.from_registry and Tasks and RL training for the training structure.