Skip to content

Commit 239e865

Browse files
KraHsuKraHsu
authored
Merge pull request #32 from KraHsu/docs/rl-runner-concept
Add RL runner bilingual concept page
2 parents f14d0ba + 8cd9936 commit 239e865

4 files changed

Lines changed: 449 additions & 0 deletions

File tree

‎CHANGELOG.md‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,11 @@ trajectory so breaking changes can land in any minor release until the 1.0 stabi
88

99
### Added
1010

11+
- `docs/concepts/rl-runner.{en,zh}.md`: bilingual concept page for `genelab.rl`
12+
(`train_task` / `play_task` signatures, `RslRlOnPolicyRunnerCfg` field
13+
layout, `AgentKind` Literal, `RslRlVecEnvWrapper` adapter, torchrun-based
14+
`--gpus N` distributed-training relaunch, `maybe_profile` flag/env-var
15+
matrix, failure modes).
1116
- `docs/concepts/mdp.{en,zh}.md`: bilingual reference for the reusable MDP term
1217
library (`genelab.mdp.{actions,commands,observations,rewards,terminations,
1318
events,curriculums,noise}`). Every re-exported function / class is tabulated

‎docs/concepts/rl-runner.en.md‎

Lines changed: 228 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,228 @@
1+
# RL runner
2+
3+
`genelab.rl` is the thin RSL-RL integration that turns a registered task into
4+
a runnable PPO training or replay session. The CLI subcommands `play` and
5+
`train` dispatch into the two entry points covered below; library callers can
6+
also reach them directly.
7+
8+
## Module surface
9+
10+
| Name | Purpose |
11+
|---|---|
12+
| `train_task(task_id, agent_cfg, **kwargs) -> Path` | Build the task's env, wrap it in `RslRlVecEnvWrapper`, and drive RSL-RL's `OnPolicyRunner.learn`. Returns the resolved log directory. |
13+
| `play_task(task_id, **kwargs) -> None` | Build the task's env (using `play_env` when set), then loop `policy(obs) → env.step` until the viewer closes or `max_steps` is reached. Supports `"zero"` / `"random"` / `"trained"` policies. |
14+
| `AgentKind` | `Literal["zero", "random", "trained"]`. |
15+
| `RslRlOnPolicyRunnerCfg` | Top-level PPO config (extends `RslRlBaseRunnerCfg`); composes `RslRlModelCfg` for actor / critic and `RslRlPpoAlgorithmCfg` for the optimiser. |
16+
| `RslRlVecEnvWrapper` | Adapter from `ManagerBasedRlEnv` to RSL-RL's `VecEnv` protocol. Dynamic subclass of `rsl_rl.env.VecEnv` when RSL-RL is installed. |
17+
| `VecEnvBase` | Minimal stand-alone VecEnv ABC for non-RSL-RL callers. |
18+
| `maybe_profile(...)` | Context manager wrapping `torch.profiler`. No-op when disabled or off rank 0. |
19+
| `profiler_enabled()` | `True` when `GENELAB_PROFILE=1`. |
20+
21+
`genelab.cli` wires `play` / `train` into `play_task` / `train_task` via
22+
`_dispatch_play` / `_dispatch_train` in `src/genelab/cli/__init__.py:384,417`.
23+
24+
## train_task
25+
26+
`train_task` resolves the env cfg off the registered task, applies CLI
27+
overrides, wraps the env, runs PPO, and returns the log directory.
28+
29+
| Argument | Default | Effect |
30+
|---|---|---|
31+
| `task_id: str` | — | Required. Looked up in the `TASKS` registry. |
32+
| `agent_cfg: RslRlOnPolicyRunnerCfg` | — | Required. The CLI pulls this off the task's `cfg.agent`; library callers pass it directly. |
33+
| `num_envs: int \| None` | `None` | Overrides `env_cfg.simulation.num_envs`. |
34+
| `max_iterations: int \| None` | `None` | Overrides `agent_cfg.max_iterations`. |
35+
| `seed: int \| None` | `None` | Sets both `env_cfg.seed` and `agent_cfg.seed`. |
36+
| `log_root: Path \| None` | `Path("logs/rsl_rl")` | Parent under which `<experiment>/<timestamp>` is created. |
37+
| `log_dir: Path \| None` | `None` | Pre-resolved final directory. Takes precedence over `log_root`. The torchrun relaunch path uses this so every rank lands in the same folder. |
38+
| `resume_from: Path \| None` | `None` | Existing checkpoint to load before `runner.learn`. |
39+
| `prof`, `prof_out`, `prof_wait`, `prof_warmup`, `prof_active`, `prof_repeat`, `prof_record_shapes`, `prof_with_stack` | `None` | Forwarded to `maybe_profile`. Each kwarg overrides the matching `GENELAB_PROFILE_*` env var; see "Profiling integration" below. |
40+
41+
Side effects on rank 0:
42+
43+
- Creates `<log_dir>/params/env.json` and `<log_dir>/params/agent.json` with a
44+
recursive JSON dump of both cfgs.
45+
- RSL-RL writes `<log_dir>/model_<iter>.pt` every `agent_cfg.save_interval`
46+
iterations.
47+
- TensorBoard event files land under `<log_dir>/` directly.
48+
49+
When profiling is on, `train_task` hooks `RslRlVecEnvWrapper.step` so the
50+
profiler schedule advances once per env step. To express the schedule in PPO
51+
iterations instead, multiply `wait` / `warmup` / `active` by
52+
`agent_cfg.num_steps_per_env`.
53+
54+
## play_task
55+
56+
`play_task` builds the env (preferring `task_cfg.play_env` when set), picks a
57+
policy according to `agent`, then loops until the viewer closes.
58+
59+
| Argument | Default | Effect |
60+
|---|---|---|
61+
| `task_id: str` | — | Required. |
62+
| `checkpoint: Path \| None` | `None` | Required when `agent="trained"`. |
63+
| `num_envs: int \| None` | `None` | Overrides `env_cfg.simulation.num_envs`. |
64+
| `agent: "zero" \| "random" \| "trained" \| None` | `None` | Defaults to `"trained"` if `checkpoint` is given, else `"zero"`. |
65+
| `agent_cfg: RslRlOnPolicyRunnerCfg \| None` | `None` | Required only for `agent="trained"`; falls back to `task.cfg.agent`. |
66+
| `deterministic: bool` | `True` | Forwarded to RSL-RL inference helpers. |
67+
| `max_steps: int \| None` | `None` | Cap the rollout. `None` runs until the viewer closes. |
68+
| `prof*` | `None` | Same semantics as `train_task`. |
69+
70+
The play loop polls `env.viewer_closed` every step, so a closed Genesis
71+
viewer terminates the rollout cleanly without raising. `env.close()` runs in
72+
a `finally` block.
73+
74+
## RslRlOnPolicyRunnerCfg
75+
76+
The top-level dataclass an extension registers as `TaskCfg.agent`. Train
77+
relies on `class_name="OnPolicyRunner"`; an extension config with any other
78+
runner class is rejected at dispatch.
79+
80+
```
81+
RslRlOnPolicyRunnerCfg
82+
├── (inherited from RslRlBaseRunnerCfg)
83+
│ seed=42, num_steps_per_env=24, max_iterations=300
84+
│ obs_groups={"actor": ("policy",), "critic": ("critic",)}
85+
│ save_interval=50, experiment_name="exp1", run_name=""
86+
│ logger="tensorboard" | "wandb"
87+
│ wandb_project="genelab", wandb_tags=()
88+
│ resume=False, load_run=".*", load_checkpoint="model_.*.pt"
89+
│ clip_actions: float | None, upload_model=False
90+
├── class_name="OnPolicyRunner"
91+
├── actor: RslRlModelCfg
92+
├── critic: RslRlModelCfg
93+
└── algorithm: RslRlPpoAlgorithmCfg
94+
```
95+
96+
`RslRlModelCfg` knobs (per actor / critic): `hidden_dims=(128, 128, 128)`,
97+
`activation="elu"`, `obs_normalization=False`, optional `cnn_cfg` /
98+
`distribution_cfg` / `rnn_type` / `rnn_hidden_dim` / `rnn_num_layers`,
99+
`class_name="MLPModel"`. Only `class_name`, `hidden_dims`, `activation`,
100+
`obs_normalization`, and `distribution_cfg` reach the MLP constructor — other
101+
fields are pruned by `_prune_model_cfg`.
102+
103+
`RslRlPpoAlgorithmCfg` knobs: `num_learning_epochs=5`, `num_mini_batches=4`,
104+
`learning_rate=1e-3`, `schedule="adaptive"|"fixed"`, `gamma=0.99`, `lam=0.95`,
105+
`entropy_coef=0.005`, `desired_kl=0.01`, `max_grad_norm=1.0`,
106+
`value_loss_coef=1.0`, `use_clipped_value_loss=True`, `clip_param=0.2`,
107+
`optimizer="adam"|"adamw"|"sgd"|"rmsprop"`, `class_name="PPO"`.
108+
109+
## AgentKind
110+
111+
`AgentKind = Literal["zero", "random", "trained"]`. `play_task` selects the
112+
policy callable from `kind`:
113+
114+
| Kind | Policy | Use case |
115+
|---|---|---|
116+
| `"zero"` | Returns `torch.zeros((num_envs, num_actions))`. | Visual smoke before any training; passive-pose viewer rollouts. |
117+
| `"random"` | Returns `2.0 * rand(...) - 1.0` (uniform `[-1, 1]`). | Sanity-check action bounds and reward shaping. |
118+
| `"trained"` | Loads the checkpoint, instantiates `OnPolicyRunner`, and calls `runner.get_inference_policy()`. | Inspect a trained policy. Requires `checkpoint`. |
119+
120+
## VecEnv wrapper
121+
122+
`RslRlVecEnvWrapper` adapts `ManagerBasedRlEnv` to the interface RSL-RL's
123+
`OnPolicyRunner` expects:
124+
125+
| Attribute / method | Source |
126+
|---|---|
127+
| `num_envs`, `device`, `max_episode_length` | `env.num_envs`, `env.device`, `env.max_episode_length`. |
128+
| `num_actions` | `env.action_manager.total_action_dim`. |
129+
| `num_obs`, `num_privileged_obs` | Last-axis dim of the `"policy"` / `"critic"` observation group, pre-warmed via one `observation_manager.compute()` call at construct. |
130+
| `reset()` | `env.reset()`. Observations are wrapped in a `TensorDict` when `tensordict` is installed; otherwise returned as a plain `dict`. |
131+
| `step(actions)` | Clips actions to `clip_actions` when set, calls `env.step`, packs `(obs, reward, dones, extras)` with `dones = terminated \| truncated` (long) and `extras["time_outs"] = truncated`. |
132+
| `episode_length_buf`, `cfg`, `unwrapped` | Pass-throughs to the underlying env. |
133+
134+
The wrapper dynamically subclasses `rsl_rl.env.VecEnv` at import time via
135+
`_attach_rsl_rl_base()`. When RSL-RL is missing (e.g. unit tests on a
136+
CPU-only host) it falls back to a plain object with the same attributes.
137+
138+
## Distributed training
139+
140+
The CLI surfaces `--gpus N` on `train`. The flag is rewritten by
141+
`_relaunch_under_torchrun` in `cli/__init__.py:457`:
142+
143+
1. Pre-resolves the log directory under `<log_root>/<experiment>/<timestamp>`
144+
on the parent process so every rank picks up the same path.
145+
2. Strips `--gpus N` from `sys.argv` to avoid an infinite relaunch loop.
146+
3. Appends `--log-dir <pre-resolved>` when the original argv did not already
147+
carry one.
148+
4. `execvp`s `python -m torch.distributed.run --standalone --nproc_per_node=N
149+
-m genelab.cli <inner>`.
150+
151+
```bash
152+
uv run genelab train Genelab-Velocity-Flat-Unitree-G1-v0 --gpus 4 --num-envs 4096
153+
```
154+
155+
Inside each rank, `genelab.rl.distributed` reads `WORLD_SIZE` / `LOCAL_RANK` /
156+
`RANK` from the environment:
157+
158+
| Helper | Returns |
159+
|---|---|
160+
| `world_size()` | `int(os.environ["WORLD_SIZE"])` or `1`. |
161+
| `local_rank()` | `int(os.environ["LOCAL_RANK"])` or `0`. |
162+
| `global_rank()` | `int(os.environ["RANK"])` or `0`. |
163+
| `is_distributed()` | `world_size() > 1`. |
164+
| `is_main_process()` | `global_rank() == 0`. |
165+
| `pin_cuda_device()` | `"cuda:0"` in distributed runs (the CLI bootstrap pins `CUDA_VISIBLE_DEVICES` per worker, so each rank sees a single device); `None` in single-GPU mode. |
166+
167+
Only the main process writes checkpoints, TensorBoard events, profiler
168+
traces, and the `params/*.json` dump. RSL-RL handles the inter-rank gradient
169+
all-reduce internally.
170+
171+
## Profiling integration
172+
173+
`maybe_profile` wraps `torch.profiler.profile` with a `schedule` and a
174+
TensorBoard trace handler. The CLI surfaces eight flags on both `play` and
175+
`train`:
176+
177+
| CLI flag | `maybe_profile` kwarg | Env var fallback | Default |
178+
|---|---|---|---|
179+
| `--prof` | `enabled` | `GENELAB_PROFILE=1` | off |
180+
| `--prof-out` | `out_dir` | `GENELAB_PROFILE_OUT` | `logs/torch_profile` |
181+
| `--prof-wait` | `wait` | `GENELAB_PROFILE_WAIT` | `10` |
182+
| `--prof-warmup` | `warmup` | `GENELAB_PROFILE_WARMUP` | `5` |
183+
| `--prof-active` | `active` | `GENELAB_PROFILE_ACTIVE` | `10` |
184+
| `--prof-repeat` | `repeat` | `GENELAB_PROFILE_REPEAT` | `2` |
185+
| `--prof-record-shapes` | `record_shapes` | `GENELAB_PROFILE_RECORD_SHAPES=1` | off |
186+
| `--prof-with-stack` | `with_stack` | `GENELAB_PROFILE_WITH_STACK=1` | off |
187+
188+
The schedule fires one trace per `wait + warmup + active` steps, repeated
189+
`repeat` times. The handler writes `*.pt.trace.json` + `*.tar.gz` under
190+
`out_dir`; `uv run genelab prof open <dir>` launches TensorBoard against it.
191+
192+
```bash
193+
uv run genelab train Genelab-Velocity-Flat-Unitree-G1-v0 \
194+
--prof --prof-active 3 --prof-repeat 1 --max-iterations 10
195+
uv run genelab prof open logs/torch_profile
196+
```
197+
198+
## Failure modes worth knowing
199+
200+
!!! warning "Mismatched agent cfg class"
201+
`_dispatch_train` rejects any `task.cfg.agent` whose type is not
202+
`RslRlOnPolicyRunnerCfg`. Extensions importing a hand-rolled runner cfg
203+
that "happens to" have the same fields are not accepted.
204+
205+
!!! warning "agent='trained' requires --checkpoint"
206+
`play_task` raises `SystemExit("agent='trained' requires a --checkpoint
207+
path")` when the checkpoint is omitted. Either pass `--checkpoint` or
208+
pick `--agent zero` / `--agent random` for the visual smoke.
209+
210+
!!! warning "torchrun relaunch loops"
211+
`_relaunch_under_torchrun` only re-execs when `--gpus N` is set *and*
212+
`TORCHELASTIC_RUN_ID` is absent. Running `torchrun ... -m genelab.cli
213+
train --gpus 4` (with `--gpus` left in argv inside the torchrun env)
214+
skips the relaunch because the env var is now present; this is intended.
215+
The infinite loop only happens if `--gpus` is forwarded *without*
216+
torchrun's env vars — typically when wrapping `genelab` inside another
217+
launcher.
218+
219+
!!! tip "Profile a short run first"
220+
Profiler traces grow fast. Start with `--prof --prof-active 3
221+
--prof-repeat 1 --max-iterations 10` to size the trace before running a
222+
full PPO schedule.
223+
224+
## See also
225+
226+
- [Play and Train](../cli/play-train.md)
227+
- [Profiling](../cli/profiling.md)
228+
- [Managers and MDP terms](managers.md)

0 commit comments

Comments
 (0)