Batch RL Training

uvms_rl is the Python training API installed by simlab. It runs many UVMS environments through the batched simulator core without launching ROS nodes. ros2_control owns the hardware integration path; RL tasks, rewards, resets, and experiment configs live in uvms_rl.

Quick Start

cd ~/ros_ws
colcon build --packages-select ros2_control_blue_reach_5 simlab
source install/setup.bash

ros2 run simlab uvms_rl_train_rsl \
  --config hover_vehicle_stage1_fixed \
  --backend gpu \
  --device cuda \
  --num-envs 2048

hover_vehicle_stage1_fixed is the first curriculum experiment. It trains a fixed-target hover policy through the GPU batch backend and writes logs and checkpoints under recordings/rl_runs unless --log-dir is supplied.

Python Training Loop

import numpy as np
from uvms_rl import UvmsBatchEnv
from uvms_rl.config import load_experiment


experiment = load_experiment("hover_vehicle")
env_cfg = experiment.config["env"]
task_cfg = experiment.config["task"]

env = UvmsBatchEnv(
    robot_count=env_cfg["robot_count"],
    control_dt=env_cfg["control_dt"],
    sim_dt=env_cfg["sim_dt"],
    max_episode_steps=env_cfg["max_episode_steps"],
    seed=env_cfg["seed"],
    backend=env_cfg.get("backend", "cpu"),
    dynamics_profile=env_cfg["dynamics_profile"],
    task=experiment.task_cls,
    task_config=task_cfg,
)

obs = env.reset()
actions = np.zeros((env.robot_count, env.action_dim), dtype=np.float32)
obs, rewards, dones, info = env.step(actions)

obs is the task observation used by the policy. Use env.sim_observations() only when you need the raw simulator state.

RSL-RL Adapter

Use RslRlUvmsEnv when training with RSL-RL. The adapter wraps UvmsBatchEnv as RSL-RL’s VecEnv interface and returns observations as a TensorDict with the policy observation group.

Install RSL-RL first if it is not already present:

python3 -m pip install rsl-rl-lib

from rsl_rl.runners import OnPolicyRunner
from uvms_rl.rsl_adapter import RslRlUvmsEnv


env = RslRlUvmsEnv.from_experiment("hover_vehicle")
runner = OnPolicyRunner(env, train_cfg, log_dir="runs/hover_vehicle", device=str(env.device))
runner.learn(num_learning_iterations=1000)

The packaged command trains from the experiment’s trainer.rsl_rl config:

ros2 run simlab uvms_rl_train_rsl --config hover_vehicle_stage1_fixed --backend gpu --device cuda

Training logs and checkpoints are written under recordings/rl_runs by default. Override the backend, number of environments, iteration count, or log directory from the command line:

ros2 run simlab uvms_rl_train_rsl \
  --config hover_vehicle_stage1_fixed \
  --backend gpu \
  --device cuda \
  --num-envs 2048 \
  --iterations 3 \
  --log-dir /tmp/uvms_rsl_hover_stage1

RSL-RL does not call reset before training, so the adapter resets once at construction. During training it performs same-step resets for completed environment rows: returned dones still mark the terminal transition, while the returned observations for those rows are already the next episode’s initial observations.

Timing Contract

The RL API separates the policy/control step from the internal simulation step. One env.step(action) applies one policy action, holds it constant, advances the simulator for one or more internal substeps, then returns one observation, reward, and done flag per environment.

control_dt = policy action interval
sim_dt     = internal simulator integration interval
substeps   = control_dt / sim_dt

For transfer to the ROS runtime, match control_dt to the action interval used by the controller path. The ROS runtime uses controller_manager.update_rate as the ros2_control loop frequency, so start with:

control_dt = 1.0 / controller_manager.update_rate

Set sim_dt to the simulator integration period. If the simulator integrates at the same rate as the controller, use sim_dt: control_dt. If it integrates faster, use the smaller simulator period and let uvms_rl substep.

Example for a 150 Hz controller and 600 Hz simulator:

env:
  control_dt: 0.006666666666666667
  sim_dt: 0.001666666666666667

The packaged hover experiments set sim_dt equal to control_dt for a simple one-substep rollout. For transfer experiments, set both values explicitly to the controller and simulator rates you want to match.

Data Contract

The first dimension is the environment batch size.

task observations: [robot_count, task.policy_observation_dim]
raw sim state:     [robot_count, 22]
actions:           [robot_count, 11]
rewards:           [robot_count]
dones:             [robot_count]

Raw simulator state layout:

x, y, z, roll, pitch, yaw,
u, v, w, p, q, r,
arm_q1..arm_q5,
arm_qd1..arm_qd5

The vehicle pose is the dynamics NED state. z is positive depth/down, so underwater hover targets should use positive target_z values.

Action layout:

vehicle_wrench: force.x, force.y, force.z, torque.x, torque.y, torque.z
arm_torque: arm_joint_1..arm_joint_5

Add an Experiment

Add experiments as folders under uvms-simlab/uvms_rl/experiments. An experiment should own its config. Reusable task implementations should live in uvms-simlab/uvms_rl/tasks so multiple curriculum stages can share the same reward/reset logic without importing from each other.

uvms_rl/experiments/my_experiment/
|-- config.yaml
`-- task.py

Create a reusable task implementation, for example uvms_rl/tasks/my_task.py:

from uvms_rl.task_base import TaskBase


class Task(TaskBase):

    @property
    def policy_observation_dim(self) -> int:
        return 33

    def reset(self, env):
        # Return initial raw simulator state with shape [N, 22].
        ...

    def reset_indices(self, env, indices):
        # Return reset raw simulator states with shape [len(indices), 22].
        # Required when using RSL-RL same-step resets.
        ...

    def policy_observation(self, env, sim_obs, actions):
        # Return policy input with shape [N, policy_observation_dim].
        ...

    def reward_done(self, env, sim_obs, actions):
        # Return rewards [N], dones [N], and metric dict.
        ...

Point the experiment folder at that task with a tiny uvms_rl/experiments/my_experiment/task.py:

from uvms_rl.tasks.my_task import Task

Add uvms_rl/experiments/my_experiment/config.yaml:

env:
  backend: cpu
  dynamics_profile: dory_alpha
  robot_count: 1024
  control_dt: 0.006666666666666667
  sim_dt: 0.001666666666666667
  max_episode_steps: 500
  seed: 7

task:
  target_x: [-2.0, 2.0]

Rebuild and run:

cd ~/ros_ws
colcon build --packages-select simlab
source install/setup.bash
ros2 run simlab uvms_rl_train_rsl --config my_experiment --backend cpu --iterations 1

Backend Boundary

The task API is independent of the dynamics backend. CPU dynamics are the portable baseline. GPU dynamics use the CUDA batch core for high-throughput rollouts under the same task, reward, reset, and trainer code.

Experiment configs should name the backend explicitly:

env:
  backend: cpu
  dynamics_profile: dory_alpha

Use backend: gpu only on machines where ros2_control_blue_reach_5 was built with CUDA dynamics enabled. It does not silently fall back to CPU; if the GPU extension or PyTorch CUDA support is missing, construction fails so the training run cannot accidentally use the wrong backend.

dynamics_profile is also required. It names the same whole-robot dynamics profile used by the hardware simulator replay path. Missing profiles fail at environment construction time instead of silently falling back to hidden defaults.

Before training against the GPU backend, validate the installed GPU dynamics once after building:

ros2 run ros2_control_blue_reach_5 uvms_gpu_dynamics_correctness
ros2 run ros2_control_blue_reach_5 uvms_gpu_dynamics_benchmark

After those pass, switch the experiment config to backend: gpu and use the same task and trainer code.