Skip to content

Commit c71c858

Browse files
committed
feat(examples): add "recall the target" RNN memory showcase
Demonstrates when recurrence decisively beats a feedforward MLP — and, by being the only such task here, makes the honest point that POMDP stabilisation alone does not need an RNN. Task: each episode a random cart target is flashed in the observation for the first 5 steps then removed; the cart must be driven to the (now hidden) target. A memoryless MLP cannot recall it and drifts to centre; an LSTM stores it in its hidden state and tracks it. Measured post-cue tracking error: LSTM 0.013 vs MLP 0.78 (targets uniform in [-0.8, 0.8]). - mdp.py: reset_cart_target / cart_target_cue (masked after cue_steps) / cart_target_tracking_l2, with a per-env target buffer stashed on the env. - single/memory_env_cfg.py: reuses the base env, adds the masked cue obs, replaces cart-centring with target tracking, removes pushes, and shortens episodes to ~100 steps. - single/rnn_cfg.py: LSTM + MLP PPO configs. num_steps_per_env=100 so truncated BPTT spans the whole cue->reach episode (the key trainability detail; a 24-step window leaves the LSTM unable to learn the recall). - tasks.py: registers GeneLab-Inverted-Pendulum-Memory-{Rnn,Mlp}-v0 via a shared _MemoryTask base; single/double/skrl tasks unchanged. - docs + README: memory showcase section with the real numbers and the trainability note; EN/ZH in parity. - tests: registration, RNN-vs-MLP cfg (RNNModel vs MLPModel, num_steps_per_env), and the masked-cue/tracking-reward env wiring — all headless. Builds on the recurrent support in dev/recurrent. Trained + evaluated end-to-end on GPU (the numbers above); example tests + ruff green.
1 parent 67f3c35 commit c71c858

9 files changed

Lines changed: 455 additions & 3 deletions

File tree

‎docs/examples/inverted-pendulum.en.md‎

Lines changed: 40 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,8 @@ tests but exercises the full GeneLab train/play path.
1010
| `GeneLab-Inverted-Pendulum-v0` | Single pole balancing (rsl_rl PPO). |
1111
| `GeneLab-Double-Inverted-Pendulum-v0` | Two-link pole balancing (rsl_rl PPO). |
1212
| `GeneLab-Inverted-Pendulum-Skrl-v0` | Single pole balancing — same env, **skrl** PPO backend. |
13+
| `GeneLab-Inverted-Pendulum-Memory-Rnn-v0` | "Recall the target" — an **LSTM** policy on a task that needs memory. |
14+
| `GeneLab-Inverted-Pendulum-Memory-Mlp-v0` | Same task, plain **MLP** — the baseline that structurally can't remember. |
1315

1416
## Installing and listing
1517

@@ -41,14 +43,51 @@ genelab eval GeneLab-Inverted-Pendulum-Skrl-v0 logs/skrl/inverted_pendulum_skrl/
4143
> `genelab train` for the skrl task needs the optional `skrl` dependency
4244
> installed; registering and listing the task does not.
4345
46+
## Recurrent (RNN) memory showcase
47+
48+
`GeneLab-Inverted-Pendulum-Memory-Rnn-v0` shows **when recurrence actually helps**. It is a
49+
"recall the target" task: each episode a random cart target is **flashed in the observation
50+
for only the first 5 steps**, then removed. The flash is too brief to drive the cart there
51+
while it is visible, so the policy must *remember* the target and move to it afterwards. A
52+
memoryless MLP, once the cue is gone, has no way to recover the target and falls back to the
53+
centre; an LSTM stores it in its hidden state and drives the cart straight to it. (Velocities
54+
stay observed, so the only thing needing memory is the target.)
55+
56+
This is the task category where recurrence is decisive — unlike plain *balancing*, which a
57+
feedforward MLP solves even with velocities or loads hidden, because feedback stabilisation is
58+
robust to unobserved state. Memory, not partial observability alone, is what an RNN buys you.
59+
60+
Train both and compare the post-cue tracking error (mean distance from the hidden target):
61+
62+
```bash
63+
genelab train GeneLab-Inverted-Pendulum-Memory-Rnn-v0 --num_envs 2048 --max_iterations 400
64+
genelab train GeneLab-Inverted-Pendulum-Memory-Mlp-v0 --num_envs 2048 --max_iterations 400
65+
```
66+
67+
| Policy | Task | Post-cue tracking error (mean \|cart − target\|) |
68+
|---|---|---|
69+
| **LSTM (recurrent)** | `...-Memory-Rnn-v0` | **≈ 0.013** — reaches and holds the hidden target |
70+
| MLP (baseline) | `...-Memory-Mlp-v0` | ≈ 0.78 — can't recall it, sits near centre |
71+
72+
(Targets are uniform in `[-0.8, 0.8]`; a policy that ignores the target scores ≈ its spread.)
73+
74+
!!! note "Why it trains (and earlier memory attempts didn't)"
75+
76+
Recurrent PPO uses truncated BPTT of length `num_steps_per_env`. The memory cfg sets it to
77+
`100` — one whole ~100-step episode — so each BPTT window spans the full cue→reach
78+
dependency and the gradient can flow from the recall back to where the target was encoded.
79+
With the default 24-step window the LSTM never learns to associate the two.
80+
4481
## Code entry points
4582

4683
| File | Role |
4784
|---|---|
4885
| `tasks.py` | Registers robots, envs, and tasks. |
4986
| `single/env_cfg.py` | Single-pole manager-based env config. |
87+
| `single/memory_env_cfg.py` | The "recall the target" env (masked target cue + tracking reward). |
88+
| `single/rnn_cfg.py` | LSTM + MLP PPO configs for the memory task. |
5089
| `double/env_cfg.py` | Double-pole manager-based env config. |
51-
| `mdp.py` | Example-specific reward and termination helpers. |
90+
| `mdp.py` | Example-specific reward / termination / memory-task terms. |
5291

5392
## See also
5493

‎docs/examples/inverted-pendulum.zh.md‎

Lines changed: 36 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,8 @@
99
| `GeneLab-Inverted-Pendulum-v0` | 单杆平衡(rsl_rl PPO)。 |
1010
| `GeneLab-Double-Inverted-Pendulum-v0` | 双连杆平衡(rsl_rl PPO)。 |
1111
| `GeneLab-Inverted-Pendulum-Skrl-v0` | 单杆平衡 —— 同一个 env,**skrl** PPO 后端。 |
12+
| `GeneLab-Inverted-Pendulum-Memory-Rnn-v0` | "记住目标" —— 需要记忆的任务,用 **LSTM** 策略。 |
13+
| `GeneLab-Inverted-Pendulum-Memory-Mlp-v0` | 同一任务的普通 **MLP** 基线 —— 结构上无法记忆。 |
1214

1315
## 安装并列出
1416

@@ -39,14 +41,47 @@ genelab eval GeneLab-Inverted-Pendulum-Skrl-v0 logs/skrl/inverted_pendulum_skrl/
3941

4042
> skrl 任务的 `genelab train` 需要安装可选依赖 `skrl`;注册与列出任务则不需要。
4143
44+
## 循环网络(RNN)记忆 showcase
45+
46+
`GeneLab-Inverted-Pendulum-Memory-Rnn-v0` 展示**循环网络真正有用的场景**。这是一个"记住目标"
47+
任务:每个 episode 随机一个小车目标,但它只在**最初 5 步出现在观测里**,之后被抹掉。这一闪太短,
48+
来不及在可见时把车开过去,所以策略必须**记住**目标、之后再去。无记忆的 MLP 在线索消失后无法
49+
恢复目标,只能退回中点;而 LSTM 把它存进 hidden state,直接把车开到目标。(速度仍可观测,所以
50+
唯一需要记忆的就是目标。)
51+
52+
这才是循环网络决定性占优的任务类别 —— 与单纯的*平衡*不同:平衡即使隐藏速度或负载,前馈 MLP
53+
也能解,因为反馈镇定对未观测状态天生鲁棒。**RNN 买到的是"记忆",而不仅仅是"部分可观测"。**
54+
55+
训练两者,对比 post-cue 跟踪误差(到隐藏目标的平均距离):
56+
57+
```bash
58+
genelab train GeneLab-Inverted-Pendulum-Memory-Rnn-v0 --num_envs 2048 --max_iterations 400
59+
genelab train GeneLab-Inverted-Pendulum-Memory-Mlp-v0 --num_envs 2048 --max_iterations 400
60+
```
61+
62+
| 策略 | 任务 | post-cue 跟踪误差(平均 \|cart − target\|) |
63+
|---|---|---|
64+
| **LSTM(循环)** | `...-Memory-Rnn-v0` | **≈ 0.013** —— 够到并保持隐藏目标 |
65+
| MLP(基线) | `...-Memory-Mlp-v0` | ≈ 0.78 —— 回忆不出,卡在中点附近 |
66+
67+
(目标在 `[-0.8, 0.8]` 均匀采样;忽略目标的策略误差 ≈ 目标分布的幅度。)
68+
69+
!!! note "为什么这次训得动(之前的记忆尝试不行)"
70+
71+
循环 PPO 用长度为 `num_steps_per_env` 的 truncated BPTT。记忆 cfg 把它设成 `100`(整整一个
72+
~100 步 episode),于是每个 BPTT 窗口覆盖完整的"闪现→够到"依赖,梯度能从回忆一路传回编码线索
73+
那一步。用默认的 24 步窗口,LSTM 永远学不会把两者关联起来。
74+
4275
## 代码入口
4376

4477
| 文件 | 作用 |
4578
|---|---|
4679
| `tasks.py` | 注册 robot、env、task。 |
4780
| `single/env_cfg.py` | 单杆 manager-based env 配置。 |
81+
| `single/memory_env_cfg.py` | "记住目标" env(带遮罩的目标线索 + 跟踪奖励)。 |
82+
| `single/rnn_cfg.py` | 记忆任务的 LSTM + MLP PPO 配置。 |
4883
| `double/env_cfg.py` | 双杆 manager-based env 配置。 |
49-
| `mdp.py` | 示例自己的 reward 与 termination helper。 |
84+
| `mdp.py` | 示例自己的 reward / termination / 记忆任务项。 |
5085

5186
## 另见
5287

‎examples/inverted_pendulum/README.md‎

Lines changed: 45 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,13 @@
11
# Inverted pendulum examples
22

3-
GeneLab extension that ships two PPO tasks for the classical cart-pole problem on flat ground:
3+
GeneLab extension that ships PPO tasks for the classical cart-pole problem on flat ground:
44

55
- **`GeneLab-Inverted-Pendulum-v0`** — balance a single inverted pole on a cart.
66
- **`GeneLab-Double-Inverted-Pendulum-v0`** — stabilise two stacked inverted poles on a cart.
7+
- **`GeneLab-Inverted-Pendulum-Memory-Rnn-v0`** — "recall the target", a task that **needs
8+
memory**, solved by an LSTM. See [Recurrent showcase](#recurrent-rnn-memory-showcase).
9+
- **`GeneLab-Inverted-Pendulum-Memory-Mlp-v0`** — the same task with a plain MLP, kept as the
10+
baseline that structurally cannot remember.
711

812
The training stack mirrors `examples/unitree/`: `ManagerBasedRlEnv` over Genesis, rsl_rl PPO,
913
`BodyVelocitySensor` for the noisy pole rate observation, and the unified `genelab train` /
@@ -35,6 +39,8 @@ uv pip install -e examples/inverted_pendulum
3539
uv run genelab list tasks
3640
# -> GeneLab-Inverted-Pendulum-v0
3741
# -> GeneLab-Double-Inverted-Pendulum-v0
42+
# -> GeneLab-Inverted-Pendulum-Memory-Rnn-v0
43+
# -> GeneLab-Inverted-Pendulum-Memory-Mlp-v0
3844
```
3945

4046
### Single inverted pendulum
@@ -53,6 +59,44 @@ uv run genelab play GeneLab-Double-Inverted-Pendulum-v0 \
5359
--checkpoint logs/rsl_rl/double_inverted_pendulum_flat/<run>/model_300.pt
5460
```
5561

62+
## Recurrent (RNN) memory showcase
63+
64+
`GeneLab-Inverted-Pendulum-Memory-{Rnn,Mlp}-v0` demonstrate **when recurrence actually helps**
65+
— and, just as importantly, when it doesn't.
66+
67+
The task is "recall the target": each episode a random cart target is **flashed in the
68+
observation for only the first 5 steps**, then removed. The flash is too brief to drive the
69+
cart there while visible, so the policy must *remember* the target and move to it afterwards.
70+
Both tasks share the same env and training budget; the only difference is the policy network:
71+
72+
- the **MLP** (`...-Memory-Mlp-v0`) has no memory — once the cue is gone it can't recover the
73+
target and drifts back to the centre;
74+
- the **LSTM** (`...-Memory-Rnn-v0`) stores the target in its hidden state and drives straight
75+
to it.
76+
77+
```bash
78+
uv run genelab train GeneLab-Inverted-Pendulum-Memory-Rnn-v0 --num_envs 2048 --max_iterations 400
79+
uv run genelab train GeneLab-Inverted-Pendulum-Memory-Mlp-v0 --num_envs 2048 --max_iterations 400
80+
```
81+
82+
Measured post-cue tracking error (mean distance from the hidden target; targets uniform in
83+
`[-0.8, 0.8]`):
84+
85+
| Policy | Tracking error | Outcome |
86+
|---|---|---|
87+
| **LSTM (recurrent)** | **≈ 0.013** | reaches & holds the hidden target |
88+
| MLP (baseline) | ≈ 0.78 | can't recall it, sits near centre |
89+
90+
This is the right kind of task to reach for an RNN. Plain **balancing** is *not* — a feedforward
91+
MLP balances the cart-pole even with velocities (or a variable pole load) hidden, because
92+
feedback stabilisation is robust to unobserved state. What an RNN buys you is **memory**, not
93+
partial observability per se.
94+
95+
> **Why it trains.** Recurrent PPO uses truncated BPTT of length `num_steps_per_env`. The
96+
> memory cfg sets it to `100` (≈ one full episode) so each BPTT window spans the whole cue→reach
97+
> dependency — the gradient flows from the recall back to where the target was encoded. With the
98+
> default 24-step window the LSTM never learns to associate the two.
99+
56100
## Notes
57101

58102
- Only the cart slide joint is PD-controlled. The pole hinges default to `kp=0, kv=0` so the

‎examples/inverted_pendulum/src/genelab_inverted_pendulum/mdp.py‎

Lines changed: 57 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -119,3 +119,60 @@ def push_cart_by_setting_joint_velocity(
119119
target_vel = current[env_ids]
120120
target_pos = env.articulations["robot"].data.joint_pos[env_ids]
121121
env.articulations["robot"].write_joint_state(target_pos, target_vel, env_ids)
122+
123+
124+
# ---- "recall the target" memory task ---------------------------------------------
125+
# A per-episode cart target is shown only for the first ``cue_steps``, then hidden. With
126+
# frequent pushes knocking the cart around, a memoryless policy cannot return to the (now
127+
# unobserved) target — it must have remembered it. This is the task recurrence is for.
128+
129+
_CART_TARGET_ATTR = "_cart_memory_target"
130+
131+
132+
def _cart_target_buf(env: "ManagerBasedRlEnv") -> torch.Tensor:
133+
"""Per-env cart target, lazily allocated and stashed on the env object.
134+
135+
Created on first access (which is the episode-0 reset, so ``reset_cart_target`` fills it
136+
before any observation/reward reads it). A different ``num_envs`` always means a freshly
137+
built env, hence a fresh attribute — no resize path is needed.
138+
"""
139+
buf = getattr(env, _CART_TARGET_ATTR, None)
140+
if buf is None:
141+
buf = torch.zeros(env.num_envs, device=env.device)
142+
setattr(env, _CART_TARGET_ATTR, buf)
143+
return buf
144+
145+
146+
def reset_cart_target(
147+
env: "ManagerBasedRlEnv",
148+
env_ids: torch.Tensor | None,
149+
target_range: tuple[float, float] = (-1.0, 1.0),
150+
) -> None:
151+
"""Reset-mode: sample a fresh hidden cart target for each reset env."""
152+
if env_ids is None or env_ids.numel() == 0:
153+
return
154+
buf = _cart_target_buf(env)
155+
lo, hi = target_range
156+
buf[env_ids] = torch.empty(int(env_ids.numel()), device=env.device).uniform_(lo, hi)
157+
158+
159+
def cart_target_cue(env: "ManagerBasedRlEnv", cue_steps: int = 30) -> torch.Tensor:
160+
"""Observation term: the cart target during the first ``cue_steps``, then ``0``.
161+
162+
After the cue window the term is identical (zero) regardless of the true target, so a
163+
memoryless policy has no way to recover it — only a recurrent policy that stored it
164+
during the cue window can keep tracking it.
165+
"""
166+
buf = _cart_target_buf(env)
167+
active = (env.episode_length_buf < cue_steps).to(buf.dtype)
168+
return (buf * active).unsqueeze(-1)
169+
170+
171+
def cart_target_tracking_l2(
172+
env: "ManagerBasedRlEnv", joint_name: str = "cart_slide"
173+
) -> torch.Tensor:
174+
"""Squared distance of the cart from its (hidden) target — the thing to minimise."""
175+
buf = _cart_target_buf(env)
176+
idx = _joint_index(env, joint_name)
177+
x = env.articulations["robot"].data.joint_pos[:, idx]
178+
return (x - buf).square()

‎examples/inverted_pendulum/src/genelab_inverted_pendulum/single/__init__.py‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,12 @@
1111
POLE_LINK,
1212
)
1313
from genelab_inverted_pendulum.single.env_cfg import inverted_pendulum_env_cfg
14+
from genelab_inverted_pendulum.single.memory_env_cfg import inverted_pendulum_memory_env_cfg
1415
from genelab_inverted_pendulum.single.ppo_cfg import inverted_pendulum_ppo_runner_cfg
16+
from genelab_inverted_pendulum.single.rnn_cfg import (
17+
inverted_pendulum_memory_mlp_runner_cfg,
18+
inverted_pendulum_memory_rnn_runner_cfg,
19+
)
1520
from genelab_inverted_pendulum.single.skrl_cfg import inverted_pendulum_skrl_agent_cfg
1621
from genelab_inverted_pendulum.single.robot import (
1722
InvertedPendulumRobotCfg,
@@ -30,6 +35,9 @@
3035
"POLE_LINK",
3136
"get_inverted_pendulum_robot_cfg",
3237
"inverted_pendulum_env_cfg",
38+
"inverted_pendulum_memory_env_cfg",
39+
"inverted_pendulum_memory_mlp_runner_cfg",
40+
"inverted_pendulum_memory_rnn_runner_cfg",
3341
"inverted_pendulum_ppo_runner_cfg",
3442
"inverted_pendulum_skrl_agent_cfg",
3543
]
Lines changed: 56 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,56 @@
1+
"""A "recall the target" cart-pole — a task that genuinely *requires* memory.
2+
3+
Unlike the POMDP balancing variants (which a memoryless policy still solves, because
4+
feedback stabilisation is robust to unobserved state), this task hides information that
5+
**cannot be recovered by feedback**: a per-episode cart target shown only for the first
6+
``cue_steps`` (a brief flash) and then removed from the observation. The cue is too short to
7+
drive the cart to the target while it is visible, so the cart is still near the centre when
8+
the target vanishes — the policy must have *remembered* it to move there afterwards. A
9+
memoryless MLP, once the cue is gone, has no idea where the target is and falls back to a
10+
fixed point (the mean target, ~centre); a recurrent policy stores the target during the
11+
flash and drives the cart to it. The cart-to-target gap between the two is the demonstration.
12+
13+
The cue is deliberately brief and there are no disturbances, so the task is *learnable*
14+
(constant pushes made earlier versions hopeless) while still being unsolvable without
15+
memory. Velocities stay observed — the only thing needing memory is the target, isolating
16+
the recurrence benefit cleanly.
17+
"""
18+
19+
from genelab.envs.manager_based_rl_env import ManagerBasedRlEnvCfg
20+
from genelab.managers import EventTermCfg, ObservationTermCfg, RewardTermCfg
21+
22+
from genelab_inverted_pendulum import mdp as ip_mdp
23+
from genelab_inverted_pendulum.single.env_cfg import inverted_pendulum_env_cfg
24+
25+
_CUE_STEPS = 5 # brief flash: long enough to read the target, too short to drive there
26+
27+
28+
def inverted_pendulum_memory_env_cfg(play: bool = False) -> ManagerBasedRlEnvCfg:
29+
"""Cart-pole that must drive the cart to a target flashed only at the episode start."""
30+
cfg = inverted_pendulum_env_cfg(play=play)
31+
# Short episodes (~100 control steps) keep the cue→reach memory dependency inside one
32+
# truncated-BPTT window, so the recurrent gradient can actually learn it (see rnn_cfg).
33+
cfg.episode_length_s = 1.0
34+
35+
# Add the (masked) target cue to both observation groups.
36+
cfg.observations_cfg["policy"].terms["target_cue"] = ObservationTermCfg(
37+
func=ip_mdp.cart_target_cue, params={"cue_steps": _CUE_STEPS}
38+
)
39+
cfg.observations_cfg["critic"].terms["target_cue"] = ObservationTermCfg(
40+
func=ip_mdp.cart_target_cue, params={"cue_steps": _CUE_STEPS}
41+
)
42+
43+
# Track the hidden target instead of holding the cart at the origin; weight it heavily so
44+
# the policy clearly prioritises reaching the remembered target over sitting at centre.
45+
cfg.rewards_cfg.pop("cart_position", None)
46+
cfg.rewards_cfg["cart_target"] = RewardTermCfg(func=ip_mdp.cart_target_tracking_l2, weight=-4.0)
47+
48+
# Sample a fresh hidden target each reset. No pushes — the brief cue already prevents the
49+
# cart from reaching (and parking at) the target while it is still visible.
50+
cfg.events_cfg["reset_cart_target"] = EventTermCfg(
51+
mode="reset",
52+
func=ip_mdp.reset_cart_target,
53+
params={"target_range": (-0.8, 0.8)}, # reachable within a 1 s episode
54+
)
55+
cfg.events_cfg.pop("push_cart", None)
56+
return cfg

0 commit comments

Comments
 (0)