From cd112dd2e642332b9ccd74ac91f6a6754f1b079c Mon Sep 17 00:00:00 2001 From: KraHsu Date: Thu, 14 May 2026 14:20:50 +0800 Subject: [PATCH] Add inverted-pendulum PPO example extension Register `GeneLab-Inverted-Pendulum-v0` and `GeneLab-Double-Inverted-Pendulum-v0` under examples/inverted_pendulum/. Single PD-controlled cart drives one or two underactuated hinge poles through rsl_rl PPO on the existing ManagerBasedRlEnv + BodyVelocitySensor stack. Play mode enables `scene.mouse_interaction` for spring-force disturbance, with matching English/Chinese walkthrough pages and a registration smoke test. --- docs/examples/inverted-pendulum.en.md | 81 ++++++++++ docs/examples/inverted-pendulum.zh.md | 77 ++++++++++ docs/examples/overview.en.md | 10 ++ docs/examples/overview.zh.md | 10 ++ examples/README.md | 32 +++- examples/inverted_pendulum/README.md | 63 ++++++++ .../assets/double_inverted_pendulum.xml | 30 ++++ .../assets/inverted_pendulum.xml | 25 +++ examples/inverted_pendulum/pyproject.toml | 21 +++ .../src/genelab_inverted_pendulum/__init__.py | 5 + .../double/__init__.py | 37 +++++ .../double/constants.py | 32 ++++ .../double/env_cfg.py | 141 +++++++++++++++++ .../double/ppo_cfg.py | 44 ++++++ .../genelab_inverted_pendulum/double/robot.py | 37 +++++ .../src/genelab_inverted_pendulum/mdp.py | 127 +++++++++++++++ .../single/__init__.py | 33 ++++ .../single/constants.py | 28 ++++ .../single/env_cfg.py | 121 +++++++++++++++ .../single/ppo_cfg.py | 44 ++++++ .../genelab_inverted_pendulum/single/robot.py | 37 +++++ .../src/genelab_inverted_pendulum/tasks.py | 144 ++++++++++++++++++ mkdocs.yml | 2 + pyproject.toml | 2 +- tests/test_inverted_pendulum_examples.py | 16 ++ 25 files changed, 1192 insertions(+), 7 deletions(-) create mode 100644 docs/examples/inverted-pendulum.en.md create mode 100644 docs/examples/inverted-pendulum.zh.md create mode 100644 examples/inverted_pendulum/README.md create mode 100644 examples/inverted_pendulum/assets/double_inverted_pendulum.xml create mode 100644 examples/inverted_pendulum/assets/inverted_pendulum.xml create mode 100644 examples/inverted_pendulum/pyproject.toml create mode 100644 examples/inverted_pendulum/src/genelab_inverted_pendulum/__init__.py create mode 100644 examples/inverted_pendulum/src/genelab_inverted_pendulum/double/__init__.py create mode 100644 examples/inverted_pendulum/src/genelab_inverted_pendulum/double/constants.py create mode 100644 examples/inverted_pendulum/src/genelab_inverted_pendulum/double/env_cfg.py create mode 100644 examples/inverted_pendulum/src/genelab_inverted_pendulum/double/ppo_cfg.py create mode 100644 examples/inverted_pendulum/src/genelab_inverted_pendulum/double/robot.py create mode 100644 examples/inverted_pendulum/src/genelab_inverted_pendulum/mdp.py create mode 100644 examples/inverted_pendulum/src/genelab_inverted_pendulum/single/__init__.py create mode 100644 examples/inverted_pendulum/src/genelab_inverted_pendulum/single/constants.py create mode 100644 examples/inverted_pendulum/src/genelab_inverted_pendulum/single/env_cfg.py create mode 100644 examples/inverted_pendulum/src/genelab_inverted_pendulum/single/ppo_cfg.py create mode 100644 examples/inverted_pendulum/src/genelab_inverted_pendulum/single/robot.py create mode 100644 examples/inverted_pendulum/src/genelab_inverted_pendulum/tasks.py create mode 100644 tests/test_inverted_pendulum_examples.py diff --git a/docs/examples/inverted-pendulum.en.md b/docs/examples/inverted-pendulum.en.md new file mode 100644 index 00000000..d0790d2a --- /dev/null +++ b/docs/examples/inverted-pendulum.en.md @@ -0,0 +1,81 @@ +# Inverted Pendulum + +`examples/inverted_pendulum/` ships two PPO tasks for the classical cart-pole problem on flat +ground. The extension mirrors `examples/unitree/`: `ManagerBasedRlEnv` over Genesis, rsl_rl PPO, +a `BodyVelocitySensor` on the pole link, and the unified `genelab train` / `genelab play` CLI. + +## Tasks + +| Task id | Problem | +|---------|---------| +| `GeneLab-Inverted-Pendulum-v0` | Single inverted pole on a cart. | +| `GeneLab-Double-Inverted-Pendulum-v0` | Two stacked inverted poles on a cart. | + +## Installation + +The extension depends on the `rl` extra (rsl_rl). Pick the `torch-*` extra that matches the +hardware. + +```bash +uv sync --extra rl --extra torch-cu128 +uv pip install -e examples/inverted_pendulum + +uv run genelab list tasks +# -> GeneLab-Inverted-Pendulum-v0 +# -> GeneLab-Double-Inverted-Pendulum-v0 +``` + +## Single inverted pendulum + +```bash +uv run genelab train GeneLab-Inverted-Pendulum-v0 \ + --num-envs 4096 --max-iterations 150 + +uv run genelab play GeneLab-Inverted-Pendulum-v0 \ + --checkpoint logs/rsl_rl/inverted_pendulum_flat//model_150.pt --vis +``` + +`--checkpoint` makes `play` route through the RL runner with `--agent trained` by default. + +## Double inverted pendulum + +```bash +uv run genelab train GeneLab-Double-Inverted-Pendulum-v0 \ + --num-envs 4096 --max-iterations 300 + +uv run genelab play GeneLab-Double-Inverted-Pendulum-v0 \ + --checkpoint logs/rsl_rl/double_inverted_pendulum_flat//model_300.pt --vis +``` + +## Sensor and underactuation + +Only the cart slide joint is PD-controlled. The pole hinges default to `kp=0, kv=0` so the +pendulum stays underactuated and the policy must learn balance through cart motion alone. A +`BodyVelocitySensor` attached to the top pole supplies a noisy angular-velocity observation +(corrupted with `Unoise` in the policy group, clean in the critic group). + +## Interactive disturbance + +Play mode launches a single environment (`num_envs=1`) and enables Genesis' +`MouseInteractionPlugin`. Left-click on the cart or pole and drag — a spring force pulls the +clicked link toward the cursor while the policy keeps balancing. Scroll wheel rotates the drag +plane around the surface normal. Release the button to remove the force. + +!!! tip "Smoke-test budget" + A 5–10 iteration run with `--num-envs 64 --max-iterations 5` is enough to validate wiring + end-to-end. The reward signal will still be noisy at that scale; convergence requires the + 150 / 300 iteration budgets above. + +## Logs + +Both tasks write to `logs/rsl_rl//_/` like the Unitree examples: + +- `params/env.json` and `params/agent.json` — frozen configs at run time. +- `model_.pt` — checkpoints saved every `save_interval` iterations. +- TensorBoard event files alongside the checkpoints. + +## See also + +- [Unitree G1 quickstart](../getting-started/quickstart.md#unitree-g1) +- [Sensors](../concepts/sensors.md) +- [Play and Train CLI](../cli/play-train.md) diff --git a/docs/examples/inverted-pendulum.zh.md b/docs/examples/inverted-pendulum.zh.md new file mode 100644 index 00000000..4d8fdaa5 --- /dev/null +++ b/docs/examples/inverted-pendulum.zh.md @@ -0,0 +1,77 @@ +# 倒立摆 + +`examples/inverted_pendulum/` 提供两个在平面上的 cart-pole 经典控制 PPO 任务。整条训练栈与 +`examples/unitree/` 对齐：基于 Genesis 的 `ManagerBasedRlEnv`、rsl_rl PPO、挂在杆上的 +`BodyVelocitySensor`，以及统一的 `genelab train` / `genelab play` CLI。 + +## 任务列表 + +| Task id | 问题 | +|---------|------| +| `GeneLab-Inverted-Pendulum-v0` | 小车 + 单杆的倒立摆。 | +| `GeneLab-Double-Inverted-Pendulum-v0` | 小车 + 串联双杆的倒立摆。 | + +## 安装 + +扩展依赖 `rl` extra（rsl_rl）。`torch-*` extra 按硬件挑选。 + +```bash +uv sync --extra rl --extra torch-cu128 +uv pip install -e examples/inverted_pendulum + +uv run genelab list tasks +# -> GeneLab-Inverted-Pendulum-v0 +# -> GeneLab-Double-Inverted-Pendulum-v0 +``` + +## 单倒立摆 + +```bash +uv run genelab train GeneLab-Inverted-Pendulum-v0 \ + --num-envs 4096 --max-iterations 150 + +uv run genelab play GeneLab-Inverted-Pendulum-v0 \ + --checkpoint logs/rsl_rl/inverted_pendulum_flat//model_150.pt --vis +``` + +传入 `--checkpoint` 会让 `play` 自动经过 RL runner，并默认使用 `--agent trained`。 + +## 双倒立摆 + +```bash +uv run genelab train GeneLab-Double-Inverted-Pendulum-v0 \ + --num-envs 4096 --max-iterations 300 + +uv run genelab play GeneLab-Double-Inverted-Pendulum-v0 \ + --checkpoint logs/rsl_rl/double_inverted_pendulum_flat//model_300.pt --vis +``` + +## 传感器与欠驱动 + +只有小车的 slide 关节通过 PD 控制。两个 pole hinge 默认 `kp=0, kv=0`，保证整体处于欠驱动状态， +策略必须通过小车水平运动间接稳定杆。顶端 pole 上挂载的 `BodyVelocitySensor` 给出一路带噪声的 +角速度观测：policy 观测组用 `Unoise` 做 corruption，critic 观测组直接读取干净值。 + +## 交互式扰动 + +`play` 默认只开 1 个环境（`num_envs=1`），并启用 Genesis 的 `MouseInteractionPlugin`。 +左键点击 cart 或 pole 并拖动，会有一根弹簧把所点击的 link 拉向光标位置；策略仍然在背后试图 +保持平衡。滚轮可绕表面法线旋转拖拽平面，松开左键即移除外力。 + +!!! tip "Smoke-test 预算" + 使用 `--num-envs 64 --max-iterations 5` 跑 5–10 次迭代足以验证整条链路。此时 reward 信号 + 仍非常嘈杂，真正收敛需要上面给出的 150 / 300 次迭代预算。 + +## 日志 + +两个任务都把日志写到 `logs/rsl_rl//_/`，结构与 Unitree 示例一致： + +- `params/env.json` 与 `params/agent.json` —— 运行时冻结的配置快照。 +- `model_.pt` —— 按 `save_interval` 保存的 checkpoint。 +- 同目录下的 TensorBoard 事件文件。 + +## See also + +- [Unitree G1 快速开始](../getting-started/quickstart.md#unitree-g1) +- [传感器](../concepts/sensors.md) +- [play 与 train CLI](../cli/play-train.md) diff --git a/docs/examples/overview.en.md b/docs/examples/overview.en.md index 742f1143..d341e4c6 100644 --- a/docs/examples/overview.en.md +++ b/docs/examples/overview.en.md @@ -3,6 +3,16 @@ The repository ships several reference extensions under `examples/`. They double as integration tests for the CLI and registry. +## inverted_pendulum + +Two PPO cart-pole tasks built on the same `ManagerBasedRlEnv` + rsl_rl stack as the Unitree +example, sized to fit in a laptop training budget: + +- **`GeneLab-Inverted-Pendulum-v0`** — single inverted pole on a cart. +- **`GeneLab-Double-Inverted-Pendulum-v0`** — two stacked inverted poles on a cart. + +Source at `examples/inverted_pendulum/`; walkthrough at [Inverted Pendulum](inverted-pendulum.md). + ## genelab_examples The canonical in-tree extension, wiring two tasks: diff --git a/docs/examples/overview.zh.md b/docs/examples/overview.zh.md index b3ebc2b9..0ce2e74e 100644 --- a/docs/examples/overview.zh.md +++ b/docs/examples/overview.zh.md @@ -2,6 +2,16 @@ 仓库在 `examples/` 下提供数个参考扩展，同时也是 CLI 与注册表的集成测试。 +## inverted_pendulum + +两个 PPO cart-pole 任务，训练栈与 Unitree 示例相同（`ManagerBasedRlEnv` + rsl_rl），训练预算 +控制在单机能跑完的量级： + +- **`GeneLab-Inverted-Pendulum-v0`** —— 小车 + 单杆倒立摆。 +- **`GeneLab-Double-Inverted-Pendulum-v0`** —— 小车 + 串联双杆倒立摆。 + +源码位于 `examples/inverted_pendulum/`；完整流程见 [倒立摆](inverted-pendulum.md)。 + ## genelab_examples 仓库内的标准扩展，接通两个任务： diff --git a/examples/README.md b/examples/README.md index 27976da2..86cfffb7 100644 --- a/examples/README.md +++ b/examples/README.md @@ -5,26 +5,36 @@ ship built-in tasks; example tasks are loaded like any other external project. ## Available Examples +- [Inverted Pendulum](inverted_pendulum/README.md): trainable single- and double-inverted-pendulum tasks that initialize Genesis while fully exercising `train` + `play`. - [GeneLab Example Extension](genelab_examples/README.md): one Python project that registers the Rubik's cube and Wuji hand tasks. - [External Project](external_project/README.md): minimal standalone Python package that extends GeneLab without editing `src/genelab/`. -The example extension registers these task IDs: +The bundled examples register these task IDs: +- `GeneLab-Inverted-Pendulum-v0` +- `GeneLab-Double-Inverted-Pendulum-v0` - `GeneLab-Rubiks-Play-v0` - `GeneLab-Wuji-Hand-Playback-v0` -List example tasks from the repository root without installing the example package: +List inverted-pendulum tasks from the repository root without installing the package: + +```bash +PYTHONPATH=examples/inverted_pendulum/src uv run genelab --import genelab_inverted_pendulum.tasks list tasks +``` + +List the Genesis demo tasks the same way: ```bash PYTHONPATH=examples/genelab_examples/src uv run genelab --import genelab_examples.tasks list tasks ``` -Install the example extension once if you want `uv run genelab list tasks` to load it through entry +Install an example extension once if you want `uv run genelab list tasks` to load it through entry points: ```bash +uv pip install -e examples/inverted_pendulum uv pip install -e examples/genelab_examples uv run genelab list tasks ``` @@ -43,13 +53,15 @@ override keys are converted to underscores. Examples: ```bash +PYTHONPATH=examples/inverted_pendulum/src uv run genelab --import genelab_inverted_pendulum.tasks train GeneLab-Inverted-Pendulum-v0 --num-envs 4096 --max-iterations 150 +PYTHONPATH=examples/inverted_pendulum/src uv run genelab --import genelab_inverted_pendulum.tasks play GeneLab-Inverted-Pendulum-v0 --checkpoint logs/rsl_rl/inverted_pendulum_flat//model_150.pt --vis PYTHONPATH=examples/genelab_examples/src uv run genelab --import genelab_examples.tasks play GeneLab-Rubiks-Play-v0 --steps 5 --env.robot.cubie_size 0.04 --env.robot.gap 0.002 PYTHONPATH=examples/genelab_examples/src uv run genelab --import genelab_examples.tasks play GeneLab-Rubiks-Play-v0 --env.robot.welded true PYTHONPATH=examples/genelab_examples/src uv run genelab --import genelab_examples.tasks play GeneLab-Wuji-Hand-Playback-v0 --env.reset_interval 0 ``` -`train` validates the task id and configuration path but currently reports that training is not -implemented: +`train` is implemented for the inverted-pendulum tasks. The Rubik's cube and Wuji hand demo tasks are +play-only and report that training is not implemented: ```bash PYTHONPATH=examples/genelab_examples/src uv run genelab --import genelab_examples.tasks train GeneLab-Rubiks-Play-v0 @@ -57,7 +69,15 @@ PYTHONPATH=examples/genelab_examples/src uv run genelab --import genelab_example ## Smoke Tests -Run short headless smoke tests after the Genesis assets and cache have initialized: +Run the inverted-pendulum smoke tests first; a tiny rsl_rl run exercises the full Genesis + +PPO pipeline: + +```bash +PYTHONPATH=examples/inverted_pendulum/src uv run genelab --import genelab_inverted_pendulum.tasks train GeneLab-Inverted-Pendulum-v0 --num-envs 64 --max-iterations 5 +PYTHONPATH=examples/inverted_pendulum/src uv run genelab --import genelab_inverted_pendulum.tasks train GeneLab-Double-Inverted-Pendulum-v0 --num-envs 64 --max-iterations 5 +``` + +Then run short headless smoke tests after the Genesis assets and cache have initialized: ```bash PYTHONPATH=examples/genelab_examples/src uv run genelab --import genelab_examples.tasks play GeneLab-Rubiks-Play-v0 --steps 5 diff --git a/examples/inverted_pendulum/README.md b/examples/inverted_pendulum/README.md new file mode 100644 index 00000000..d7b6463e --- /dev/null +++ b/examples/inverted_pendulum/README.md @@ -0,0 +1,63 @@ +# Inverted pendulum examples + +GeneLab extension that ships two PPO tasks for the classical cart-pole problem on flat ground: + +- **`GeneLab-Inverted-Pendulum-v0`** — balance a single inverted pole on a cart. +- **`GeneLab-Double-Inverted-Pendulum-v0`** — stabilise two stacked inverted poles on a cart. + +The training stack mirrors `examples/unitree/`: `ManagerBasedRlEnv` over Genesis, rsl_rl PPO, +`BodyVelocitySensor` for the noisy pole rate observation, and the unified `genelab train` / +`genelab play` CLI. + +## Layout + +``` +examples/inverted_pendulum/ +├── pyproject.toml +├── README.md +├── assets/ # cart-pole MJCFs +│ ├── inverted_pendulum.xml +│ └── double_inverted_pendulum.xml +└── src/genelab_inverted_pendulum/ + ├── tasks.py # registers both tasks + ├── mdp.py # cart-pole-specific reward / termination / event terms + ├── single/ # single-pendulum config (robot + env + PPO) + └── double/ # double-pendulum config (robot + env + PPO) +``` + +## Quickstart + +```bash +# From the GeneLab repo root +uv sync --extra rl --extra torch-cu128 # pick whichever torch flavor fits your GPU +uv pip install -e examples/inverted_pendulum + +uv run genelab list tasks +# -> GeneLab-Inverted-Pendulum-v0 +# -> GeneLab-Double-Inverted-Pendulum-v0 +``` + +### Single inverted pendulum + +```bash +uv run genelab train GeneLab-Inverted-Pendulum-v0 --num_envs 4096 --max_iterations 150 +uv run genelab play GeneLab-Inverted-Pendulum-v0 \ + --checkpoint logs/rsl_rl/inverted_pendulum_flat//model_150.pt +``` + +### Double inverted pendulum + +```bash +uv run genelab train GeneLab-Double-Inverted-Pendulum-v0 --num_envs 4096 --max_iterations 300 +uv run genelab play GeneLab-Double-Inverted-Pendulum-v0 \ + --checkpoint logs/rsl_rl/double_inverted_pendulum_flat//model_300.pt +``` + +## Notes + +- Only the cart slide joint is PD-controlled. The pole hinges default to `kp=0, kv=0` so the + pendulum stays underactuated and the policy must learn balance via cart motion alone. +- The observation group corrupts joint position, joint velocity, and pole angular velocity with + `Unoise`. The critic group sees the same features without corruption. +- Logs land under `logs/rsl_rl//_/` exactly like the Unitree examples, + with `params/env.json`, `params/agent.json`, and `model_.pt` files. diff --git a/examples/inverted_pendulum/assets/double_inverted_pendulum.xml b/examples/inverted_pendulum/assets/double_inverted_pendulum.xml new file mode 100644 index 00000000..01e5d1f8 --- /dev/null +++ b/examples/inverted_pendulum/assets/double_inverted_pendulum.xml @@ -0,0 +1,30 @@ + + diff --git a/examples/inverted_pendulum/assets/inverted_pendulum.xml b/examples/inverted_pendulum/assets/inverted_pendulum.xml new file mode 100644 index 00000000..feeef6f0 --- /dev/null +++ b/examples/inverted_pendulum/assets/inverted_pendulum.xml @@ -0,0 +1,25 @@ + + diff --git a/examples/inverted_pendulum/pyproject.toml b/examples/inverted_pendulum/pyproject.toml new file mode 100644 index 00000000..c4b542bd --- /dev/null +++ b/examples/inverted_pendulum/pyproject.toml @@ -0,0 +1,21 @@ +[project] +name = "genelab-inverted-pendulum" +version = "0.1.0" +description = "Inverted-pendulum PPO tasks for GeneLab." +readme = "README.md" +requires-python = ">=3.12" +dependencies = ["genelab[rl]"] + +[project.entry-points."genelab.extensions"] +genelab_inverted_pendulum = "genelab_inverted_pendulum.tasks:register" + +[build-system] +requires = ["hatchling"] +build-backend = "hatchling.build" + +[tool.hatch.build.targets.wheel] +packages = ["src/genelab_inverted_pendulum"] +include = ["assets/**"] + +[tool.uv.sources] +genelab = { path = "../..", editable = true } diff --git a/examples/inverted_pendulum/src/genelab_inverted_pendulum/__init__.py b/examples/inverted_pendulum/src/genelab_inverted_pendulum/__init__.py new file mode 100644 index 00000000..7f2f22c2 --- /dev/null +++ b/examples/inverted_pendulum/src/genelab_inverted_pendulum/__init__.py @@ -0,0 +1,5 @@ +"""Inverted-pendulum RL tasks for GeneLab.""" + +__all__ = ["__version__"] + +__version__ = "0.1.0" diff --git a/examples/inverted_pendulum/src/genelab_inverted_pendulum/double/__init__.py b/examples/inverted_pendulum/src/genelab_inverted_pendulum/double/__init__.py new file mode 100644 index 00000000..528e3d78 --- /dev/null +++ b/examples/inverted_pendulum/src/genelab_inverted_pendulum/double/__init__.py @@ -0,0 +1,37 @@ +"""Double inverted-pendulum robot, env, and PPO configuration.""" + +from genelab_inverted_pendulum.double.constants import ( + CART_ACTION_SCALE, + CART_JOINT, + DEFAULT_JOINT_POS, + DOUBLE_INVERTED_PENDULUM_MJCF, + JOINT_KP, + JOINT_KV, + POLE_1_JOINT, + POLE_2_JOINT, + POLE_2_LINK, + POLE_HINGE_JOINTS, +) +from genelab_inverted_pendulum.double.env_cfg import double_inverted_pendulum_env_cfg +from genelab_inverted_pendulum.double.ppo_cfg import double_inverted_pendulum_ppo_runner_cfg +from genelab_inverted_pendulum.double.robot import ( + DoubleInvertedPendulumRobotCfg, + get_double_inverted_pendulum_robot_cfg, +) + +__all__ = [ + "CART_ACTION_SCALE", + "CART_JOINT", + "DEFAULT_JOINT_POS", + "DOUBLE_INVERTED_PENDULUM_MJCF", + "DoubleInvertedPendulumRobotCfg", + "JOINT_KP", + "JOINT_KV", + "POLE_1_JOINT", + "POLE_2_JOINT", + "POLE_2_LINK", + "POLE_HINGE_JOINTS", + "double_inverted_pendulum_env_cfg", + "double_inverted_pendulum_ppo_runner_cfg", + "get_double_inverted_pendulum_robot_cfg", +] diff --git a/examples/inverted_pendulum/src/genelab_inverted_pendulum/double/constants.py b/examples/inverted_pendulum/src/genelab_inverted_pendulum/double/constants.py new file mode 100644 index 00000000..3ba049bc --- /dev/null +++ b/examples/inverted_pendulum/src/genelab_inverted_pendulum/double/constants.py @@ -0,0 +1,32 @@ +"""Double inverted-pendulum constants: MJCF path, joint gains, action scale.""" + +from pathlib import Path +from typing import Final + +# .../double/constants.py → .../assets/double_inverted_pendulum.xml +DOUBLE_INVERTED_PENDULUM_MJCF: Final = ( + Path(__file__).resolve().parents[3] / "assets" / "double_inverted_pendulum.xml" +).resolve() + +CART_JOINT: Final = "cart_slide" +POLE_1_JOINT: Final = "pole_1_hinge" +POLE_2_JOINT: Final = "pole_2_hinge" +POLE_HINGE_JOINTS: Final = (POLE_1_JOINT, POLE_2_JOINT) +POLE_2_LINK: Final = "pole_2" + +JOINT_KP: Final[dict[str, float]] = {CART_JOINT: 100.0} +JOINT_KV: Final[dict[str, float]] = {CART_JOINT: 10.0} + +CART_ACTION_SCALE: Final[dict[str, float]] = {CART_JOINT: 1.0} + +DEFAULT_JOINT_POS: Final[dict[str, float]] = { + CART_JOINT: 0.0, + POLE_1_JOINT: 0.0, + POLE_2_JOINT: 0.0, +} + +CART_POSITION_LIMIT: Final = 2.4 +POLE_1_ANGLE_LIMIT: Final = 0.5 +POLE_2_ANGLE_LIMIT: Final = 0.7 + +INIT_BASE_HEIGHT: Final = 0.12 diff --git a/examples/inverted_pendulum/src/genelab_inverted_pendulum/double/env_cfg.py b/examples/inverted_pendulum/src/genelab_inverted_pendulum/double/env_cfg.py new file mode 100644 index 00000000..274150f4 --- /dev/null +++ b/examples/inverted_pendulum/src/genelab_inverted_pendulum/double/env_cfg.py @@ -0,0 +1,141 @@ +"""Flat-ground double inverted-pendulum env config (manager-based + Genesis).""" + +from genelab import mdp +from genelab.configs import SceneCfg +from genelab.envs.manager_based_rl_env import ManagerBasedRlEnvCfg +from genelab.managers import ( + EventTermCfg, + ObservationGroupCfg, + ObservationTermCfg, + RewardTermCfg, + TerminationTermCfg, +) +from genelab.mdp.actions.joint_position import JointPositionActionCfg +from genelab.mdp.noise import Unoise +from genelab.sensor import BodyVelocitySensorCfg + +from genelab_inverted_pendulum import mdp as ip_mdp +from genelab_inverted_pendulum.double.constants import ( + CART_ACTION_SCALE, + CART_JOINT, + CART_POSITION_LIMIT, + POLE_1_ANGLE_LIMIT, + POLE_1_JOINT, + POLE_2_ANGLE_LIMIT, + POLE_2_JOINT, + POLE_2_LINK, + POLE_HINGE_JOINTS, +) +from genelab_inverted_pendulum.double.robot import get_double_inverted_pendulum_robot_cfg + + +def _obs_terms() -> dict[str, ObservationTermCfg]: + return { + "joint_pos": ObservationTermCfg( + func=mdp.joint_pos_rel, + noise=Unoise(-0.005, 0.005), + ), + "joint_vel": ObservationTermCfg( + func=mdp.joint_vel_rel, + scale=0.1, + noise=Unoise(-0.05, 0.05), + ), + "pole_tip_ang_vel": ObservationTermCfg( + func=mdp.sensor_data, + params={"sensor_name": "pole_tip_ang_vel"}, + scale=0.2, + noise=Unoise(-0.1, 0.1), + ), + "actions": ObservationTermCfg(func=mdp.last_action), + } + + +def double_inverted_pendulum_env_cfg(play: bool = False) -> ManagerBasedRlEnvCfg: + """Flat-ground double-inverted-pendulum env config.""" + robot_entity_cfg = get_double_inverted_pendulum_robot_cfg().to_entity_cfg() + + cfg = ManagerBasedRlEnvCfg( + scene=SceneCfg( + num_envs=4096 if not play else 1, + dt=0.005, + substeps=1, + env_spacing=(3.0, 3.0), + vis=play, + mouse_interaction=play, + sensors=( + BodyVelocitySensorCfg( + name="pole_tip_ang_vel", + link_name=POLE_2_LINK, + measure="ang_vel", + ), + ), + ), + decimation=2, + episode_length_s=12.0, + device="cuda", + robot=robot_entity_cfg, + actions_cfg={ + "cart": JointPositionActionCfg( + asset_name="robot", + joint_names=(CART_JOINT,), + scale=dict(CART_ACTION_SCALE), + use_default_offset=True, + ), + }, + observations_cfg={ + "policy": ObservationGroupCfg(enable_corruption=True, terms=_obs_terms()), + "critic": ObservationGroupCfg(enable_corruption=False, terms=_obs_terms()), + }, + rewards_cfg={ + "alive": RewardTermCfg(func=ip_mdp.alive_bonus, weight=1.0), + "poles_upright": RewardTermCfg( + func=ip_mdp.double_pole_upright, + weight=4.0, + params={"joint_names": POLE_HINGE_JOINTS}, + ), + "poles_aligned": RewardTermCfg( + func=ip_mdp.double_pole_alignment, + weight=-0.5, + params={"joint_names": POLE_HINGE_JOINTS}, + ), + "cart_position": RewardTermCfg(func=ip_mdp.cart_position_l2, weight=-0.05), + "cart_velocity": RewardTermCfg(func=ip_mdp.cart_velocity_l2, weight=-0.005), + "poles_velocity": RewardTermCfg( + func=ip_mdp.double_pole_velocity_l2, + weight=-0.005, + params={"joint_names": POLE_HINGE_JOINTS}, + ), + "action_rate": RewardTermCfg(func=mdp.action_rate_l2, weight=-0.005), + }, + terminations_cfg={ + "time_out": TerminationTermCfg(func=mdp.time_out, time_out=True), + "poles_fell": TerminationTermCfg( + func=ip_mdp.any_pole_angle_exceeds, + params={ + "limits": { + POLE_1_JOINT: POLE_1_ANGLE_LIMIT, + POLE_2_JOINT: POLE_2_ANGLE_LIMIT, + }, + }, + ), + "cart_out": TerminationTermCfg( + func=ip_mdp.cart_position_exceeds, + params={"limit": CART_POSITION_LIMIT, "joint_name": CART_JOINT}, + ), + }, + events_cfg={ + "reset_joints": EventTermCfg( + mode="reset", + func=mdp.reset_joints_to_default, + params={"pos_jitter": 0.05, "vel_jitter": 0.05}, + ), + }, + ) + if not play: + cfg.events_cfg["push_cart"] = EventTermCfg( + mode="interval", + interval_range_s=(3.0, 6.0), + func=ip_mdp.push_cart_by_setting_joint_velocity, + params={"velocity_range": (-0.3, 0.3), "joint_name": CART_JOINT}, + ) + return cfg diff --git a/examples/inverted_pendulum/src/genelab_inverted_pendulum/double/ppo_cfg.py b/examples/inverted_pendulum/src/genelab_inverted_pendulum/double/ppo_cfg.py new file mode 100644 index 00000000..7be5e6b7 --- /dev/null +++ b/examples/inverted_pendulum/src/genelab_inverted_pendulum/double/ppo_cfg.py @@ -0,0 +1,44 @@ +"""PPO runner config for the double inverted-pendulum task.""" + +from genelab.rl import RslRlModelCfg, RslRlOnPolicyRunnerCfg, RslRlPpoAlgorithmCfg + + +def double_inverted_pendulum_ppo_runner_cfg() -> RslRlOnPolicyRunnerCfg: + return RslRlOnPolicyRunnerCfg( + seed=42, + num_steps_per_env=24, + max_iterations=300, + save_interval=50, + experiment_name="double_inverted_pendulum_flat", + logger="tensorboard", + clip_actions=100.0, + actor=RslRlModelCfg( + hidden_dims=(256, 128), + activation="elu", + obs_normalization=True, + distribution_cfg={ + "class_name": "GaussianDistribution", + "init_std": 1.0, + "std_type": "scalar", + }, + ), + critic=RslRlModelCfg( + hidden_dims=(256, 128), + activation="elu", + obs_normalization=True, + ), + algorithm=RslRlPpoAlgorithmCfg( + learning_rate=1.0e-3, + schedule="adaptive", + gamma=0.99, + lam=0.95, + clip_param=0.2, + entropy_coef=0.005, + num_learning_epochs=5, + num_mini_batches=4, + desired_kl=0.01, + max_grad_norm=1.0, + value_loss_coef=1.0, + use_clipped_value_loss=True, + ), + ) diff --git a/examples/inverted_pendulum/src/genelab_inverted_pendulum/double/robot.py b/examples/inverted_pendulum/src/genelab_inverted_pendulum/double/robot.py new file mode 100644 index 00000000..c72b4578 --- /dev/null +++ b/examples/inverted_pendulum/src/genelab_inverted_pendulum/double/robot.py @@ -0,0 +1,37 @@ +"""Double inverted-pendulum robot factory.""" + +from dataclasses import dataclass, field + +from genelab.envs.manager_based_rl_env import RobotEntityCfg + +from genelab_inverted_pendulum.double.constants import ( + CART_ACTION_SCALE, + DEFAULT_JOINT_POS, + DOUBLE_INVERTED_PENDULUM_MJCF, + INIT_BASE_HEIGHT, + JOINT_KP, + JOINT_KV, +) + + +@dataclass +class DoubleInvertedPendulumRobotCfg: + """User-facing double-pendulum config. Wraps a ``RobotEntityCfg`` for the env.""" + + mjcf_path: str = field(default_factory=lambda: str(DOUBLE_INVERTED_PENDULUM_MJCF)) + init_height: float = INIT_BASE_HEIGHT + + def to_entity_cfg(self) -> RobotEntityCfg: + return RobotEntityCfg( + mjcf_path=self.mjcf_path, + init_pos=(0.0, 0.0, self.init_height), + init_quat=(1.0, 0.0, 0.0, 0.0), + default_joint_pos=dict(DEFAULT_JOINT_POS), + joint_kp=dict(JOINT_KP), + joint_kv=dict(JOINT_KV), + action_scale=dict(CART_ACTION_SCALE), + ) + + +def get_double_inverted_pendulum_robot_cfg() -> DoubleInvertedPendulumRobotCfg: + return DoubleInvertedPendulumRobotCfg() diff --git a/examples/inverted_pendulum/src/genelab_inverted_pendulum/mdp.py b/examples/inverted_pendulum/src/genelab_inverted_pendulum/mdp.py new file mode 100644 index 00000000..f36bebae --- /dev/null +++ b/examples/inverted_pendulum/src/genelab_inverted_pendulum/mdp.py @@ -0,0 +1,127 @@ +"""Cart-pole reward / termination / event terms reused by single and double tasks.""" + +from typing import TYPE_CHECKING + +import torch + +if TYPE_CHECKING: + from genelab.envs.manager_based_rl_env import ManagerBasedRlEnv + + +_JOINT_INDEX_CACHE: dict[int, dict[str, int]] = {} + + +def _joint_index(env: "ManagerBasedRlEnv", joint_name: str) -> int: + cache = _JOINT_INDEX_CACHE.setdefault(id(env), {}) + idx = cache.get(joint_name) + if idx is None: + idx = env.joint_names.index(joint_name) + cache[joint_name] = idx + return idx + + +def alive_bonus(env: "ManagerBasedRlEnv") -> torch.Tensor: + return torch.ones(env.num_envs, device=env.device) + + +def cart_position_l2(env: "ManagerBasedRlEnv", joint_name: str = "cart_slide") -> torch.Tensor: + return env.robot_state.joint_pos[:, _joint_index(env, joint_name)].square() + + +def cart_velocity_l2(env: "ManagerBasedRlEnv", joint_name: str = "cart_slide") -> torch.Tensor: + return env.robot_state.joint_vel[:, _joint_index(env, joint_name)].square() + + +def pole_upright(env: "ManagerBasedRlEnv", joint_name: str = "pole_hinge") -> torch.Tensor: + pole_angle = env.robot_state.joint_pos[:, _joint_index(env, joint_name)] + return torch.cos(pole_angle) + + +def pole_angle_l2(env: "ManagerBasedRlEnv", joint_name: str = "pole_hinge") -> torch.Tensor: + return env.robot_state.joint_pos[:, _joint_index(env, joint_name)].square() + + +def pole_velocity_l2(env: "ManagerBasedRlEnv", joint_name: str = "pole_hinge") -> torch.Tensor: + return env.robot_state.joint_vel[:, _joint_index(env, joint_name)].square() + + +def double_pole_upright( + env: "ManagerBasedRlEnv", + joint_names: tuple[str, ...] = ("pole_1_hinge", "pole_2_hinge"), +) -> torch.Tensor: + """Average of ``cos(angle)`` across the listed pole hinges.""" + indices = [_joint_index(env, n) for n in joint_names] + angles = env.robot_state.joint_pos[:, indices] + return torch.cos(angles).mean(dim=-1) + + +def double_pole_alignment( + env: "ManagerBasedRlEnv", + joint_names: tuple[str, ...] = ("pole_1_hinge", "pole_2_hinge"), +) -> torch.Tensor: + """``(pole_2 - pole_1)^2`` so PPO learns to keep the chain colinear.""" + indices = [_joint_index(env, n) for n in joint_names] + angles = env.robot_state.joint_pos[:, indices] + return (angles[:, 1] - angles[:, 0]).square() + + +def double_pole_velocity_l2( + env: "ManagerBasedRlEnv", + joint_names: tuple[str, ...] = ("pole_1_hinge", "pole_2_hinge"), +) -> torch.Tensor: + indices = [_joint_index(env, n) for n in joint_names] + return env.robot_state.joint_vel[:, indices].square().sum(dim=-1) + + +def cart_position_exceeds( + env: "ManagerBasedRlEnv", + limit: float, + joint_name: str = "cart_slide", +) -> torch.Tensor: + return env.robot_state.joint_pos[:, _joint_index(env, joint_name)].abs() > limit + + +def pole_angle_exceeds( + env: "ManagerBasedRlEnv", + limit: float, + joint_name: str = "pole_hinge", +) -> torch.Tensor: + return env.robot_state.joint_pos[:, _joint_index(env, joint_name)].abs() > limit + + +def any_pole_angle_exceeds( + env: "ManagerBasedRlEnv", + limits: dict[str, float], +) -> torch.Tensor: + """``done`` if any of the listed pole hinges exceeds its per-joint limit.""" + done = torch.zeros(env.num_envs, dtype=torch.bool, device=env.device) + for joint_name, limit in limits.items(): + idx = _joint_index(env, joint_name) + done |= env.robot_state.joint_pos[:, idx].abs() > limit + return done + + +def push_cart_by_setting_joint_velocity( + env: "ManagerBasedRlEnv", + env_ids: torch.Tensor | None, + velocity_range: tuple[float, float], + joint_name: str = "cart_slide", +) -> None: + """Interval-mode disturbance: overwrite the cart's slide velocity.""" + if env_ids is None or env_ids.numel() == 0: + return + n = int(env_ids.numel()) + idx = _joint_index(env, joint_name) + lo, hi = velocity_range + sample = torch.empty(n, device=env.device).uniform_(lo, hi) + current = env.robot_state.joint_vel.clone() + current[env_ids, idx] = sample + set_vel = getattr(env.robot, "set_dofs_velocity", None) + if set_vel is None: + return + actuated_idx = getattr(env, "_actuated_dof_idx", None) + target = current[env_ids] + try: + set_vel(target, actuated_idx, envs_idx=env_ids) + except TypeError: + set_vel(target, actuated_idx) diff --git a/examples/inverted_pendulum/src/genelab_inverted_pendulum/single/__init__.py b/examples/inverted_pendulum/src/genelab_inverted_pendulum/single/__init__.py new file mode 100644 index 00000000..955b4621 --- /dev/null +++ b/examples/inverted_pendulum/src/genelab_inverted_pendulum/single/__init__.py @@ -0,0 +1,33 @@ +"""Single inverted-pendulum robot, env, and PPO configuration.""" + +from genelab_inverted_pendulum.single.constants import ( + CART_ACTION_SCALE, + CART_JOINT, + DEFAULT_JOINT_POS, + INVERTED_PENDULUM_MJCF, + JOINT_KP, + JOINT_KV, + POLE_JOINT, + POLE_LINK, +) +from genelab_inverted_pendulum.single.env_cfg import inverted_pendulum_env_cfg +from genelab_inverted_pendulum.single.ppo_cfg import inverted_pendulum_ppo_runner_cfg +from genelab_inverted_pendulum.single.robot import ( + InvertedPendulumRobotCfg, + get_inverted_pendulum_robot_cfg, +) + +__all__ = [ + "CART_ACTION_SCALE", + "CART_JOINT", + "DEFAULT_JOINT_POS", + "INVERTED_PENDULUM_MJCF", + "InvertedPendulumRobotCfg", + "JOINT_KP", + "JOINT_KV", + "POLE_JOINT", + "POLE_LINK", + "get_inverted_pendulum_robot_cfg", + "inverted_pendulum_env_cfg", + "inverted_pendulum_ppo_runner_cfg", +] diff --git a/examples/inverted_pendulum/src/genelab_inverted_pendulum/single/constants.py b/examples/inverted_pendulum/src/genelab_inverted_pendulum/single/constants.py new file mode 100644 index 00000000..b34f37fa --- /dev/null +++ b/examples/inverted_pendulum/src/genelab_inverted_pendulum/single/constants.py @@ -0,0 +1,28 @@ +"""Single inverted-pendulum constants: MJCF path, joint gains, action scale.""" + +from pathlib import Path +from typing import Final + +# .../single/constants.py → .../assets/inverted_pendulum.xml +INVERTED_PENDULUM_MJCF: Final = ( + Path(__file__).resolve().parents[3] / "assets" / "inverted_pendulum.xml" +).resolve() + +CART_JOINT: Final = "cart_slide" +POLE_JOINT: Final = "pole_hinge" +POLE_LINK: Final = "pole" + +# Underactuated: only the cart slide receives PD. The pole hinge defaults to kp=0/kv=0. +JOINT_KP: Final[dict[str, float]] = {CART_JOINT: 80.0} +JOINT_KV: Final[dict[str, float]] = {CART_JOINT: 8.0} + +# JointPositionAction scale — policy action ∈ [-1, 1] maps to ±1.0 m setpoint about the cart origin. +CART_ACTION_SCALE: Final[dict[str, float]] = {CART_JOINT: 1.0} + +DEFAULT_JOINT_POS: Final[dict[str, float]] = {CART_JOINT: 0.0, POLE_JOINT: 0.0} + +# Termination limits. +CART_POSITION_LIMIT: Final = 2.4 +POLE_ANGLE_LIMIT: Final = 0.4 # ~23° + +INIT_BASE_HEIGHT: Final = 0.12 diff --git a/examples/inverted_pendulum/src/genelab_inverted_pendulum/single/env_cfg.py b/examples/inverted_pendulum/src/genelab_inverted_pendulum/single/env_cfg.py new file mode 100644 index 00000000..d08f7a46 --- /dev/null +++ b/examples/inverted_pendulum/src/genelab_inverted_pendulum/single/env_cfg.py @@ -0,0 +1,121 @@ +"""Flat-ground single inverted-pendulum env config (manager-based + Genesis).""" + +from genelab import mdp +from genelab.configs import SceneCfg +from genelab.envs.manager_based_rl_env import ManagerBasedRlEnvCfg +from genelab.managers import ( + EventTermCfg, + ObservationGroupCfg, + ObservationTermCfg, + RewardTermCfg, + TerminationTermCfg, +) +from genelab.mdp.actions.joint_position import JointPositionActionCfg +from genelab.mdp.noise import Unoise +from genelab.sensor import BodyVelocitySensorCfg + +from genelab_inverted_pendulum import mdp as ip_mdp +from genelab_inverted_pendulum.single.constants import ( + CART_ACTION_SCALE, + CART_JOINT, + CART_POSITION_LIMIT, + POLE_ANGLE_LIMIT, + POLE_JOINT, + POLE_LINK, +) +from genelab_inverted_pendulum.single.robot import get_inverted_pendulum_robot_cfg + + +def _obs_terms() -> dict[str, ObservationTermCfg]: + return { + "joint_pos": ObservationTermCfg( + func=mdp.joint_pos_rel, + noise=Unoise(-0.005, 0.005), + ), + "joint_vel": ObservationTermCfg( + func=mdp.joint_vel_rel, + scale=0.1, + noise=Unoise(-0.05, 0.05), + ), + "pole_ang_vel": ObservationTermCfg( + func=mdp.sensor_data, + params={"sensor_name": "pole_ang_vel"}, + scale=0.2, + noise=Unoise(-0.1, 0.1), + ), + "actions": ObservationTermCfg(func=mdp.last_action), + } + + +def inverted_pendulum_env_cfg(play: bool = False) -> ManagerBasedRlEnvCfg: + """Flat-ground inverted-pendulum env config.""" + robot_entity_cfg = get_inverted_pendulum_robot_cfg().to_entity_cfg() + + cfg = ManagerBasedRlEnvCfg( + scene=SceneCfg( + num_envs=4096 if not play else 1, + dt=0.005, + substeps=1, + env_spacing=(2.5, 2.5), + vis=play, + mouse_interaction=play, + sensors=( + BodyVelocitySensorCfg( + name="pole_ang_vel", + link_name=POLE_LINK, + measure="ang_vel", + ), + ), + ), + decimation=2, + episode_length_s=10.0, + device="cuda", + robot=robot_entity_cfg, + actions_cfg={ + "cart": JointPositionActionCfg( + asset_name="robot", + joint_names=(CART_JOINT,), + scale=dict(CART_ACTION_SCALE), + use_default_offset=True, + ), + }, + observations_cfg={ + "policy": ObservationGroupCfg(enable_corruption=True, terms=_obs_terms()), + "critic": ObservationGroupCfg(enable_corruption=False, terms=_obs_terms()), + }, + rewards_cfg={ + "alive": RewardTermCfg(func=ip_mdp.alive_bonus, weight=1.0), + "pole_upright": RewardTermCfg(func=ip_mdp.pole_upright, weight=3.0), + "pole_angle": RewardTermCfg(func=ip_mdp.pole_angle_l2, weight=-0.5), + "cart_position": RewardTermCfg(func=ip_mdp.cart_position_l2, weight=-0.05), + "cart_velocity": RewardTermCfg(func=ip_mdp.cart_velocity_l2, weight=-0.005), + "pole_velocity": RewardTermCfg(func=ip_mdp.pole_velocity_l2, weight=-0.01), + "action_rate": RewardTermCfg(func=mdp.action_rate_l2, weight=-0.005), + }, + terminations_cfg={ + "time_out": TerminationTermCfg(func=mdp.time_out, time_out=True), + "pole_fell": TerminationTermCfg( + func=ip_mdp.pole_angle_exceeds, + params={"limit": POLE_ANGLE_LIMIT, "joint_name": POLE_JOINT}, + ), + "cart_out": TerminationTermCfg( + func=ip_mdp.cart_position_exceeds, + params={"limit": CART_POSITION_LIMIT, "joint_name": CART_JOINT}, + ), + }, + events_cfg={ + "reset_joints": EventTermCfg( + mode="reset", + func=mdp.reset_joints_to_default, + params={"pos_jitter": 0.05, "vel_jitter": 0.05}, + ), + }, + ) + if not play: + cfg.events_cfg["push_cart"] = EventTermCfg( + mode="interval", + interval_range_s=(2.0, 5.0), + func=ip_mdp.push_cart_by_setting_joint_velocity, + params={"velocity_range": (-0.5, 0.5), "joint_name": CART_JOINT}, + ) + return cfg diff --git a/examples/inverted_pendulum/src/genelab_inverted_pendulum/single/ppo_cfg.py b/examples/inverted_pendulum/src/genelab_inverted_pendulum/single/ppo_cfg.py new file mode 100644 index 00000000..6932b5bf --- /dev/null +++ b/examples/inverted_pendulum/src/genelab_inverted_pendulum/single/ppo_cfg.py @@ -0,0 +1,44 @@ +"""PPO runner config for the single inverted-pendulum task.""" + +from genelab.rl import RslRlModelCfg, RslRlOnPolicyRunnerCfg, RslRlPpoAlgorithmCfg + + +def inverted_pendulum_ppo_runner_cfg() -> RslRlOnPolicyRunnerCfg: + return RslRlOnPolicyRunnerCfg( + seed=42, + num_steps_per_env=24, + max_iterations=150, + save_interval=25, + experiment_name="inverted_pendulum_flat", + logger="tensorboard", + clip_actions=100.0, + actor=RslRlModelCfg( + hidden_dims=(128, 128), + activation="elu", + obs_normalization=True, + distribution_cfg={ + "class_name": "GaussianDistribution", + "init_std": 1.0, + "std_type": "scalar", + }, + ), + critic=RslRlModelCfg( + hidden_dims=(128, 128), + activation="elu", + obs_normalization=True, + ), + algorithm=RslRlPpoAlgorithmCfg( + learning_rate=1.0e-3, + schedule="adaptive", + gamma=0.99, + lam=0.95, + clip_param=0.2, + entropy_coef=0.005, + num_learning_epochs=5, + num_mini_batches=4, + desired_kl=0.01, + max_grad_norm=1.0, + value_loss_coef=1.0, + use_clipped_value_loss=True, + ), + ) diff --git a/examples/inverted_pendulum/src/genelab_inverted_pendulum/single/robot.py b/examples/inverted_pendulum/src/genelab_inverted_pendulum/single/robot.py new file mode 100644 index 00000000..4c966bf3 --- /dev/null +++ b/examples/inverted_pendulum/src/genelab_inverted_pendulum/single/robot.py @@ -0,0 +1,37 @@ +"""Single inverted-pendulum robot factory.""" + +from dataclasses import dataclass, field + +from genelab.envs.manager_based_rl_env import RobotEntityCfg + +from genelab_inverted_pendulum.single.constants import ( + CART_ACTION_SCALE, + DEFAULT_JOINT_POS, + INIT_BASE_HEIGHT, + INVERTED_PENDULUM_MJCF, + JOINT_KP, + JOINT_KV, +) + + +@dataclass +class InvertedPendulumRobotCfg: + """User-facing single-pendulum config. Wraps a ``RobotEntityCfg`` for the env.""" + + mjcf_path: str = field(default_factory=lambda: str(INVERTED_PENDULUM_MJCF)) + init_height: float = INIT_BASE_HEIGHT + + def to_entity_cfg(self) -> RobotEntityCfg: + return RobotEntityCfg( + mjcf_path=self.mjcf_path, + init_pos=(0.0, 0.0, self.init_height), + init_quat=(1.0, 0.0, 0.0, 0.0), + default_joint_pos=dict(DEFAULT_JOINT_POS), + joint_kp=dict(JOINT_KP), + joint_kv=dict(JOINT_KV), + action_scale=dict(CART_ACTION_SCALE), + ) + + +def get_inverted_pendulum_robot_cfg() -> InvertedPendulumRobotCfg: + return InvertedPendulumRobotCfg() diff --git a/examples/inverted_pendulum/src/genelab_inverted_pendulum/tasks.py b/examples/inverted_pendulum/src/genelab_inverted_pendulum/tasks.py new file mode 100644 index 00000000..1c8ef49c --- /dev/null +++ b/examples/inverted_pendulum/src/genelab_inverted_pendulum/tasks.py @@ -0,0 +1,144 @@ +"""GeneLab extension entry point: registers single + double inverted-pendulum tasks.""" + +from genelab.configs import TaskCfg +from genelab.registry import ( + ENVS, + ROBOTS, + TASKS, + register_env, + register_robot, + register_task, +) + +from genelab_inverted_pendulum.double import ( + DoubleInvertedPendulumRobotCfg, + double_inverted_pendulum_env_cfg, + double_inverted_pendulum_ppo_runner_cfg, + get_double_inverted_pendulum_robot_cfg, +) +from genelab_inverted_pendulum.single import ( + InvertedPendulumRobotCfg, + get_inverted_pendulum_robot_cfg, + inverted_pendulum_env_cfg, + inverted_pendulum_ppo_runner_cfg, +) + +INVERTED_PENDULUM_TASK_ID = "GeneLab-Inverted-Pendulum-v0" +DOUBLE_PENDULUM_TASK_ID = "GeneLab-Double-Inverted-Pendulum-v0" +ROBOT_NAME_SINGLE = "inverted-pendulum" +ROBOT_NAME_DOUBLE = "double-inverted-pendulum" +ENV_NAME_SINGLE = "inverted-pendulum-env" +ENV_NAME_DOUBLE = "double-inverted-pendulum-env" + + +def _build_single_env(play: bool = False): + from genelab.envs.manager_based_rl_env import ManagerBasedRlEnv + + return ManagerBasedRlEnv(inverted_pendulum_env_cfg(play=play)) + + +def _build_double_env(play: bool = False): + from genelab.envs.manager_based_rl_env import ManagerBasedRlEnv + + return ManagerBasedRlEnv(double_inverted_pendulum_env_cfg(play=play)) + + +class InvertedPendulumTask: + """Trainable single inverted-pendulum task. Delegates to ``genelab.rl.runner``.""" + + def __init__(self) -> None: + self.cfg = TaskCfg( + name=INVERTED_PENDULUM_TASK_ID, + env_name=ENV_NAME_SINGLE, + robot_name=ROBOT_NAME_SINGLE, + env=inverted_pendulum_env_cfg(play=False), + play_env=inverted_pendulum_env_cfg(play=True), + agent=inverted_pendulum_ppo_runner_cfg(), + trainable=True, + ) + + def play(self) -> None: + from genelab.rl import play_task + + play_task(self.cfg.name, checkpoint=None) + + def train(self) -> None: + from genelab.rl import RslRlOnPolicyRunnerCfg, train_task + + agent = self.cfg.agent + if not isinstance(agent, RslRlOnPolicyRunnerCfg): + raise TypeError(f"agent cfg has unexpected type {type(agent).__name__}") + train_task(self.cfg.name, agent) + + +class DoubleInvertedPendulumTask: + """Trainable double inverted-pendulum task.""" + + def __init__(self) -> None: + self.cfg = TaskCfg( + name=DOUBLE_PENDULUM_TASK_ID, + env_name=ENV_NAME_DOUBLE, + robot_name=ROBOT_NAME_DOUBLE, + env=double_inverted_pendulum_env_cfg(play=False), + play_env=double_inverted_pendulum_env_cfg(play=True), + agent=double_inverted_pendulum_ppo_runner_cfg(), + trainable=True, + ) + + def play(self) -> None: + from genelab.rl import play_task + + play_task(self.cfg.name, checkpoint=None) + + def train(self) -> None: + from genelab.rl import RslRlOnPolicyRunnerCfg, train_task + + agent = self.cfg.agent + if not isinstance(agent, RslRlOnPolicyRunnerCfg): + raise TypeError(f"agent cfg has unexpected type {type(agent).__name__}") + train_task(self.cfg.name, agent) + + +def register() -> None: + if ROBOT_NAME_SINGLE not in ROBOTS: + register_robot( + ROBOT_NAME_SINGLE, + get_inverted_pendulum_robot_cfg, + description="Cart with a single hinge-mounted inverted pole.", + cfg_type=InvertedPendulumRobotCfg, + ) + if ROBOT_NAME_DOUBLE not in ROBOTS: + register_robot( + ROBOT_NAME_DOUBLE, + get_double_inverted_pendulum_robot_cfg, + description="Cart with two serially linked inverted poles.", + cfg_type=DoubleInvertedPendulumRobotCfg, + ) + if ENV_NAME_SINGLE not in ENVS: + register_env( + ENV_NAME_SINGLE, + lambda: _build_single_env(play=False), + description="Single inverted-pendulum balancing on a flat plane.", + cfg_type=type(None), + ) + if ENV_NAME_DOUBLE not in ENVS: + register_env( + ENV_NAME_DOUBLE, + lambda: _build_double_env(play=False), + description="Double inverted-pendulum balancing on a flat plane.", + cfg_type=type(None), + ) + if INVERTED_PENDULUM_TASK_ID not in TASKS: + register_task( + INVERTED_PENDULUM_TASK_ID, + InvertedPendulumTask, + description="PPO inverted-pendulum balancing (single pole).", + cfg_type=TaskCfg, + ) + if DOUBLE_PENDULUM_TASK_ID not in TASKS: + register_task( + DOUBLE_PENDULUM_TASK_ID, + DoubleInvertedPendulumTask, + description="PPO double-inverted-pendulum balancing (two stacked poles).", + cfg_type=TaskCfg, + ) diff --git a/mkdocs.yml b/mkdocs.yml index dc596d37..06b7066f 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -100,6 +100,7 @@ plugins: Sensors: 传感器 Extensions: 扩展加载 Examples: 示例 + Inverted Pendulum: 倒立摆 API Reference: API 参考 Contributing: 贡献指南 - mkdocstrings: @@ -163,5 +164,6 @@ nav: - Extensions: concepts/extensions.md - Examples: - Overview: examples/overview.md + - Inverted Pendulum: examples/inverted-pendulum.md - API Reference: api/reference.md - Contributing: contributing.md diff --git a/pyproject.toml b/pyproject.toml index 4898a639..f3ff4c8a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -115,5 +115,5 @@ src = ["src"] [tool.pytest.ini_options] addopts = ["-ra"] -pythonpath = ["examples/genelab_examples/src", "examples/unitree/src"] +pythonpath = ["examples/genelab_examples/src", "examples/unitree/src", "examples/inverted_pendulum/src"] testpaths = ["tests"] diff --git a/tests/test_inverted_pendulum_examples.py b/tests/test_inverted_pendulum_examples.py new file mode 100644 index 00000000..726b59c4 --- /dev/null +++ b/tests/test_inverted_pendulum_examples.py @@ -0,0 +1,16 @@ +"""Registration smoke test for the inverted-pendulum extension.""" + +from genelab.registry import ENVS, ROBOTS, TASKS, load_extension_module + + +def test_inverted_pendulum_extension_registers() -> None: + load_extension_module("genelab_inverted_pendulum.tasks") + + assert "GeneLab-Inverted-Pendulum-v0" in TASKS.names() + assert "GeneLab-Double-Inverted-Pendulum-v0" in TASKS.names() + + assert "inverted-pendulum" in ROBOTS.names() + assert "double-inverted-pendulum" in ROBOTS.names() + + assert "inverted-pendulum-env" in ENVS.names() + assert "double-inverted-pendulum-env" in ENVS.names()