Skip to content

Tasks and RL training

Three names make up a training job:

  • a simulo.Task is your learning problem: the scene, what the robot observes, the reward, when an episode ends, and how an environment resets;
  • a simulo.LearningEnv runs many copies of that task side by side on one GPU;
  • a simulo.RLTrainer trains a policy on the environment and saves its checkpoints.

The task knows nothing about the training algorithm, so the same task class can be used by every job file in your project. That is why the starter keeps it in task.py and imports it from both preview.py and train.py. Preview the task before training, and after every task or robot change. See Preview.

This is the task from the simulo create starter, shortened:

from __future__ import annotations
import simulo
cartpole = simulo.Asset.from_registry("simulo/robot/cartpole:v1")
app = simulo.App("my-app")
with app.runtime.imports():
import torch # resolved only on the GPU machine
class StarterTask(simulo.Task):
observation_dim = 4
action_dim = 1
episode_length_s = 5.0
def build(self, scene: simulo.Scene) -> None:
scene.add(simulo.Terrain.plane(name="ground"), 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")
def get_observations(self) -> torch.Tensor: ...
def get_rewards(self) -> torch.Tensor: ...
def get_dones(self) -> tuple[dict[str, torch.Tensor], torch.Tensor]: ...
def apply_actions(self, actions: torch.Tensor) -> None:
self.robot.set_joint_effort_target(100.0 * actions, joint_ids=self._cart_dof_idx)
def reset_idx(self, env_ids: torch.Tensor) -> None:
self.robot.reset(env_ids)

observation_dim and action_dim are the sizes of what the robot observes and the actions it takes. A checkpoint records both, and continuing from a checkpoint requires them to stay the same.

Method Called Purpose
build(scene) Once, when the environment is built Declare what exists: ground, lights, the robot, props, sensors.
on_start(env) Once, after the simulation is live Look up joint indices and anything else that needs the built scene.
get_observations() Every step Return the observation batch, one row per environment.
get_rewards() Every step Return one reward per environment.
get_dones() Every step Return (terminated, truncated). terminated is a dictionary of named conditions that end an episode before its time limit. truncated is one flag per environment.
apply_actions(actions) Every step Apply the policy’s actions to the robot. The actions are your task’s own copy, so scaling them in place, as in actions *= 0.05, is safe.
on_post_physics_step() Every step, after physics and before get_dones() Optional. Work that must happen exactly once per step, such as advancing a state machine or updating an actuator.
reset_idx(env_ids) When episodes end Reset only the environments listed in env_ids.

Each step runs in this order: apply_actions(), the physics, on_post_physics_step(), get_dones(), get_rewards(), get_observations(), and then reset_idx() for the environments whose episodes ended. The first observation of a run is different: reset_idx() runs for every environment and get_observations() follows straight away, with no physics step in between.

Scene, Robot, Prop & World covers what goes inside build().

An early end is an episode that ends before its time limit because something happened, such as a fall. get_dones() returns (terminated, truncated). terminated is a dictionary: each key names one early-end condition, and each value is a boolean tensor with one entry per environment that is true where the condition applies. Return {} when the task never ends episodes early. truncated is one boolean tensor with one entry per environment, usually for the time limit.

def get_dones(self) -> tuple[dict[str, torch.Tensor], torch.Tensor]:
# _pole_angle() and _cart_pos() are your own helpers in this task.
pole_fell = self._pole_angle().abs() > 0.5
off_track = self._cart_pos().abs() > 2.4
time_up = self.episode_length_buf >= self.max_episode_length - 1
return {"pole fell": pole_fell, "cart left the track": off_track}, time_up

Training ends an environment’s episode when any dictionary value is true. The report names the first true condition in the task’s dictionary order. Use a short past-tense phrase a person can read in that report, such as "pole fell" or "block dropped", not a variable name or metric. A success rule calls task.ended_by("block lifted"), so changing a condition name breaks that rule.

Each dictionary must have:

  • at most 32 unique names;
  • names that are non-empty, at most 64 characters, and have no surrounding spaces; and
  • the same names on every step.

"time limit" and "cut short by the task" are reserved in any letter case because the report uses them for a time-limit ending and an ending the task cut short. An invalid dictionary stops the job with a message that names the condition.

If the task uses TerminationSet, give every condition a name and return self.terminations.check() directly from get_dones(). Do not wrap that result in another dictionary.

After reset_idx() returns, get_observations() must describe the new episode. Some tasks keep copies of values such as a block’s position, refresh them after each physics step, and read them in get_observations(). A reset moves the robot and props, but those copies still hold the old episode’s positions, so the first observation of the new episode is wrong. At the end of reset_idx(), refresh every copy that get_observations() reads.

The simplest way to meet this is to keep no copies. Read the robot’s state where you use it, in each method that needs it:

def hand_position(self):
position, _ = self.robot.get_body_pose_in_base_frame("panda_hand")
return position
def get_rewards(self):
distance = torch.norm(self.goal_pos - self.hand_position(), dim=-1)
...

After Robot.reset() writes the robot’s joint state and root pose, a later get_body_pose_in_base_frame() read needs no task-level refresh: the runtime invalidates its cached body pose and recomputes it on demand. This is current within the normal environment update lifecycle. Keep a copy only for a value that is expensive to compute and that several methods share. Compute it in on_post_physics_step(), and again at the end of reset_idx() for the environments being reset.

Why this imports on a machine without PyTorch

Section titled “Why this imports on a machine without PyTorch”

simulo run imports your file on your own machine, where PyTorch is usually not installed. Three things keep that import cheap:

  • from __future__ import annotations stores every annotation as a string that is never evaluated, so -> torch.Tensor needs no torch.
  • simulo.Task on your machine is a lightweight base class. Your class is defined and type-checked normally; only the GPU machine swaps in the full implementation.
  • The one heavy import sits inside with app.runtime.imports():. On your machine it is recorded and skipped; on the GPU machine it is a normal import. See Runtimes.

build() and the other task methods run only on the GPU machine, so they can use simulo.Scene, simulo.Robot, and torch freely.

An evaluation judges each completed episode with the one eval.py rule that has @app.success. The rule runs on the final step before that environment resets and returns one boolean per environment. Evaluation explains how the rule becomes a success report. See Evaluate a trained policy for the complete rule and workflow.

Leave task unannotated in an @app.success rule. During type checking, simulo.Task is a lightweight protocol that does not expose these helpers, although the running task does. All three helpers raise RuntimeError outside an @app.success function.

Helper What it does Rule
task.survived() Returns true when no named end condition fired and the episode reached its real time limit. A task-cut-short truncated ending does not count as survival.
task.ended_by(name) Returns true when the named condition from get_dones() was true on the final step. An unknown name raises ValueError; use a name only when it means success for your task.
task.check(name, value) Records each environment’s true or false value for the report and returns value unchanged. value must be a bool tensor of shape (num_envs,); every call needs a unique name.

Name each check so that a yes is good, such as pole within 10° of upright, not pole tilted. get_dones() supplies the named conditions that ended_by() accepts. A time-limit ending is not survival unless the episode reached the real limit. An episode your task cut short with truncated does not count as survived.

Your rule may read public task state and robot state, including task.robot.state, task.robot.find_joints(...), and task.robot.get_body_pose_in_base_frame(...). It may also read public task attributes such as task.goal_pos. Do not read underscore-prefixed task attributes such as task._pole_dof_idx: they are implementation details that a task may rename without notice. If a rule needs a per-episode value such as a random goal, expose it as a public task attribute.

For a reach task, read the hand position and that public goal attribute:

@app.success
def reached(task):
hand, _ = task.robot.get_body_pose_in_base_frame("panda_hand")
return task.check("hand within 2 cm of the goal", (hand - task.goal_pos).norm(dim=-1) < 0.02)

For a task that ends as soon as it succeeds, the rule can use the condition name directly:

@app.success
def lifted(task):
return task.ended_by("block lifted")

Keep success-rule helpers in eval.py. Changing task.py after training can make an evaluation describe a task the checkpoint did not train against.

A training job builds the environment, trains, and closes both:

@app.job(type="train", system=simulo.SystemType.TIER_1, timeout=10 * 60)
def train(num_envs: int = 16, max_iterations: int = 2) -> dict:
env = simulo.LearningEnv(
task=StarterTask(), 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)
try:
stats = trainer.train(max_iterations=max_iterations)
return {"num_envs": num_envs, **stats}
finally:
trainer.close()
env.close()

LearningEnv settings that matter most:

Argument Meaning
num_envs How many copies of the task run side by side. More copies means more experience per iteration and more GPU memory.
dt The physics time step, in seconds.
physics_steps_per_action How many physics steps run for each action the policy takes.
env_spacing How far apart the copies sit, so they don’t collide.
seed The random seed, for repeatable runs.

RLTrainer.train(max_iterations=...) runs the training loop and returns a dict of summary statistics, such as best_reward. Returning it from the job makes it what simulo result prints.

Inside a type="train" job, train() also saves the policy’s best and latest checkpoints automatically. There is nothing to save yourself. Policies and checkpoints explains when saves happen and how best is chosen.

The object a training job builds, LearningEnv, is not one of the parallel environments. There is exactly one of it, and it drives all num_envs copies in lockstep. Every observation, reward, and robot state it hands back has one row per copy.

simulo.RLPlayer runs a trained policy without training it, for watching or recording. It is the building block for play jobs: simulo run play.py --policy policy_x:best steps the same task with the checkpoint’s actions instead of training it. See Play a policy and record it.