Tasks and RL training
Three names make up a training job:
- a
simulo.Taskis your learning problem: the scene, what the robot observes, the reward, when an episode ends, and how an environment resets; - a
simulo.LearningEnvruns many copies of that task side by side on one GPU; - a
simulo.RLTrainertrains 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.
Write a task
Section titled “Write a task”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.
The task lifecycle
Section titled “The task lifecycle”| 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().
Episode end conditions
Section titled “Episode end conditions”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_upTraining 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 annotationsstores every annotation as a string that is never evaluated, so-> torch.Tensorneeds notorch.simulo.Taskon 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.
Judging success
Section titled “Judging success”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.successdef 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.successdef 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.
Train it
Section titled “Train it”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.
Playing a policy back
Section titled “Playing a policy back”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.