diff --git a/.agents/skills/write-motrixlab-task-docs/SKILL.md b/.agents/skills/write-motrixlab-task-docs/SKILL.md index 3f00bf66..b648075f 100644 --- a/.agents/skills/write-motrixlab-task-docs/SKILL.md +++ b/.agents/skills/write-motrixlab-task-docs/SKILL.md @@ -1,17 +1,26 @@ --- name: write-motrixlab-task-docs -description: Write, restructure, or review bilingual Sphinx/MyST user documentation for MotrixLab task environments, from simple single-page tasks to reusable task families. Use when documenting a registered environment, its runtime contract, configuration, presets, training evidence, or an implemented extension workflow under docs/source/{zh_CN,en}/user_guide/envs/. +description: Write, restructure, or review bilingual Sphinx/MyST user documentation for MotrixLab — task environments (from simple single-page tasks to reusable task families) and framework-level tutorial pages (concepts, building environments, training, advanced topics) under docs/source/{zh_CN,en}/user_guide/. Use when documenting a registered environment, its runtime contract, configuration, presets, training evidence, or an implemented extension workflow, or when authoring or restructuring tutorial pages and tutorial navigation. --- # Write MotrixLab Task Docs -Create user-facing task-environment documentation from current repository evidence. Keep Chinese and English pages aligned, -state exact runtime semantics, and validate the rendered Sphinx output. +Create user-facing MotrixLab documentation from current repository evidence. Keep Chinese and English pages aligned, +state exact runtime semantics, and validate the rendered Sphinx output. Two surfaces are covered: + +- **Task-environment pages** (`user_guide/envs/`): follow [references/writing-standard.md](references/writing-standard.md). +- **Framework tutorial pages** (`user_guide/tutorial/`): follow + [references/tutorial-standard.md](references/tutorial-standard.md) — layered information architecture, macro-before-detail + ordering, SVG pipeline diagram rules (compact snake layout, chip sub-items, light/dark pairs), and toctree hygiene. + +The evidence, bilingual, and validation rules below apply to both surfaces. ## Read the standard -Read [references/writing-standard.md](references/writing-standard.md) completely before drafting or restructuring a task -environment page. Apply only sections supported by the target environment; do not add empty boilerplate. +Read the standard for the target surface completely before drafting or restructuring: +[references/writing-standard.md](references/writing-standard.md) for task-environment pages, +[references/tutorial-standard.md](references/tutorial-standard.md) for framework tutorial pages. Apply only sections +supported by the target; do not add empty boilerplate. ## Establish evidence @@ -69,7 +78,8 @@ lifecycle, or configuration-schema chapter by default. Edit `docs/source/zh_CN/` and `docs/source/en/` together unless the user explicitly scopes the work to one language. Preserve technical identifiers across languages and translate meaning rather than sentence structure. Use established terminology from -the neighboring pages. +the neighboring pages. Keep section order, figures, tables, and toctree/link changes mirrored on both sides — restructure one +language only together with the other. ## Handle media and performance evidence diff --git a/.agents/skills/write-motrixlab-task-docs/references/tutorial-standard.md b/.agents/skills/write-motrixlab-task-docs/references/tutorial-standard.md new file mode 100644 index 00000000..7546d295 --- /dev/null +++ b/.agents/skills/write-motrixlab-task-docs/references/tutorial-standard.md @@ -0,0 +1,112 @@ +# MotrixLab Framework Tutorial Documentation Standard + +Applies to framework-level tutorial pages under `docs/source/{zh_CN,en}/user_guide/tutorial/` — concepts, +building environments, training, and advanced topics. Task-environment pages follow `writing-standard.md` +instead; the language, validation, and bilingual rules of that standard apply to tutorials too. + +## Contents + +1. Layered information architecture +2. Page-internal order +3. Diagrams (SVG) +4. Navigation and toctree hygiene +5. Build hygiene specific to tutorials + +## 1. Layered information architecture + +Write from the macro picture down to details, and let the directory tree express that order: + +- Each layer (for example `building_envs/`, `training/`, `advanced/`) has an `index.md` that answers only + three questions: what problem this layer solves, which steps it contains, and where to go next. Details + sink into child pages. +- Prefer a small number of layers keyed to the user's journey (concepts → build → train → advanced) over a + flat list of peer pages. +- A thin topic is a section in its parent index, never a standalone page. Before creating a page, state the + question it answers; if a short section answers it, do not create the page. +- The macro element of a topic (for example the shared lifecycle) comes before the details that depend on + it. In a page, the big-picture section precedes config tables and contract details. + +One figure, one question: an overview figure draws only the shared skeleton (for example the +`ArrayEnv` step pipeline with framework-owned stages). Workflow-specific annotations — which stages a +`DirectEnv` hook fills, which stages manager kernels drive — belong in that workflow's own page, not in +the shared overview. + +## 2. Page-internal order + +Use this order by default: + +1. a complete minimal working example (or, for concept pages, the macro diagram); +2. concept expansion, section per element; +3. contract details (semantics, invariants, common pitfalls); +4. reference tables and field lists, as lookup material at the end. + +Do not open a page with comparison tables, field tables, or contract enumerations before the reader has a +mental model. Keep a short bridge sentence at each transition ("the next section builds the big picture +before the details of …"). + +Reference material the reader copies or looks up — directory trees, code, per-field tables — stays text. + +## 3. Diagrams (SVG) + +Draw as SVG: architecture layers, step pipelines, boundary/topology diagrams, workflow comparisons. Keep +as text: repository/runs directory trees, code, field and hook reference tables. + +Layout rules: + +- Canvas width matches the content column: at most ~800 px, so an embedded figure renders near 1:1 and + labels stay readable. A six-stage pipeline uses a 2×3 snake layout (three columns per row, second row + flowing right-to-left) instead of one long row. +- Stage boxes hold only the stage title. Sub-items (term callbacks, hook names) are individual rounded + chips stacked under their stage, connected by short vertical lines, color-coded by role with a legend. + Do not pack many small text lines inside one box. +- Route cross-references (for example "a mid-transition request joins the next reset") without crossing + chips; a dashed annotation chip placed under the target chip is usually clearer than a long connector. +- Figure labels use English technical identifiers; captions and the surrounding prose are bilingual. + +Theme and rendering rules: + +- Provide a light/dark pair and embed with the `only-light` / `only-dark` classes: + + ````markdown + ```{image} /_static/images/tutorial/
-light.svg + :alt: + :class: only-light + ``` + + ```{image} /_static/images/tutorial/
-dark.svg + :alt: + :class: only-dark + ``` + ```` + +- Use solid per-theme fills, not `rgba()` translucency — translucent fills composite unreliably across + SVG renderers. Draw explicit arrowhead triangles instead of `marker-end`. Avoid clipped edge labels. +- Store sources under `docs/source/_static/images/tutorial/`. +- Before embedding, rasterize both variants and actually look at them (dark on a dark background): check + arrows render, no text is clipped or overlapped, and contrast holds. Never claim an unviewed figure is + correct. + +## 4. Navigation and toctree hygiene + +- A toctree `:caption:` and the title of the index page below it must not be identical; the sidebar would + render the same name twice (caption → page → children). +- An index page must not repeat the sidebar as an in-page text tree. A short layered summary (one line per + child, what problem it solves) plus path guidance is the replacement. +- Keep toctree entries, page titles, and cross-reference link texts consistent with the actual `#` heading + of the target page. +- When moving or renaming pages: use `git mv`, then grep the whole `docs/source` tree for old filenames + and stale relative depths (`../../../configs`-style links change when a page moves a level). Update + `zh_CN` and `en` toctrees and links in the same change. + +## 5. Build hygiene specific to tutorials + +`docs/source/conf.py` copies the selected language directory over the build source root with +`dirs_exist_ok=True` and never deletes removed files. Stale copies of deleted or moved pages surface as +"document isn't included in any toctree" warnings that do not come from your change. Clean before building: + +```bash +git clean -fdX docs/source +``` + +Then run the standard strict builds from `writing-standard.md` section 10, and inspect the rendered +sidebar: nesting depth, entry titles, and that the new pages appear where the reading path claims. diff --git a/docs/source/_static/images/tutorial/direct-env-lifecycle-dark.svg b/docs/source/_static/images/tutorial/direct-env-lifecycle-dark.svg new file mode 100644 index 00000000..1822069d --- /dev/null +++ b/docs/source/_static/images/tutorial/direct-env-lifecycle-dark.svg @@ -0,0 +1,43 @@ + + One control step of DirectEnv + + apply_action + + hook + + + + physics_step + + sim.step × substeps + + + + compute_transition + + hook + + + + compute_observation + + hook + + auto-reset + + reset(env_ids) + + truncation check + + episode_steps ≥ max + + time_outs mask + + + + + + ArrayEnv framework + + Direct hook (environment implementation) + diff --git a/docs/source/_static/images/tutorial/direct-env-lifecycle-light.svg b/docs/source/_static/images/tutorial/direct-env-lifecycle-light.svg new file mode 100644 index 00000000..9705a73c --- /dev/null +++ b/docs/source/_static/images/tutorial/direct-env-lifecycle-light.svg @@ -0,0 +1,43 @@ + + One control step of DirectEnv + + apply_action + + hook + + + + physics_step + + sim.step × substeps + + + + compute_transition + + hook + + + + compute_observation + + hook + + auto-reset + + reset(env_ids) + + truncation check + + episode_steps ≥ max + + time_outs mask + + + + + + ArrayEnv framework + + Direct hook (environment implementation) + diff --git a/docs/source/_static/images/tutorial/env-lifecycle-dark.svg b/docs/source/_static/images/tutorial/env-lifecycle-dark.svg new file mode 100644 index 00000000..5b5231d3 --- /dev/null +++ b/docs/source/_static/images/tutorial/env-lifecycle-dark.svg @@ -0,0 +1,25 @@ + + One control step of ArrayEnv + + apply_action + + + + physics_step + + + + compute_transition + + + + compute_observation + + auto-reset + + truncation check + + + + + diff --git a/docs/source/_static/images/tutorial/env-lifecycle-light.svg b/docs/source/_static/images/tutorial/env-lifecycle-light.svg new file mode 100644 index 00000000..64797889 --- /dev/null +++ b/docs/source/_static/images/tutorial/env-lifecycle-light.svg @@ -0,0 +1,25 @@ + + One control step of ArrayEnv + + apply_action + + + + physics_step + + + + compute_transition + + + + compute_observation + + auto-reset + + truncation check + + + + + diff --git a/docs/source/_static/images/tutorial/manager-env-lifecycle-dark.svg b/docs/source/_static/images/tutorial/manager-env-lifecycle-dark.svg new file mode 100644 index 00000000..be081512 --- /dev/null +++ b/docs/source/_static/images/tutorial/manager-env-lifecycle-dark.svg @@ -0,0 +1,58 @@ + + +One control step of ManagerEnv + + + +apply_action +action terms · process() +write ctrl targets + +physics_step +sim.step × substeps + +compute_transition +evaluate kernel +update · rewards · terminations · advance +command · on_transition() + + + + + + + + + + +truncation check +done = terminated | time_outs + +done? + +done rows + +alive rows + + +auto-reset +host-side resets +command reset(ctx) · action reset(env_ids) +reset kernel +reset_env hooks → sim_reset terms +sim_reset_requested rows join this run + + + +compute_observation +observe kernel +observation terms → policy / value + + +→ next control step + + +ArrayEnv framework stage +Manager kernel (@dispatch) +Manager host-side callback + diff --git a/docs/source/_static/images/tutorial/manager-env-lifecycle-light.svg b/docs/source/_static/images/tutorial/manager-env-lifecycle-light.svg new file mode 100644 index 00000000..6918e36a --- /dev/null +++ b/docs/source/_static/images/tutorial/manager-env-lifecycle-light.svg @@ -0,0 +1,58 @@ + + +One control step of ManagerEnv + + + +apply_action +action terms · process() +write ctrl targets + +physics_step +sim.step × substeps + +compute_transition +evaluate kernel +update · rewards · terminations · advance +command · on_transition() + + + + + + + + + + +truncation check +done = terminated | time_outs + +done? + +done rows + +alive rows + + +auto-reset +host-side resets +command reset(ctx) · action reset(env_ids) +reset kernel +reset_env hooks → sim_reset terms +sim_reset_requested rows join this run + + + +compute_observation +observe kernel +observation terms → policy / value + + +→ next control step + + +ArrayEnv framework stage +Manager kernel (@dispatch) +Manager host-side callback + diff --git a/docs/source/_static/images/tutorial/rl-loop-dark.svg b/docs/source/_static/images/tutorial/rl-loop-dark.svg new file mode 100644 index 00000000..7c908573 --- /dev/null +++ b/docs/source/_static/images/tutorial/rl-loop-dark.svg @@ -0,0 +1,36 @@ + + The simulation RL loop + + Policy + network π(aₜ | oₜ) + updated by the trainer + + Environment + N parallel simulated envs + physics step · reward · termination + + Trainer / RL algorithm + collects transitions + updates policy parameters + + + action aₜ = π(oₜ) + + + observation oₜ + + + + reward rₜ · terminated / truncated + transitions (oₜ, aₜ, rₜ, done) + + + + update parameters + + Learning side (updated by training) + + Environment side + + Training orchestration + diff --git a/docs/source/_static/images/tutorial/rl-loop-light.svg b/docs/source/_static/images/tutorial/rl-loop-light.svg new file mode 100644 index 00000000..698c6852 --- /dev/null +++ b/docs/source/_static/images/tutorial/rl-loop-light.svg @@ -0,0 +1,36 @@ + + The simulation RL loop + + Policy + network π(aₜ | oₜ) + updated by the trainer + + Environment + N parallel simulated envs + physics step · reward · termination + + Trainer / RL algorithm + collects transitions + updates policy parameters + + + action aₜ = π(oₜ) + + + observation oₜ + + + + reward rₜ · terminated / truncated + transitions (oₜ, aₜ, rₜ, done) + + + + update parameters + + Learning side (updated by training) + + Environment side + + Training orchestration + diff --git a/docs/source/en/user_guide/envs/whole_body_tracking/index.md b/docs/source/en/user_guide/envs/whole_body_tracking/index.md index 0b4c7f09..f65be639 100644 --- a/docs/source/en/user_guide/envs/whole_body_tracking/index.md +++ b/docs/source/en/user_guide/envs/whole_body_tracking/index.md @@ -171,5 +171,5 @@ python scripts/play.py env=ENV_ID num_envs=16 `play.py` selects the best policy from the latest metadata-backed run for the environment. The WBT play config starts at motion frame 0, disables reset noise and adaptive sampling, and removes the 10-second training time limit. At the clip end it -restarts from frame 0. See [Training Artifacts: the runs Directory and Checkpoint Structure](../../tutorial/runs_and_checkpoints.md) +restarts from frame 0. See [Training Artifacts: the runs Directory and Checkpoint Structure](../../tutorial/training/runs_and_checkpoints.md) for run and checkpoint selection. diff --git a/docs/source/en/user_guide/getting_started/hello_motrixlab.md b/docs/source/en/user_guide/getting_started/hello_motrixlab.md index ad71cc98..b592aa8e 100644 --- a/docs/source/en/user_guide/getting_started/hello_motrixlab.md +++ b/docs/source/en/user_guide/getting_started/hello_motrixlab.md @@ -82,5 +82,5 @@ Next, you can try modifying parameters to observe physical effects under differe ## Next Steps - Learn about the [Basic Framework](../tutorial/basic_frame.md) -- Study [Physics Environment Configuration](../tutorial/physics_environment.md) +- Study [Physics Environment Configuration](../tutorial/building_envs/scene.md) - Browse more [Environments](../envs/index.md) diff --git a/docs/source/en/user_guide/index.md b/docs/source/en/user_guide/index.md index 1cb7fe81..6a37762a 100644 --- a/docs/source/en/user_guide/index.md +++ b/docs/source/en/user_guide/index.md @@ -13,26 +13,13 @@ robots ```{toctree} :caption: Tutorials -:maxdepth: 1 +:maxdepth: 2 +tutorial/index tutorial/basic_frame -tutorial/physics_environment -tutorial/training_environment_config -tutorial/rewards -tutorial/training_and_result -tutorial/runs_and_checkpoints - -``` - -```{toctree} -:caption: Advanced Topics -:maxdepth: 1 - -tutorial/export_onnx -tutorial/motrix_deploy -tutorial/input_devices_and_bindings -tutorial/custom_training_backend - +tutorial/building_envs/index +tutorial/training/index +tutorial/advanced/index ``` ```{toctree} diff --git a/docs/source/en/user_guide/tutorial/custom_training_backend.md b/docs/source/en/user_guide/tutorial/advanced/custom_training_backend.md similarity index 100% rename from docs/source/en/user_guide/tutorial/custom_training_backend.md rename to docs/source/en/user_guide/tutorial/advanced/custom_training_backend.md diff --git a/docs/source/en/user_guide/tutorial/export_onnx.md b/docs/source/en/user_guide/tutorial/advanced/export_onnx.md similarity index 96% rename from docs/source/en/user_guide/tutorial/export_onnx.md rename to docs/source/en/user_guide/tutorial/advanced/export_onnx.md index 1a567f89..b3da493b 100644 --- a/docs/source/en/user_guide/tutorial/export_onnx.md +++ b/docs/source/en/user_guide/tutorial/advanced/export_onnx.md @@ -14,13 +14,13 @@ The following training configurations are currently supported: | RSL-RL | Torch | PPO | Supports an `MLPModel` actor with a single `policy` observation group; state-dependent standard deviation is not supported yet | | Motrix | Torch | FastSAC | Exports the actor's deterministic inference path | -SKRL/JAX checkpoints cannot be exported yet. The command accepts a complete run directory, not a checkpoint file detached from its `metadata.json` and `checkpoints/manifest.json`. See [Training Artifacts: runs Directories and Checkpoints](runs_and_checkpoints.md) for details about run directories and best-policy selection. +SKRL/JAX checkpoints cannot be exported yet. The command accepts a complete run directory, not a checkpoint file detached from its `metadata.json` and `checkpoints/manifest.json`. See [Training Artifacts: runs Directories and Checkpoints](../training/runs_and_checkpoints.md) for details about run directories and best-policy selection. ## Install Dependencies No extra step is needed: the ONNX export and inference dependencies (`onnx`, `onnxruntime`) are part of the default runtime environment installed by `sh install.sh` (see -[Installation](../getting_started/installation.md)). +[Installation](../../getting_started/installation.md)). ## Export the Model diff --git a/docs/source/en/user_guide/tutorial/advanced/index.md b/docs/source/en/user_guide/tutorial/advanced/index.md new file mode 100644 index 00000000..42233a61 --- /dev/null +++ b/docs/source/en/user_guide/tutorial/advanced/index.md @@ -0,0 +1,21 @@ +# Advanced Topics + +Capabilities beyond the basic training flow, read as needed: + +- [Exporting ONNX Policies](export_onnx.md): export a trained policy into a standalone + ONNX model for inference-side integration; +- [Go2 Flat Walking: from Training to Hardware Deployment](motrix_deploy.md): a complete + walkthrough of training, artifact export, MuJoCo checks, and hardware runs; +- [Command Input Architecture](input_devices_and_bindings.md): decouples "how input is + read" from "what targets the policy needs", connecting input devices and task commands; +- [Adding a Custom Training Backend](custom_training_backend.md): integrate a new RL + framework, algorithm, or training backend. + +```{toctree} +:hidden: + +export_onnx +motrix_deploy +input_devices_and_bindings +custom_training_backend +``` diff --git a/docs/source/en/user_guide/tutorial/input_devices_and_bindings.md b/docs/source/en/user_guide/tutorial/advanced/input_devices_and_bindings.md similarity index 100% rename from docs/source/en/user_guide/tutorial/input_devices_and_bindings.md rename to docs/source/en/user_guide/tutorial/advanced/input_devices_and_bindings.md diff --git a/docs/source/en/user_guide/tutorial/input_devices_and_bindings/built_in.md b/docs/source/en/user_guide/tutorial/advanced/input_devices_and_bindings/built_in.md similarity index 100% rename from docs/source/en/user_guide/tutorial/input_devices_and_bindings/built_in.md rename to docs/source/en/user_guide/tutorial/advanced/input_devices_and_bindings/built_in.md diff --git a/docs/source/en/user_guide/tutorial/input_devices_and_bindings/extending.md b/docs/source/en/user_guide/tutorial/advanced/input_devices_and_bindings/extending.md similarity index 100% rename from docs/source/en/user_guide/tutorial/input_devices_and_bindings/extending.md rename to docs/source/en/user_guide/tutorial/advanced/input_devices_and_bindings/extending.md diff --git a/docs/source/en/user_guide/tutorial/motrix_deploy.md b/docs/source/en/user_guide/tutorial/advanced/motrix_deploy.md similarity index 100% rename from docs/source/en/user_guide/tutorial/motrix_deploy.md rename to docs/source/en/user_guide/tutorial/advanced/motrix_deploy.md diff --git a/docs/source/en/user_guide/tutorial/basic_frame.md b/docs/source/en/user_guide/tutorial/basic_frame.md index b012acea..fdfac1f0 100644 --- a/docs/source/en/user_guide/tutorial/basic_frame.md +++ b/docs/source/en/user_guide/tutorial/basic_frame.md @@ -1,119 +1,78 @@ # Basic Framework -MotrixLab separates environment implementation, training methods, configuration, and command-line orchestration. This section explains how those pieces fit together before you build a custom environment or training backend. - -## Repository Layers - -```text -MotrixLab/ -├── motrix_envs/ # Environment configs, implementations, and registry -│ └── src/motrix_envs/ -├── motrix_rl/ # RL frameworks, providers, trainers, and run artifacts -│ └── src/motrix_rl/ -├── configs/ -│ ├── algo_base/ # Complete typed defaults for each RL provider -│ └── task// # Per-environment training recipes -└── scripts/ - ├── train.py # Hydra training entry point - ├── play.py # Metadata-backed policy playback - └── view.py # Random-action environment preview -``` - -The main runtime flow is: - -```text -task=/. - │ - ▼ -Hydra composes root config + algorithm base + Task + CLI overrides - │ - ▼ -runner resolves an AgentProvider and creates a Trainer - │ - ▼ -Trainer creates the registered environment and executes train/play - │ - ▼ -runs/... stores metadata, resolved Task config, and checkpoints -``` - -## Core Components - -### Environment Layer - -An environment normally consists of: - -- An `EnvCfg` dataclass registered with `@registry.envcfg("name")`. -- An environment implementation registered with `@registry.env("name")`. -- Task logic for observations, rewards, termination, reset, and action application. - -The environment registry owns environment names and simulation-backend implementations. `scripts/view.py`, trainers, and playback all create environments through this same registry. +MotrixLab separates environment implementation, simulation backends, training methods, configuration, and +command-line orchestration into independent layers. This page unfolds top-down: first the typical +simulation-RL loop, then which part of that loop each MotrixLab package covers, and finally the +capabilities of each part. -### RL Framework and Provider Layer +## The typical simulation RL loop -`RlFramework` defines an RL framework namespace such as `skrl`, `rslrl`, or `motrix`. Each framework contains one or more `AgentProvider` implementations. A provider declares: +Simulation reinforcement learning is driven by the interaction loop between a **policy** and an +**environment**, with a **trainer** updating the policy on top of that loop: -- Its algorithm name, such as `ppo` or `fastsac`. -- Its training backend, such as `jax` or `torch`. -- The typed algorithm configuration schema it accepts. -- Its checkpoint format and how to create a trainer. - -Frameworks and providers are registered in Python because they represent executable capabilities. See [Adding a Custom Training Backend](custom_training_backend.md) for the extension interface. - -### Hydra Configuration Layer - -Training values live in YAML rather than Python Task subclasses: +```{image} /_static/images/tutorial/rl-loop-light.svg +:alt: The simulation RL loop: the policy outputs actions to the environment, the environment returns observations, rewards, and termination flags, and the trainer collects transitions and updates the policy parameters +:class: only-light +``` -- `configs/algo_base/..yaml` supplies the complete provider-owned algorithm defaults. -- `configs/task//..yaml` selects an environment and stores task-specific tuning. -- An optional `..yaml` Task contains only backend-specific differences. -- CLI `key=value` arguments apply temporary overrides after composition. +```{image} /_static/images/tutorial/rl-loop-dark.svg +:alt: The simulation RL loop: the policy outputs actions to the environment, the environment returns observations, rewards, and termination flags, and the trainer collects transitions and updates the policy parameters +:class: only-dark +``` -The provider's dataclass schema validates field names and types, while YAML remains the source of truth for values. Task files are discovered by scanning `configs/task/`; there is no RL configuration decorator or Python Task registry. +- The policy takes the environment observation oₜ as input and outputs an action aₜ. +- The environment runs N simulated instances in parallel, applies the action, advances the physics, and + produces the reward rₜ, terminated (failure) / truncated (time limit), and the next observation oₜ₊₁. +- The next observation oₜ₊₁ returns to the policy as the input for the next action; the reward and the + termination flags **never enter the policy network** — together with (oₜ, aₜ) they form the transitions + handed to the trainer, which updates the policy parameters with the selected RL algorithm, repeating + until convergence. -### Runner and Trainer Layer +One timing detail: the observation the environment produces in response to aₜ becomes the policy's +input at step t+1; both interaction edges in the figure are labeled with the current cycle's oₜ, aₜ, rₜ. +With this loop in mind, every MotrixLab package has a place in the figure. -The shared runner handles framework-neutral orchestration: +## Which part of the loop each package covers -1. Read `task.env`, `task.rllib`, `task.algo`, and `task.train_backend` from the composed config. -2. Resolve a compatible provider and training backend. -3. Create a run directory and write `metadata.json` plus `task_config.yaml`. -4. Build a `TrainerContext` and ask the provider to create its trainer. -5. Execute training or playback and register checkpoint artifacts. +MotrixLab is a UV workspace made of nine packages, grouped by the loop above: -The trainer owns framework-specific model construction, optimization, checkpoint serialization, and inference. It should use the environment registry instead of coupling itself to a concrete environment class. +| Package | Loop stage | Responsibility | +| ----------------------- | ------------------------------ | -------------------------------------------------------------------- | +| `motrix_env_core` | Policy–environment interaction | Backend-agnostic environment framework: `EnvCfg`, registry, frontends, lifecycle | +| `motrix_envs` | Policy–environment interaction | Built-in environments, robot models, and task assets | +| `motrix_rl` | Policy training and updates | RL framework integrations (providers, trainers) and training tools | +| `configs/`, `scripts/` | Configuration and orchestration | Hydra algorithm base configs and Task recipes; train / play / view / export entry points | +| `motrix_deploy*` | Policy deployment | Framework-agnostic artifacts and runtime contracts, MuJoCo replay and Unitree hardware backends | -## Training Workflow +A simulation backend (such as `motrix_env_motrixsim`) is isolated behind the `SimBackend` interface — +when using an environment you normally do not need to care which one it is. To select or integrate a +backend, see [SimBackend: Decoupling from the Simulator](building_envs/sim_backend.md). -For example: +## One full training run ```bash python scripts/train.py task=cartpole/skrl.ppo num_envs=1024 ``` -This command performs the following steps: - -1. Hydra composes `configs/train.yaml`, `configs/algo_base/skrl.ppo.yaml`, and `configs/task/cartpole/skrl.ppo.yaml`. -2. `num_envs=1024` overrides the composed Task value for this run only. -3. The runner resolves the SKRL PPO provider and an available JAX or Torch backend. -4. The trainer creates the registered `cartpole` environment and starts optimization. -5. Run metadata, the resolved Task snapshot, logs, and checkpoint manifests are written under `runs/cartpole/`. +1. Hydra composes `configs/train.yaml`, `configs/algo_base/skrl.ppo.yaml`, and + `configs/task/cartpole/skrl.ppo.yaml`; `num_envs=1024` only overrides this run. +2. The runner resolves the SKRL PPO provider and automatically picks an available JAX/Torch backend. +3. The trainer creates the registered `cartpole` environment through the registry and starts optimizing. +4. Run metadata, the final Task snapshot, logs, and the checkpoint manifest land in `runs/cartpole/`. -## Multi-Framework Support - -The same environment can have multiple Task recipes without changing its implementation: +The same environment can have several Task recipes without touching the environment implementation: ```bash python scripts/train.py task=cartpole/skrl.ppo python scripts/train.py task=cartpole/rslrl.ppo ``` -SKRL supports JAX and Torch providers, RSLRL uses Torch, and `motrix.fastsac` selects its synchronous or asynchronous Torch trainer through `algo.asynchronous`. The selected Task and provider determine the algorithm configuration and output metadata. - -## Why This Separation Matters +## What the layering buys you 1. **Environment reuse**: one registered environment can be trained by multiple RL frameworks. -2. **Typed configuration**: provider schemas reject misspelled or incompatible YAML/CLI values before training. -3. **Reproducibility**: each run stores the resolved Task configuration and provider identity. -4. **Extensibility**: new environments add registry entries and Task YAML; new RL integrations add providers and trainers. -5. **Consistent artifacts**: playback and resume use metadata and checkpoint manifests instead of guessing file names. +2. **Typed configuration**: provider schemas reject misspelled or mistyped YAML/CLI values before training. +3. **Reproducible experiments**: every run stores its final Task config and provider identity. +4. **Multiple backends**: the backend is a config-level choice; environment implementations never see the + concrete simulator. +5. **Easy extension**: a new environment needs a registration plus a Task YAML; a new RL integration needs + a provider and trainer; a new simulator needs a registered SimBackend. diff --git a/docs/source/en/user_guide/tutorial/building_envs/direct_env.md b/docs/source/en/user_guide/tutorial/building_envs/direct_env.md new file mode 100644 index 00000000..d80c7868 --- /dev/null +++ b/docs/source/en/user_guide/tutorial/building_envs/direct_env.md @@ -0,0 +1,194 @@ +# Writing DirectEnv Environments + +`DirectEnv` is MotrixLab's direct-workflow frontend: the environment implementation holds +the backend-neutral simulator interface `SimBackend` directly (accessed as `self.sim`) +and implements action application, rewards, termination, and observations inside fixed +lifecycle hooks. It suits tasks with highly custom logic that do not need declarative +per-term composition of rewards, observations, and terminations. + +This page starts with a complete minimal example to build the overall picture, then +unfolds the concepts behind it; for when to choose DirectEnv over the manager workflow, +see the [Building Environments Overview](index.md). +## A minimal example + +The built-in `cartpole` environment (full source under +`motrix_envs/src/motrix_envs/basic/cartpole/`) shows the pattern. A DirectEnv environment +consists of a **config class** and an **environment class**, each registered separately: + +```python +import os + +import gymnasium as gym +import numpy as np + +from motrix_env_core import registry +from motrix_env_core.array.env import ArrayEnvState +from motrix_env_core.config import configclass +from motrix_env_core.config.scene import SceneCfg +from motrix_env_core.direct.env import DirectEnv, DirectEnvCfg +from motrix_env_core.sim import DofPositionQuery, DofVelocityQuery +from motrix_env_core.sim.write import CtrlTargetsWrite, JointPositionWrite, JointVelocityWrite + +model_file = os.path.dirname(__file__) + "/cartpole.xml" + +_SIM_DATA_QUERIES = { + "dof_pos": DofPositionQuery(), + "dof_vel": DofVelocityQuery(), +} + + +@registry.envcfg("cartpole") +@configclass +class CartPoleEnvCfg(DirectEnvCfg): + """Move a cart to keep an inverted pendulum upright. + + zh_CN: 移动小车以保持倒立摆直立。 + """ + + scene: SceneCfg = SceneCfg(file=model_file) + max_episode_seconds: float = 10 + reset_noise_scale: float = 0.01 + + +@registry.env("cartpole") +class CartPoleEnv(DirectEnv): + def __init__(self, cfg: CartPoleEnvCfg, num_envs=1, backend: str | None = None): + super().__init__(cfg, num_envs, backend=backend) + self.sim_data = self.sim.compile_reads(_SIM_DATA_QUERIES) + self._ctrl_writes = self.sim.compile_writes({"ctrl": CtrlTargetsWrite()}) + self._reset_program = self.sim.compile_writes( + { + "position": JointPositionWrite(("slider", "hinge")), + "velocity": JointVelocityWrite(("slider", "hinge")), + }, + reset=True, + ) + self._action_space = gym.spaces.Box(-3.0, 3.0, (1,), dtype=np.float32) + self._observation_space = gym.spaces.Box(-np.inf, np.inf, (4,), dtype=np.float32) + + @property + def observation_space(self): + return self._observation_space + + @property + def action_space(self): + return self._action_space + + def reset(self, env_ids: np.ndarray): + rows = len(env_ids) + scale = self._cfg.reset_noise_scale + self._reset_program.buffer("position")[env_ids] = np.random.uniform( + -scale, scale, (rows, 2) + ).astype(np.float32) + self._reset_program.buffer("velocity")[env_ids] = np.random.uniform( + -scale, scale, (rows, 2) + ).astype(np.float32) + self._reset_program.execute(env_ids) + self.sim_data.execute(env_ids) + return {} + + def apply_action(self, actions: np.ndarray, state: ArrayEnvState): + self._ctrl_writes.buffer("ctrl")[:] = actions.astype(np.float32, copy=False) + self._ctrl_writes.execute() + return state + + def compute_transition(self, state: ArrayEnvState): + self.sim_data.execute() + dof_pos = self.sim_data["dof_pos"] + cart_pos, angle = dof_pos[:, 0], dof_pos[:, 1] + state.reward = np.ones((self.num_envs,), dtype=np.float32) + state.terminated = ( + np.isnan(angle) + | (np.abs(angle) > 0.2) + | (cart_pos < -0.8) + | (cart_pos > 0.8) + ) + return state + + def compute_observation(self, state: ArrayEnvState): + obs = np.concatenate([self.sim_data["dof_pos"], self.sim_data["dof_vel"]], axis=-1) + return state.replace(obs=obs) +``` + +Once registered, preview it with `python scripts/view.py env=cartpole` and start +training by creating a Task as described in +[Task Configuration](../training/task_config.md). + +Three core elements of DirectEnv are visible in the example: + +1. The **config** inherits `DirectEnvCfg` and declares the scene and task parameters; +2. The **environment class** compiles queries and read/write programs against `self.sim` + in its constructor — see [SimBackend: Decoupling from the Simulator](sim_backend.md); +3. The **task logic** lives in the four lifecycle hooks `reset` / `apply_action` / + `compute_transition` / `compute_observation`. + +The next section builds the big picture with the lifecycle figure, before the details of +the config and the SimBackend programs. + +## Lifecycle + +`DirectEnv` inherits `ArrayEnv`: the step / auto-reset / truncation lifecycle is owned by +the base class and must not be re-implemented in subclasses. In the `step(actions)` +sequence below, the four orange-highlighted stages are exactly what a direct workflow +fills in: + +```{image} /_static/images/tutorial/direct-env-lifecycle-light.svg +:alt: DirectEnv single control step pipeline: apply_action (hook), physics_step, compute_transition (hook), truncation check, auto-reset (reset(env_ids)), compute_observation (hook) +:class: only-light +``` + +```{image} /_static/images/tutorial/direct-env-lifecycle-dark.svg +:alt: DirectEnv single control step pipeline: apply_action (hook), physics_step, compute_transition (hook), truncation check, auto-reset (reset(env_ids)), compute_observation (hook) +:class: only-dark +``` + +Hooks a subclass implements: + +| Hook | Responsibility | +| ------------------------------------------- | -------------------------------------------------------------------------------------- | +| `reset(env_ids)` | Write reset state (randomized initial poses, ...) for the selected rows and return an info dict; observations are produced afterwards by `compute_observation` | +| `apply_action(actions, state)` | Write the action into the simulator (usually ctrl targets) | +| `compute_transition(state)` | Execute the read program and derive `state.reward`, `state.terminated`, ...; this is the only full data refresh of a step and must **not** write `state.obs` | +| `compute_observation(state)` | Assemble `state.obs` purely from already refreshed simulator data, without further reads | +| `observation_space` / `action_space` | Defined as properties, usually pre-built in `__init__` | + +Semantics: + +- `terminated` marks episode-ending conditions such as task failure; `truncated` marks + the time limit at `max_episode_steps`. `ArrayEnv` combines both into `done` and + triggers auto-reset; `info["time_outs"]` flags rows that timed out without failing. +- The environment dimension must use vectorized NumPy operations; plain loops are only + allowed over a fixed number of joints, feet, or terms. +- Constants (initial poses, space definitions, query names) are precomputed in + `__init__` or config construction, not recreated inside the step loop. + + +## Configuration + +`DirectEnvCfg` inherits `EnvCfg`; direct-workflow configs are written as +`class MyEnvCfg(..., DirectEnvCfg)`. `EnvCfg` provides these common fields: + +| Field | Meaning | +| --------------------- | -------------------------------------------------------------------- | +| `scene` | `SceneCfg`, the scene and model source; required, `validate()` fails otherwise | +| `sim` | `SimCfg`, simulation parameters (`dt`, `solver_iterations`, ...) | +| `ctrl_dt` | Control step in seconds, default `0.01` | +| `max_episode_seconds` | Maximum episode length in seconds, converted to `max_episode_steps = max_episode_seconds / ctrl_dt`; `None` means unlimited | +| `render_spacing` | Spacing between environments in grid rendering | + +Each control step advances the simulation by `sim_substeps = round(ctrl_dt / sim.dt)` +physics substeps, so `sim.dt` must be less than or equal to `ctrl_dt`. Model files and +simulation parameters are covered in [Setting Up the Physics Scene](scene.md). + + + +## Registration rules + +- **The config is registered before the environment class**: `@registry.env("name")` + requires a config with the same name registered via `@registry.envcfg("name")`. +- The config registration target may be a config class or a zero-argument factory with + a return type annotation (useful for returning customized config instances). +- The registry infers the frontend type from the class hierarchy (`DirectEnv` belongs + to the `"np"` data backend); never specify it manually. +- The first docstring line and the `zh_CN:` line of the config serve as the + environment's English and Chinese descriptions shown in the environment overview. diff --git a/docs/source/en/user_guide/tutorial/building_envs/index.md b/docs/source/en/user_guide/tutorial/building_envs/index.md new file mode 100644 index 00000000..2aa49cd0 --- /dev/null +++ b/docs/source/en/user_guide/tutorial/building_envs/index.md @@ -0,0 +1,100 @@ +# Building Environments Overview + +This section explains how to write your own reinforcement learning environment in +MotrixLab. Before the details, three building blocks of the big picture. +## Anatomy of an environment + +An environment consists of two parts, each registered under the environment name: + +- **Config class**: a `@configclass` data class declaring the scene, simulation + parameters, and task parameters, registered via `@registry.envcfg("name")`; +- **Environment class**: implements the task logic, registered via + `@registry.env("name")`. + +Once registered, `scripts/view.py`, the trainers, and the replay flow all create the +environment by name through the same registry. + + +## Lifecycle overview + +Both workflows share the same vectorized environment lifecycle; `step(actions)` follows +a fixed sequence: + +```{image} /_static/images/tutorial/env-lifecycle-light.svg +:alt: ArrayEnv single control step pipeline: apply_action, physics_step, compute_transition, truncation check, auto-reset, compute_observation +:class: only-light +``` + +```{image} /_static/images/tutorial/env-lifecycle-dark.svg +:alt: ArrayEnv single control step pipeline: apply_action, physics_step, compute_transition, truncation check, auto-reset, compute_observation +:class: only-dark +``` + +Every stage is orchestrated by the `ArrayEnv` base class; environment implementations only fill in +hooks and must not re-implement the lifecycle. Semantics: + +- `terminated` marks episode-ending conditions such as task failure; `truncated` marks + the time limit, and `info["time_outs"]` flags rows that timed out without failing; +- environments that are done are reset automatically at the end of each step, and + observations are recomputed after the reset. + +Who implements each stage in the two workflows: DirectEnv hook by hook in +[Writing DirectEnv Environments](direct_env.md#lifecycle), and ManagerEnv driven by the config +groups in [Writing ManagerEnv Environments](manager_env.md#lifecycle). + + +## Choosing between the two workflows + +MotrixLab provides two environment workflows; the difference is where the task logic +lives: + +| Aspect | DirectEnv (direct workflow) | ManagerEnv (manager workflow) | +| ----------- | -------------------------------------- | ---------------------------------------------------- | +| Config base | `DirectEnvCfg` | `ManagerBasedEnvCfg` | +| Task logic | Hand-written inside environment hooks | Declared per term in config; compiled into a kernel | +| Rewards | Hand-written array math in `compute_transition` | `rewards` / `terminations` config groups | +| Observations| Hand-written assembly in `compute_observation` | `observations` config groups (`policy`/`value`) | +| Reset logic | Hand-written `reset(env_ids)` override | `sim_reset` config group + command reset hooks | +| Typical use | Simple or highly customized tasks | Compositional tasks such as locomotion | + +For a simple task whose logic fits on one screen, the direct workflow is the least +ceremony; once reward, observation, and termination terms multiply and need reuse across +environments, the manager workflow's declarative composition is easier to maintain. + + +## Rewards and terminations + +The two workflows only differ in how rewards are **written**, not in how to design them: + +- DirectEnv hand-writes the reward array math in `compute_transition`; ManagerEnv + declares terms in the `rewards` config group and sums them weighted by `weight`. + Terminations work the same way, as the `terminated` mask in `compute_transition` or + the `terminations` config group. +- Design guidance: let each reward term own one clear objective; prefer smooth + functions (for example exponential decay) over hard thresholds; keep weights in + config so terms can be tuned one by one; check for reward loopholes that allow high + reward without the intended behavior. + +For concrete examples, see the built-in environments and the rewards-and-terminations +section of [Writing ManagerEnv Environments](manager_env.md). + + +## Reading path + +1. [Writing DirectEnv Environments](direct_env.md): understand the environment skeleton + from a minimal example; +2. [Writing ManagerEnv Environments](manager_env.md): switch to the manager workflow + when you need declarative composition; +3. [SceneCfg: Setting Up the Physics Scene](scene.md): model files and simulation + parameters in detail; +4. [SimBackend: Decoupling from the Simulator](sim_backend.md): the environment–simulator + boundary, backend selection and integration. + +```{toctree} +:hidden: + +direct_env +manager_env +scene +sim_backend +``` diff --git a/docs/source/en/user_guide/tutorial/building_envs/manager_env.md b/docs/source/en/user_guide/tutorial/building_envs/manager_env.md new file mode 100644 index 00000000..94b8b7ce --- /dev/null +++ b/docs/source/en/user_guide/tutorial/building_envs/manager_env.md @@ -0,0 +1,314 @@ +# Writing ManagerEnv Environments + +`ManagerEnv` is the manager-workflow frontend: actions, commands, resets, observations, +rewards, and terminations are not hand-written in the environment class. Instead they are +declared term by term in the config groups of `ManagerBasedEnvCfg`, and the manager +compiler generates a fused Numba task program from them. It suits tasks with many terms +that reuse composable mdp terms, such as locomotion and whole-body tracking. + +For choosing between the workflows, see the +[Building Environments Overview](index.md). This page first shows a minimal skeleton to +establish the whole picture, then walks through each term group: what it does and how to +define one. + +## A minimal skeleton + +The manager workflow does not subclass `ManagerEnv`: write a `ManagerBasedEnvCfg` +config, then register the generic environment class under the environment name: + +```python +from motrix_env_core import registry +from motrix_env_core.config import configclass +from motrix_env_core.config.scene import SceneCfg +from motrix_env_core.manager import ( + ManagerBasedEnvCfg, + ManagerEnv, +) + +# —— the config groups are declared below and omitted here for brevity —— + +@registry.envcfg("my-task") +@configclass +class MyTaskEnvCfg(ManagerBasedEnvCfg): + scene: SceneCfg = SceneCfg(file="my_model.xml") + actions: ActionsCfg = ActionsCfg() # at least one action term is required + observations: ObservationsCfg = ObservationsCfg() # the policy group is required + rewards: RewardsCfg = RewardsCfg() + terminations: TerminationsCfg = TerminationsCfg() + sim_reset: ManagerResetCfg = ResetCfg() + + +registry.env("my-task")(ManagerEnv) +``` + +Each config group is a `@configclass` data class whose field names are the term names +and whose values are term config objects. The sections below cover each group: what the +term does and how to define it. + +## Lifecycle + +`ManagerEnv` also inherits `ArrayEnv`; the step sequence and the auto-reset / truncation +semantics match DirectEnv. The figure below integrates the term callbacks into the +pipeline: green parts run inside manager kernels (`@dispatch`), purple parts are the +manager's host-side Python callbacks, and blue parts are the framework — the +orchestrating stages and its own data plumbing: + +```{image} /_static/images/tutorial/manager-env-lifecycle-light.svg +:alt: ManagerEnv single control step pipeline with a done-row branch: apply_action → physics_step → compute_transition → truncation check → done? splits done rows into auto-reset microflow and alive rows toward compute_observation +:class: only-light +``` + +```{image} /_static/images/tutorial/manager-env-lifecycle-dark.svg +:alt: ManagerEnv single control step pipeline with a done-row branch: apply_action → physics_step → compute_transition → truncation check → done? splits done rows into auto-reset microflow and alive rows toward compute_observation +:class: only-dark +``` + +Two differences from DirectEnv: term callbacks are not hand-written in the environment +class — they are declared in the config groups and compiled into fused kernels by the +manager compiler; kernel compilation and warm-up happen at the first `init_state()` +(startup), never inside the step loop. + +## Action terms (actions) + +**What they do**: each control step, an action term transforms the policy's action batch +into actuator controls and writes them to the simulator. Each term's action space +(one-dimensional `Box`, float32) is concatenated in declaration order into the +environment's `action_space`, and the policy action is split the same way when +distributed to the terms. + +**How to define one**: an action term has two parts — + +1. An `ActionCfg` subclass declaring static parameters and the actuator route; its + `__call__(env, actuators)` returns the runtime term. With `actuator_names` an empty + tuple `()` controls all actuators, explicit names control only the listed ones, and + one actuator may belong to a single term. `None` means the term writes no actuator + targets. +2. An `ActionTerm` subclass decorated with `@kernel_data`: it holds persistent state + (`np.ndarray` fields, one row per environment) and shared model data (`SharedArray` + fields), and implements three methods — `action_space()` returns this term's action + space, `process(actions)` transforms the action batch into routed actuator controls, + and `reset(env_ids)` clears persistent state. + +```python +@configclass(kw_only=True) +class MyJointActionCfg(ActionCfg): + actuator_names: tuple[str, ...] = () + scale: float = 0.5 + + def __call__(self, env: ManagerEnv, actuators) -> ActionTerm: + return MyJointAction(scale=np.float32(self.scale), num_envs=env.num_envs, + actuators=actuators) + + +@kernel_data +class MyJointAction(ActionTerm): + current: np.ndarray # persistent state: previous action (num_envs rows) + scale: SharedArray # static parameter: passed in by __call__ + + def action_space(self, env, actuators) -> gym.spaces.Box: + ... # derived from the actuators' ctrl_range + + def process(self, actions: np.ndarray) -> np.ndarray: + ... # transform actions, return this term's routed actuator controls + + def reset(self, env_ids: np.ndarray) -> None: + self.current[env_ids] = 0.0 +``` + +## Observation terms (observations) + +**What they do**: assemble simulator state and commands into the inputs of the policy +and value networks. Only the `policy` and `value` groups exist, and `policy` is +required: `policy` feeds the actor (noise allowed), `value` feeds the critic (usually +noise-free and may include privileged information). Term outputs are concatenated in +declaration order into the observation vector. + +**How to define one**: an observation term is an `ObservationTermCfg` subclass whose +`__call__(env)` returns `ObsTerm(size, dispatch, *args)`. `size` is the term's output +width; the dispatch has the shape `def xxx_obs(ctx, out, *args) -> None` and writes the +term's observation into `out` — `ctx.sim["key"]` reads data declared in `queries`, and +noise is passed in as a scalar argument and added inside the dispatch. + +```python +@dispatch +def projected_gravity_obs(ctx: ManagerContext, out: np.ndarray, + noise_amplitude: np.float32) -> None: + ... # compute and write out (width 3) + add_uniform_noise(out, noise_amplitude, ctx.rand.state) + + +@configclass(kw_only=True) +class ProjectedGravityObsCfg(ObservationTermCfg): + noise: UniformNoiseCfg = UniformNoiseCfg() + + def __call__(self, env: ManagerEnv) -> ObsTerm: + return ObsTerm(3, projected_gravity_obs, np.float32(self.noise.amplitude)) +``` + +Simulation data an observation depends on is declared through the environment's +`queries` group or the term's `required_sim_queries()` — see the Queries section below. + +## Reward terms (rewards) + +**What they do**: each control step, a reward term scores the current state with a +scalar; the total reward is the sum of the terms weighted by `weight` — a positive +`weight` rewards, a negative one penalizes. + +**How to define one**: a reward term is a `RewardTermCfg` subclass (the base carries +`weight: float`) whose `__call__(env)` returns `RewardTerm(dispatch, *args)`. The +dispatch has the shape `def xxx_reward(ctx, *args) -> float` and reads declared data +through `ctx.sim["key"]`. + +```python +@dispatch +def base_height_reward(ctx: ManagerContext, target_z: np.float32, + sigma: np.float32) -> float: + error = ctx.sim["robot_base_pos"][2] - target_z + return math.exp(-(error * error) / (sigma * sigma)) + + +@configclass(kw_only=True) +class BaseHeightRewardCfg(RewardTermCfg): + target_z: float + sigma: float + + def __call__(self, env: ManagerEnv) -> RewardTerm: + return RewardTerm(base_height_reward, + np.float32(self.target_z), np.float32(self.sigma)) +``` + +## Termination terms (terminations) + +**What they do**: each control step, a termination term outputs a boolean per +environment; any true term ends the episode as `terminated` (as opposed to the +time-limit `truncated`). + +**How to define one**: a termination term is a `TerminationTermCfg` subclass whose +`__call__(env)` returns `TerminationTerm(dispatch, *args, metric_names=(...))`. The +dispatch has the shape `def xxx_termination(ctx, *args) -> bool`; the optional +`metric_names` register per-environment quantities the dispatch writes into +`ctx.metrics` as readable metrics. + +```python +@dispatch +def bad_orientation_termination(ctx: ManagerContext, threshold: np.float32) -> bool: + ... + return tilt_sq > threshold * threshold + + +@configclass(kw_only=True) +class BadOrientationTerminationCfg(TerminationTermCfg): + threshold: float + + def __call__(self, env: ManagerEnv) -> TerminationTerm: + return TerminationTerm(bad_orientation_termination, + np.float32(self.threshold), + metric_names=("base_tilt",)) +``` + +## Reset terms (sim_reset) + +**What they do**: on episode resets (and on mid-transition recomputation requested by a +command term, see the commands section) they rewrite the simulator state of the selected +environment rows — initial poses, velocities, joint angles, ... — executed by the reset +kernel in declaration order. + +**How to define one**: a reset term is a `ResetTermCfg` subclass whose `__call__(env)` +returns `ResetTerm(dispatch, *args, writes={...})`. `writes` declares which simulator +state the term writes (`BodyPositionWrite`, `JointVelocityWrite`, ...; the dict keys are +the write channels inside the dispatch); the dispatch computes the values to write. + +```python +@dispatch +def _reset_body_pos(ctx: ManagerContext, sim_writes, spawn, noise) -> None: + position = sim_writes["position"] + ... # write the randomized initial position + + +@configclass(kw_only=True) +class BodyPosResetCfg(ResetTermCfg): + spawn: tuple[float, float, float] + noise: tuple[float, float, float] = (0.02, 0.02, 0.005) + + def __call__(self, env: ManagerEnv) -> ResetTerm: + return ResetTerm( + _reset_body_pos, + tuple(np.asarray(self.spawn, dtype=np.float32)), + tuple(np.asarray(self.noise, dtype=np.float32)), + writes={"position": BodyPositionWrite((base_link,))}, + ) +``` + +## Command terms (commands) + +**What they do**: maintain command state and advance it every cycle — velocity command +sampling, curriculum, reference-motion progression, ... A command term also produces +per-environment metrics: write `ctx.metrics["name"]` inside a dispatch and read it back +through `env.state.metrics` / `env.metrics`. Most tasks can reuse built-in commands and +never need a custom one. + +**How to define one**: a command term is a `CommandCfg` subclass whose `__call__(env)` +returns a `CommandTerm` subclass decorated with `@kernel_data`, implementing four +callbacks: + +| Callback | Runs in | Responsibility | +| ----------------- | ----------------- | ------------------------------------------------------------ | +| `update(ctx)` | evaluate kernel | update this cycle's derived command data | +| `advance(ctx)` | transition kernel | advance persistent command state; setting `ctx.sim_reset_requested` requests the row's simulator state be recomputed at the end of the step | +| `reset(ctx)` | host side | prepare host data for the rows being reset (including termination statistics) | +| `reset_env(ctx)` | reset kernel | reset persistent command state | + +## Queries (queries) + +**What they do**: declare the simulation data and structural information terms need; +the framework compiles the declarations into batched reads. `queries.data` holds +simulation data queries (`JointPositionQuery`, `LinkPositionQuery`, ...) and +`queries.model` holds model queries (`ActuatorKpQuery`, +`BodyJointPositionLimitsQuery`, ...). + +**How to declare them**: usually filled in the config's `__post_init__` from the robot +configuration: + +```python +def __post_init__(self) -> None: + joint_names = ... # resolved from scene.objs.robot + base_link = ... # base link name + self.queries.data["robot_dof_pos"] = JointPositionQuery(joints=joint_names) + self.queries.model["robot_joint_position_limits"] = BodyJointPositionLimitsQuery(body=base_link) +``` + +Observation terms can also declare their needs via `required_sim_queries(env_cfg)`; +all declarations merge by key, and declarations sharing a key must be exactly equal or +the config fails. How queries are compiled and executed is covered in +[SimBackend: Decoupling from the Simulator](sim_backend.md). + +## Registration and variants + +- **The config is registered before the environment class**; the registry infers the + frontend type from the class hierarchy (`ManagerEnv` and `DirectEnv` both belong to + the `"np"` data backend), so never specify it manually. +- The first docstring line and the `zh_CN:` line serve as the environment's English and + Chinese descriptions. +- Multiple presets of one environment are derived through config factories: a variant + overrides only what differs from the shared config (for example lowering one reward + weight). + +The built-in `microduck-ball-balance` environment (source under +`motrix_envs/src/motrix_envs/locomotion/ball_balance/microduck.py`) shows the factory +form: + +```python +@registry.envcfg("microduck-ball-balance") +def make_microduck_ball_balance_cfg() -> MicroduckBallBalanceEnvCfg: + """Balance on top of a basketball with Microduck. + + zh_CN: 让 Microduck 双脚站在篮球上并保持平衡。 + """ + return MicroduckBallBalanceEnvCfg() + + +registry.env("microduck-ball-balance")(ManagerEnv) +``` + +Once registered, preview it with `python scripts/view.py env=` and start training +by creating a Task as described in [Task Configuration](../training/task_config.md). diff --git a/docs/source/en/user_guide/tutorial/physics_environment.md b/docs/source/en/user_guide/tutorial/building_envs/scene.md similarity index 94% rename from docs/source/en/user_guide/tutorial/physics_environment.md rename to docs/source/en/user_guide/tutorial/building_envs/scene.md index fef7588f..11f866a0 100644 --- a/docs/source/en/user_guide/tutorial/physics_environment.md +++ b/docs/source/en/user_guide/tutorial/building_envs/scene.md @@ -1,6 +1,6 @@ -# Physics Environment Configuration +# SceneCfg: Setting Up the Physics Scene -Physics environment configuration defines simulation parameters and scene settings in reinforcement learning training. +The physics scene defines simulation parameters and scene settings for reinforcement learning training. MotrixLab uses [MotrixSim](https://motrixsim.readthedocs.io/en/latest/user_guide/index.html) as the physics simulation backend. ## Supported File Formats diff --git a/docs/source/en/user_guide/tutorial/building_envs/sim_backend.md b/docs/source/en/user_guide/tutorial/building_envs/sim_backend.md new file mode 100644 index 00000000..2778731a --- /dev/null +++ b/docs/source/en/user_guide/tutorial/building_envs/sim_backend.md @@ -0,0 +1,32 @@ +# SimBackend: Decoupling from the Simulator + +All simulation interaction in `DirectEnv` and `ManagerEnv` goes through the backend-neutral +`SimBackend` interface (held as `self.sim`). No concrete simulator type crosses the core boundary — +the same environment code runs on any backend implementing that interface. + +## Compiling three kinds of programs + +Environments compile three kinds of programs against `self.sim` at construction; each step then only +writes small buffers and calls `execute()`: + +- `compile_model(queries)`: model queries — structural information that does not change during + simulation (actuator lists, joint limits, ...); +- `compile_reads(queries)`: simulation data reads (joint positions, velocities, link poses, ...); + each `execute(env_ids)` refreshes the batch, and results are indexed by the keys used at declaration; +- `compile_writes({name: SimWrite})`: write programs (carried by the `write_compiler`). Control targets use `CtrlTargetsWrite`; + reset writes (initial poses, ...) pass `reset=True`. Fill `buffer(name)` first, then call + `execute(env_ids)`. + +Division of labor between the workflows: a `DirectEnv` compiles these programs in its own constructor +(see the minimal example in [Writing DirectEnv Environments](direct_env.md)); a `ManagerEnv` generates +them from the config's `queries` group and its terms, so users normally never touch them directly +(see [Writing ManagerEnv Environments](manager_env.md)). + +## Choosing and registering a backend + +- The backend is a construction-time string: creating an environment with `backend=None` selects the + registered default backend (currently `motrixsim`); pass a backend name explicitly, or use + `registry.make(..., sim="motrixsim")`. +- Backends are lazily discovered through the `motrix_env.sim_backends` entry-point group: a third-party + simulator plugs in by registering a `SimBackend` factory in that group, at zero import cost. +- `motrix_env_mujoco` only compiles MuJoCo scenes and is not a training backend. diff --git a/docs/source/en/user_guide/tutorial/index.md b/docs/source/en/user_guide/tutorial/index.md new file mode 100644 index 00000000..20f3b7f9 --- /dev/null +++ b/docs/source/en/user_guide/tutorial/index.md @@ -0,0 +1,20 @@ +# Tutorial Overview + +This section is organized to build the big picture first and unfold details +progressively; the sidebar on the left shows the full reading structure: + +- **Basic Framework**: MotrixLab's layered architecture and how environments, RL + providers, Hydra configs, and the Trainer cooperate. +- **Building Environments**: write environments first (the DirectEnv and ManagerEnv + guides), then the two cross-cutting concepts — SceneCfg (scene and simulation config) + and SimBackend (decoupling from the simulator); start from the overview (anatomy, + lifecycle, workflow choice). +- **Training and Results**: from creating a Task config and running training to + analyzing the `runs/` artifacts and checkpoints. +- **Advanced Topics**: standalone capabilities to dig into as needed — ONNX export, + hardware deployment, the command input architecture, and custom training backends. + +If you have not run your first training yet, start with the +[Getting Started](../getting_started/installation.md) guide; if you only want to +implement environments, jump straight to the +[Building Environments Overview](building_envs/index.md). diff --git a/docs/source/en/user_guide/tutorial/rewards.md b/docs/source/en/user_guide/tutorial/rewards.md deleted file mode 100644 index fb75b2f7..00000000 --- a/docs/source/en/user_guide/tutorial/rewards.md +++ /dev/null @@ -1,51 +0,0 @@ -# Reward Function Design - -The reward function tells the agent what behaviors are desired and is a core part of reinforcement learning environment design. - -## Position of Reward Function in Training Loop - -In MotrixLab's `DirectEnv` workflow, reward calculation occurs in the `compute_transition` phase of the `step` function: - -```python -# Execution flow of DirectEnv.step() -def step(self, actions: np.ndarray) -> ArrayEnvState: - # 1. Preparation phase: Clear rewards and state - self._prev_physics_step() # reward = 0.0, terminated = False, truncated = False - - # 2. Apply actions - self._state = self.apply_action(actions, self._state) - - # 3. Physics simulation - self.physics_step() # Execute physics simulation - - # 4. Compute transition ← Reward function is calculated here - self._state = self.compute_transition(self._state) - - # 5. Post-processing - self._update_truncate() # Check time truncation - self._reset_done_envs() # Reset completed environments - self._state = self.compute_observation(self._state) # Compute post-reset observations - - return self._state -``` - -You need to implement reward calculation logic in the `compute_transition` method of subclasses. For specific reward function design ideas, please refer to the training examples. - -### Reward Component Design Principles - -1. **Separation of Concerns**: Each reward function should handle a specific goal -2. **Weight Configuration**: Manage weights of different components through configuration files -3. **Normalization**: Keep reward values within reasonable ranges -4. **Smoothness**: Avoid hard thresholds, use exponential functions for smooth transitions - -This approach makes reward functions modular, facilitating debugging and adjustment of individual component weights. - -## Design Principles - -1. **Clear Goal Orientation**: Reward functions should directly reflect task goals -2. **Reasonable Reward Range**: Avoid overly large or small reward values to maintain training stability -3. **Balance Exploration and Exploitation**: Appropriately reward behaviors close to goals, avoiding sparse rewards -4. **Avoid Reward Hacking**: Check if agents can obtain high rewards through unintended means -5. **Debug-Friendly**: Output reward decomposition information during development for optimization - -By correctly implementing reward calculation in the `compute_transition` method, you can design effective learning signals for various robot tasks. diff --git a/docs/source/en/user_guide/tutorial/training/index.md b/docs/source/en/user_guide/tutorial/training/index.md new file mode 100644 index 00000000..661cf45b --- /dev/null +++ b/docs/source/en/user_guide/tutorial/training/index.md @@ -0,0 +1,28 @@ +# Training Overview + +Once an environment is registered, training an RL task takes three steps: + +```text +Create a Task config pick the environment, RL framework, algorithm, hyperparameters + ↓ +python scripts/train.py task=... Hydra composes the config and starts training + ↓ +runs/{env}/... stores metadata, logs, and checkpoints +``` + +The pages follow the order of use: + +- [Task Configuration and CLI Overrides](task_config.md): create the Task file, tune + run and algorithm parameters, override values with `key=value`; +- [Running Training and Analyzing Results](training_and_result.md): start training, + read TensorBoard logs, replay policies with play; +- [Training Artifacts: the runs Directory and Checkpoint Structure](runs_and_checkpoints.md): + the `runs/` layout, best-policy selection, and how to resume. + +```{toctree} +:hidden: + +task_config +training_and_result +runs_and_checkpoints +``` diff --git a/docs/source/en/user_guide/tutorial/runs_and_checkpoints.md b/docs/source/en/user_guide/tutorial/training/runs_and_checkpoints.md similarity index 100% rename from docs/source/en/user_guide/tutorial/runs_and_checkpoints.md rename to docs/source/en/user_guide/tutorial/training/runs_and_checkpoints.md diff --git a/docs/source/en/user_guide/tutorial/training_environment_config.md b/docs/source/en/user_guide/tutorial/training/task_config.md similarity index 96% rename from docs/source/en/user_guide/tutorial/training_environment_config.md rename to docs/source/en/user_guide/tutorial/training/task_config.md index aca806f7..63242a57 100644 --- a/docs/source/en/user_guide/tutorial/training_environment_config.md +++ b/docs/source/en/user_guide/tutorial/training/task_config.md @@ -146,7 +146,7 @@ algo: timesteps: 20000 ``` -See [`configs/algo_base/skrl.ppo.yaml`](../../../../configs/algo_base/skrl.ppo.yaml) for the complete field set and per-field comments. +See [`configs/algo_base/skrl.ppo.yaml`](../../../../../configs/algo_base/skrl.ppo.yaml) for the complete field set and per-field comments. ### RSLRL PPO @@ -167,7 +167,7 @@ algo: entropy_coef: 0.005 ``` -See [`configs/algo_base/rslrl.ppo.yaml`](../../../../configs/algo_base/rslrl.ppo.yaml) for the complete field set and per-field comments. +See [`configs/algo_base/rslrl.ppo.yaml`](../../../../../configs/algo_base/rslrl.ppo.yaml) for the complete field set and per-field comments. ### Motrix FastSAC @@ -182,7 +182,7 @@ FastSAC uses one algorithm identity for both execution topologies. Its main fiel - `weight_publish_interval` and `weight_poll_interval`: policy synchronization cadence. - `max_ingest_per_iter` and `idle_sleep_s`: learner ingestion and idle backoff. -See [`configs/algo_base/motrix.fastsac.yaml`](../../../../configs/algo_base/motrix.fastsac.yaml) for the complete field set and per-field comments. +See [`configs/algo_base/motrix.fastsac.yaml`](../../../../../configs/algo_base/motrix.fastsac.yaml) for the complete field set and per-field comments. ## Override parameters from the CLI diff --git a/docs/source/en/user_guide/tutorial/training_and_result.md b/docs/source/en/user_guide/tutorial/training/training_and_result.md similarity index 98% rename from docs/source/en/user_guide/tutorial/training_and_result.md rename to docs/source/en/user_guide/tutorial/training/training_and_result.md index 35279c37..b743457d 100644 --- a/docs/source/en/user_guide/tutorial/training_and_result.md +++ b/docs/source/en/user_guide/tutorial/training/training_and_result.md @@ -26,7 +26,7 @@ Built-in RL methods and their training backends: | `rslrl.ppo` | `torch` | RSLRL PPO | | `motrix.fastsac` | `torch` | FastSAC; `algo.asynchronous` selects the execution topology | -Run `python scripts/train.py --help` to list the Tasks available in the current checkout. See [Task Configuration and CLI Overrides](training_environment_config.md) for the Task file layout and override rules. +Run `python scripts/train.py --help` to list the Tasks available in the current checkout. See [Task Configuration and CLI Overrides](task_config.md) for the Task file layout and override rules. ### Selecting Training and Simulation Backends diff --git a/docs/source/zh_CN/user_guide/envs/whole_body_tracking/index.md b/docs/source/zh_CN/user_guide/envs/whole_body_tracking/index.md index a55bef78..69be464f 100644 --- a/docs/source/zh_CN/user_guide/envs/whole_body_tracking/index.md +++ b/docs/source/zh_CN/user_guide/envs/whole_body_tracking/index.md @@ -165,4 +165,4 @@ python scripts/play.py env=ENV_ID num_envs=16 `play.py` 自动选择该环境最新一次 metadata-backed run 的最佳策略。WBT 的 play 配置从 motion 第 0 帧开始,关闭 reset noise 和自适应采样,并移除 10 秒训练时限;播放到 clip 末尾后从第 0 帧重新开始。训练产物和 checkpoint 选择规则见 -[训练产物:runs 目录与 checkpoint 结构](../../tutorial/runs_and_checkpoints.md)。 +[训练产物:runs 目录与 checkpoint 结构](../../tutorial/training/runs_and_checkpoints.md)。 diff --git a/docs/source/zh_CN/user_guide/getting_started/hello_motrixlab.md b/docs/source/zh_CN/user_guide/getting_started/hello_motrixlab.md index 7760d884..629f73b6 100644 --- a/docs/source/zh_CN/user_guide/getting_started/hello_motrixlab.md +++ b/docs/source/zh_CN/user_guide/getting_started/hello_motrixlab.md @@ -82,5 +82,5 @@ python scripts/play.py env=cartpole policy=/path/to/run/checkpoints/policy-file ## 下一步 - 了解 [基础框架](../tutorial/basic_frame.md) -- 学习 [物理环境配置](../tutorial/physics_environment.md) +- 学习 [物理环境配置](../tutorial/building_envs/scene.md) - 浏览更多 [环境](../envs/index.md) diff --git a/docs/source/zh_CN/user_guide/index.md b/docs/source/zh_CN/user_guide/index.md index 9a11f41b..abe55415 100644 --- a/docs/source/zh_CN/user_guide/index.md +++ b/docs/source/zh_CN/user_guide/index.md @@ -13,26 +13,13 @@ robots ```{toctree} :caption: 使用教程 -:maxdepth: 1 +:maxdepth: 2 +tutorial/index tutorial/basic_frame -tutorial/physics_environment -tutorial/training_environment_config -tutorial/rewards -tutorial/training_and_result -tutorial/runs_and_checkpoints - -``` - -```{toctree} -:caption: 进阶主题 -:maxdepth: 1 - -tutorial/export_onnx -tutorial/motrix_deploy -tutorial/input_devices_and_bindings -tutorial/custom_training_backend - +tutorial/building_envs/index +tutorial/training/index +tutorial/advanced/index ``` ```{toctree} diff --git a/docs/source/zh_CN/user_guide/tutorial/custom_training_backend.md b/docs/source/zh_CN/user_guide/tutorial/advanced/custom_training_backend.md similarity index 100% rename from docs/source/zh_CN/user_guide/tutorial/custom_training_backend.md rename to docs/source/zh_CN/user_guide/tutorial/advanced/custom_training_backend.md diff --git a/docs/source/zh_CN/user_guide/tutorial/export_onnx.md b/docs/source/zh_CN/user_guide/tutorial/advanced/export_onnx.md similarity index 97% rename from docs/source/zh_CN/user_guide/tutorial/export_onnx.md rename to docs/source/zh_CN/user_guide/tutorial/advanced/export_onnx.md index acece4ec..076bc6e4 100644 --- a/docs/source/zh_CN/user_guide/tutorial/export_onnx.md +++ b/docs/source/zh_CN/user_guide/tutorial/advanced/export_onnx.md @@ -14,12 +14,12 @@ MotrixLab 可以把一次训练的最佳策略导出为独立的 ONNX 模型。 | RSL-RL | Torch | PPO | 支持单 `policy` observation group 的 `MLPModel` actor;暂不支持 state-dependent standard deviation | | Motrix | Torch | FastSAC | 导出 actor 的确定性推理路径 | -SKRL/JAX checkpoint 暂不支持导出。导出命令接收完整的 run 目录,不接收脱离 `metadata.json` 和 `checkpoints/manifest.json` 的单个 checkpoint 文件。有关 run 目录和最佳策略的说明,参见[训练产物:runs 目录与 checkpoint 结构](runs_and_checkpoints.md)。 +SKRL/JAX checkpoint 暂不支持导出。导出命令接收完整的 run 目录,不接收脱离 `metadata.json` 和 `checkpoints/manifest.json` 的单个 checkpoint 文件。有关 run 目录和最佳策略的说明,参见[训练产物:runs 目录与 checkpoint 结构](../training/runs_and_checkpoints.md)。 ## 安装依赖 无需额外步骤:ONNX 导出与推理所需的依赖(`onnx`、`onnxruntime`)已包含在 `sh install.sh` -安装的默认运行环境中(参见[安装环境](../getting_started/installation.md))。 +安装的默认运行环境中(参见[安装环境](../../getting_started/installation.md))。 ## 导出模型 diff --git a/docs/source/zh_CN/user_guide/tutorial/advanced/index.md b/docs/source/zh_CN/user_guide/tutorial/advanced/index.md new file mode 100644 index 00000000..090cd0f1 --- /dev/null +++ b/docs/source/zh_CN/user_guide/tutorial/advanced/index.md @@ -0,0 +1,20 @@ +# 进阶主题 + +基础训练流程之外的能力,按需选读: + +- [导出 ONNX 策略](export_onnx.md):把训练好的策略导出为独立的 ONNX 模型, + 用于推理侧集成; +- [Go2 平地行走:从训练到真机部署](motrix_deploy.md):完整演示训练策略、导出 artifact、 + MuJoCo 检查与真机运行; +- [指令输入架构](input_devices_and_bindings.md):把"如何读取输入"与 + "策略需要什么目标"解耦,接入键盘等输入设备与任务指令; +- [扩展自定义训练后端](custom_training_backend.md):接入新的 RL 框架、算法或训练后端。 + +```{toctree} +:hidden: + +export_onnx +motrix_deploy +input_devices_and_bindings +custom_training_backend +``` diff --git a/docs/source/zh_CN/user_guide/tutorial/input_devices_and_bindings.md b/docs/source/zh_CN/user_guide/tutorial/advanced/input_devices_and_bindings.md similarity index 100% rename from docs/source/zh_CN/user_guide/tutorial/input_devices_and_bindings.md rename to docs/source/zh_CN/user_guide/tutorial/advanced/input_devices_and_bindings.md diff --git a/docs/source/zh_CN/user_guide/tutorial/input_devices_and_bindings/built_in.md b/docs/source/zh_CN/user_guide/tutorial/advanced/input_devices_and_bindings/built_in.md similarity index 100% rename from docs/source/zh_CN/user_guide/tutorial/input_devices_and_bindings/built_in.md rename to docs/source/zh_CN/user_guide/tutorial/advanced/input_devices_and_bindings/built_in.md diff --git a/docs/source/zh_CN/user_guide/tutorial/input_devices_and_bindings/extending.md b/docs/source/zh_CN/user_guide/tutorial/advanced/input_devices_and_bindings/extending.md similarity index 100% rename from docs/source/zh_CN/user_guide/tutorial/input_devices_and_bindings/extending.md rename to docs/source/zh_CN/user_guide/tutorial/advanced/input_devices_and_bindings/extending.md diff --git a/docs/source/zh_CN/user_guide/tutorial/motrix_deploy.md b/docs/source/zh_CN/user_guide/tutorial/advanced/motrix_deploy.md similarity index 100% rename from docs/source/zh_CN/user_guide/tutorial/motrix_deploy.md rename to docs/source/zh_CN/user_guide/tutorial/advanced/motrix_deploy.md diff --git a/docs/source/zh_CN/user_guide/tutorial/basic_frame.md b/docs/source/zh_CN/user_guide/tutorial/basic_frame.md index 15cdb6db..7a35d1ae 100644 --- a/docs/source/zh_CN/user_guide/tutorial/basic_frame.md +++ b/docs/source/zh_CN/user_guide/tutorial/basic_frame.md @@ -1,119 +1,71 @@ # 基础框架 -MotrixLab 将环境实现、训练方法、配置和命令行编排分成独立层。本节介绍这些部分如何协作,便于后续开发自定义环境或训练后端。 - -## 仓库分层 - -```text -MotrixLab/ -├── motrix_envs/ # 环境配置、实现与注册表 -│ └── src/motrix_envs/ -├── motrix_rl/ # RL framework、provider、trainer 与训练产物 -│ └── src/motrix_rl/ -├── configs/ -│ ├── algo_base/ # 各 RL provider 的完整类型化默认配置 -│ └── task// # 各环境的训练配方 -└── scripts/ - ├── train.py # Hydra 训练入口 - ├── play.py # 基于 metadata 的策略回放入口 - └── view.py # 随机动作环境预览入口 -``` - -主要运行流程如下: - -```text -task=<环境>/<框架>.<算法> - │ - ▼ -Hydra 组合根配置、算法基础配置、Task 与 CLI 覆盖 - │ - ▼ -runner 解析 AgentProvider 并创建 Trainer - │ - ▼ -Trainer 创建已注册环境并执行 train/play - │ - ▼ -runs/... 保存 metadata、最终 Task 配置与 checkpoint -``` - -## 核心组件 - -### 环境层 - -一个环境通常包含: - -- 使用 `@registry.envcfg("name")` 注册的 `EnvCfg` 数据类。 -- 使用 `@registry.env("name")` 注册的环境实现。 -- observation、reward、termination、reset 和 action application 等任务逻辑。 - -环境注册表维护环境名及其仿真后端实现。`scripts/view.py`、训练器和回放流程都通过同一注册表创建环境。 +MotrixLab 将环境实现、仿真后端、训练方法、配置和命令行编排分成独立层。本页自顶向下展开: +先看仿真强化学习的典型流程,再看每个 MotrixLab package 对应该流程的哪部分,最后逐层介绍各部分的能力。 -### RL Framework 与 Provider 层 +## 仿真强化学习的典型流程 -`RlFramework` 定义 `skrl`、`rslrl`、`motrix` 等 RL 框架命名空间。每个 framework 包含一个或多个 `AgentProvider`。Provider 声明: +仿真强化学习由**策略(Policy)**与**环境(Environment)**的交互循环驱动,**训练器(Trainer)** +在这个循环之上更新策略: -- 算法名,例如 `ppo` 或 `fastsac`。 -- 训练后端,例如 `jax` 或 `torch`。 -- 它接受的类型化算法配置 schema。 -- checkpoint 格式以及如何创建 trainer。 - -Framework 与 provider 代表可执行能力,因此在 Python 中注册。扩展接口见[添加自定义训练后端](custom_training_backend.md)。 - -### Hydra 配置层 - -训练参数保存在 YAML 中,不再使用 Python Task 子类: +```{image} /_static/images/tutorial/rl-loop-light.svg +:alt: 仿真强化学习循环:策略输出动作给环境,环境返回观测、奖励与终止标志,训练器收集转移数据并更新策略参数 +:class: only-light +``` -- `configs/algo_base/..yaml` 提供 provider 所有算法字段的完整默认值。 -- `configs/task//..yaml` 选择环境并保存任务调优参数。 -- 可选的 `..yaml` Task 只保存后端差异。 -- CLI 的 `key=value` 参数在组合完成后应用临时覆盖。 +```{image} /_static/images/tutorial/rl-loop-dark.svg +:alt: 仿真强化学习循环:策略输出动作给环境,环境返回观测、奖励与终止标志,训练器收集转移数据并更新策略参数 +:class: only-dark +``` -Provider 的数据类 schema 负责校验字段名和类型,YAML 是配置值的唯一来源。Task 通过扫描 `configs/task/` 自动发现,不再存在 RL 配置装饰器或 Python Task 注册表。 +- 策略以环境观测 oₜ 为输入,输出动作 aₜ; +- 环境向量化地并行运行 N 个仿真实例,施加动作并推进物理仿真,产生奖励 rₜ、 + terminated(失败终止)/ truncated(超时截断)与下一观测 oₜ₊₁; +- 下一观测 oₜ₊₁ 回到策略,作为下一步动作的输入;奖励与终止标志**不进入策略网络**, + 它们连同 (oₜ, aₜ) 组成转移数据交给训练器,按所选 RL 算法更新策略参数,如此往复直到收敛。 -### Runner 与 Trainer 层 +时序上有一个细节:环境因 aₜ 而产生的观测要到第 t+1 步才成为策略的输入;图中两条交互边统一用 +当前周期的 oₜ、aₜ、rₜ 标注。理解了这个循环,MotrixLab 的每个 package 都能在图中找到自己的位置。 -公共 runner 负责与框架无关的编排: +## MotrixLab package 与流程的对应 -1. 从组合配置读取 `task.env`、`task.rllib`、`task.algo` 和 `task.train_backend`。 -2. 解析兼容的 provider 和训练后端。 -3. 创建 run 目录并写入 `metadata.json` 与 `task_config.yaml`。 -4. 构造 `TrainerContext`,再由 provider 创建 trainer。 -5. 执行训练或回放,并登记 checkpoint artifact。 +MotrixLab 是一个 UV workspace,由九个 package 组成。按上述流程划分: -Trainer 负责框架特有的模型创建、优化、checkpoint 序列化和推理。它应通过环境注册表创建环境,而不是依赖某个具体环境类。 +| Package | 流程环节 | 职责 | +| ------------------- | ------------------ | ---------------------------------------------------------------- | +| `motrix_env_core` | 策略与环境的交互 | backend 无关的环境框架:`EnvCfg`、registry、环境前端与生命周期 | +| `motrix_envs` | 策略与环境的交互 | 内置环境、机器人模型与任务资产 | +| `motrix_rl` | 策略的训练与更新 | RL 框架集成(provider、trainer)与训练工具 | +| `configs/`、`scripts/` | 配置与编排 | Hydra 算法基础配置与 Task 配方;train / play / view / export 入口 | +| `motrix_deploy*` | 策略的部署 | 框架无关 artifact 与运行时契约、MuJoCo 回放与 Unitree 真机后端 | -## 训练流程 +仿真后端(如 `motrix_env_motrixsim`)通过 `SimBackend` 接口隔离在环境框架之下,使用环境时通常 +无需关心它;需要选择或接入仿真后端时,见[SimBackend:与仿真器解耦](building_envs/sim_backend.md)。 -例如: +## 一次训练的完整流程 ```bash python scripts/train.py task=cartpole/skrl.ppo num_envs=1024 ``` -该命令会依次执行: - -1. Hydra 组合 `configs/train.yaml`、`configs/algo_base/skrl.ppo.yaml` 和 `configs/task/cartpole/skrl.ppo.yaml`。 -2. `num_envs=1024` 只覆盖本次运行的 Task 值。 -3. runner 解析 SKRL PPO provider,并从可用的 JAX/Torch 后端中选择一个。 -4. trainer 创建已注册的 `cartpole` 环境并开始优化。 -5. run metadata、最终 Task 快照、日志和 checkpoint manifest 写入 `runs/cartpole/`。 +1. Hydra 组合 `configs/train.yaml`、`configs/algo_base/skrl.ppo.yaml` 和 + `configs/task/cartpole/skrl.ppo.yaml`,`num_envs=1024` 只覆盖本次运行的值。 +2. runner 解析 SKRL PPO provider,并从可用的 JAX/Torch 后端中自动选择。 +3. trainer 通过环境注册表创建 `cartpole` 环境并开始优化。 +4. run metadata、最终 Task 快照、日志和 checkpoint manifest 写入 `runs/cartpole/`。 -## 多框架支持 - -同一个环境可以拥有多份 Task 配方,不需要修改环境实现: +同一环境可以拥有多份 Task 配方,不需要修改环境实现: ```bash python scripts/train.py task=cartpole/skrl.ppo python scripts/train.py task=cartpole/rslrl.ppo ``` -SKRL 提供 JAX 与 Torch provider,RSLRL 使用 Torch;`motrix.fastsac` 通过 `algo.asynchronous` 选择同步或异步 Torch trainer。所选 Task 与 provider 共同决定算法配置和输出 metadata。 - ## 分层带来的优势 1. **环境复用**:一个注册环境可以由多个 RL 框架训练。 2. **类型化配置**:provider schema 会在训练前拒绝拼写错误或类型不兼容的 YAML/CLI 值。 3. **实验可复现**:每个 run 保存最终 Task 配置和 provider 身份。 -4. **易于扩展**:新增环境只需注册环境并添加 Task YAML;新增 RL 集成则添加 provider 与 trainer。 -5. **训练产物一致**:回放和续训依赖 metadata 与 checkpoint manifest,不猜测文件名。 +4. **多后端**:backend 是配置级选择,环境实现不感知具体仿真器。 +5. **易于扩展**:新增环境只需注册环境并添加 Task YAML;新增 RL 集成则添加 provider 与 trainer; + 新增仿真后端则注册新的 SimBackend。 diff --git a/docs/source/zh_CN/user_guide/tutorial/building_envs/direct_env.md b/docs/source/zh_CN/user_guide/tutorial/building_envs/direct_env.md new file mode 100644 index 00000000..42f3c0de --- /dev/null +++ b/docs/source/zh_CN/user_guide/tutorial/building_envs/direct_env.md @@ -0,0 +1,186 @@ +# 编写 DirectEnv 环境 + +`DirectEnv` 是 MotrixLab 的直接工作流前端:环境实现直接持有仿真后端接口 `SimBackend`(通过 `self.sim` 访问), +在固定的生命周期钩子中实现动作施加、奖励、终止与观测。它适合任务逻辑高度定制、 +不需要 reward/observation/termination 逐项声明式组合的环境。 + +本页先给出一个完整的最小示例建立整体印象,再逐层展开背后的概念;何时选择 DirectEnv +而不是 Manager 工作流,见[环境构建总览](index.md)。 +## 最小示例 + +以内置 `cartpole` 环境为例(完整源码见 `motrix_envs/src/motrix_envs/basic/cartpole/`)。 +一个 DirectEnv 环境由**配置类**和**环境类**两部分组成,分别注册后即可训练和预览: + +```python +import os + +import gymnasium as gym +import numpy as np + +from motrix_env_core import registry +from motrix_env_core.array.env import ArrayEnvState +from motrix_env_core.config import configclass +from motrix_env_core.config.scene import SceneCfg +from motrix_env_core.direct.env import DirectEnv, DirectEnvCfg +from motrix_env_core.sim import DofPositionQuery, DofVelocityQuery +from motrix_env_core.sim.write import CtrlTargetsWrite, JointPositionWrite, JointVelocityWrite + +model_file = os.path.dirname(__file__) + "/cartpole.xml" + +_SIM_DATA_QUERIES = { + "dof_pos": DofPositionQuery(), + "dof_vel": DofVelocityQuery(), +} + + +@registry.envcfg("cartpole") +@configclass +class CartPoleEnvCfg(DirectEnvCfg): + """Move a cart to keep an inverted pendulum upright. + + zh_CN: 移动小车以保持倒立摆直立。 + """ + + scene: SceneCfg = SceneCfg(file=model_file) + max_episode_seconds: float = 10 + reset_noise_scale: float = 0.01 + + +@registry.env("cartpole") +class CartPoleEnv(DirectEnv): + def __init__(self, cfg: CartPoleEnvCfg, num_envs=1, backend: str | None = None): + super().__init__(cfg, num_envs, backend=backend) + self.sim_data = self.sim.compile_reads(_SIM_DATA_QUERIES) + self._ctrl_writes = self.sim.compile_writes({"ctrl": CtrlTargetsWrite()}) + self._reset_program = self.sim.compile_writes( + { + "position": JointPositionWrite(("slider", "hinge")), + "velocity": JointVelocityWrite(("slider", "hinge")), + }, + reset=True, + ) + self._action_space = gym.spaces.Box(-3.0, 3.0, (1,), dtype=np.float32) + self._observation_space = gym.spaces.Box(-np.inf, np.inf, (4,), dtype=np.float32) + + @property + def observation_space(self): + return self._observation_space + + @property + def action_space(self): + return self._action_space + + def reset(self, env_ids: np.ndarray): + rows = len(env_ids) + scale = self._cfg.reset_noise_scale + self._reset_program.buffer("position")[env_ids] = np.random.uniform( + -scale, scale, (rows, 2) + ).astype(np.float32) + self._reset_program.buffer("velocity")[env_ids] = np.random.uniform( + -scale, scale, (rows, 2) + ).astype(np.float32) + self._reset_program.execute(env_ids) + self.sim_data.execute(env_ids) + return {} + + def apply_action(self, actions: np.ndarray, state: ArrayEnvState): + self._ctrl_writes.buffer("ctrl")[:] = actions.astype(np.float32, copy=False) + self._ctrl_writes.execute() + return state + + def compute_transition(self, state: ArrayEnvState): + self.sim_data.execute() + dof_pos = self.sim_data["dof_pos"] + cart_pos, angle = dof_pos[:, 0], dof_pos[:, 1] + state.reward = np.ones((self.num_envs,), dtype=np.float32) + state.terminated = ( + np.isnan(angle) + | (np.abs(angle) > 0.2) + | (cart_pos < -0.8) + | (cart_pos > 0.8) + ) + return state + + def compute_observation(self, state: ArrayEnvState): + obs = np.concatenate([self.sim_data["dof_pos"], self.sim_data["dof_vel"]], axis=-1) + return state.replace(obs=obs) +``` + +注册完成后,`python scripts/view.py env=cartpole` 可以预览,按 +[Task 配置](../training/task_config.md)创建训练 Task 即可开始训练。 + +从示例可以读出 DirectEnv 的三个核心要素: + +1. **配置**继承 `DirectEnvCfg`,声明场景与任务参数; +2. **环境类**在构造函数中对 `self.sim` 编译查询与读写程序——这层边界见 + [SimBackend:与仿真器解耦](sim_backend.md); +3. **任务逻辑**写在 `reset` / `apply_action` / `compute_transition` / `compute_observation` + 四个生命周期钩子里。 + +下一节先用生命周期图建立整体认识,再逐节展开配置与 SimBackend 程序的细节。 + +## 生命周期 + +`DirectEnv` 继承 `ArrayEnv`,step / auto-reset / truncation 生命周期由基类管理, +子类不要重复实现。`step(actions)` 的固定时序为,橙色标注的四个阶段 +就是 Direct 工作流需要填充的位置: + +```{image} /_static/images/tutorial/direct-env-lifecycle-light.svg +:alt: DirectEnv 单个控制步流水线:apply_action(hook)→ physics_step → compute_transition(hook)→ 截断判定 → auto-reset(reset(env_ids))→ compute_observation(hook) +:class: only-light +``` + +```{image} /_static/images/tutorial/direct-env-lifecycle-dark.svg +:alt: DirectEnv 单个控制步流水线:apply_action(hook)→ physics_step → compute_transition(hook)→ 截断判定 → auto-reset(reset(env_ids))→ compute_observation(hook) +:class: only-dark +``` + +子类需要实现的钩子: + +| 钩子 | 职责 | +| ------------------------------------ | ---------------------------------------------------------------------------------------- | +| `reset(env_ids)` | 为选中的环境写入重置状态(随机化初始姿态等),返回 info 字典;观测交给后续的 `compute_observation` | +| `apply_action(actions, state)` | 将动作写入仿真(通常是 ctrl 目标) | +| `compute_transition(state)` | 执行读取程序并派生 `state.reward`、`state.terminated` 等;这是每步唯一的全量数据刷新点,**不得**写 `state.obs` | +| `compute_observation(state)` | 纯粹用已刷新的仿真数据拼装 `state.obs`,不再执行读取 | +| `observation_space` / `action_space` | 作为 property 定义,通常在 `__init__` 中预构造 | + +语义约定: + +- `terminated` 表示任务失败等回合终止条件;`truncated` 表示达到 `max_episode_steps` + 的时间截断。两者由 `ArrayEnv` 合成 `done` 并触发 auto-reset; + `info["time_outs"]` 标记"截断但未失败"的行。 +- 环境维度必须使用 NumPy 向量化操作;只有遍历固定数量的关节、脚或 term 时才允许普通循环。 +- 常量(初始姿态、空间定义、query 名称)在 `__init__` 或配置构造阶段预计算, + 不在 step 循环中重复创建。 + + +## 配置 + +`DirectEnvCfg` 继承 `EnvCfg`,配置写成 `class MyEnvCfg(..., DirectEnvCfg)`。 +`EnvCfg` 提供以下公共字段: + +| 字段 | 含义 | +| --------------------- | -------------------------------------------------------------------- | +| `scene` | `SceneCfg`,场景与模型来源;必须配置,否则 `validate()` 失败 | +| `sim` | `SimCfg`,仿真参数(`dt`、`solver_iterations` 等) | +| `ctrl_dt` | 控制步长(秒),默认 `0.01` | +| `max_episode_seconds` | 回合最长时长(秒),换算为 `max_episode_steps = max_episode_seconds / ctrl_dt`;`None` 表示不限制 | +| `render_spacing` | 多环境网格渲染时的间距 | + +每个控制步内仿真推进 `sim_substeps = round(ctrl_dt / sim.dt)` 个物理子步, +因此 `sim.dt` 必须小于等于 `ctrl_dt`。场景文件与仿真参数的详细说明见 +[搭建物理场景](scene.md)。 + + + +## 注册约定 + +- **配置先于环境类注册**:`@registry.env("name")` 要求同名配置已通过 + `@registry.envcfg("name")` 注册,否则报错。 +- 配置注册对象可以是配置类,也可以是带返回类型标注的零参数 factory + (适合返回经过定制的配置实例)。 +- registry 根据环境类的继承关系推断前端类型(`DirectEnv` 属于 `"np"` 数据后端), + 不要手动指定。 +- 配置 docstring 的首行与 `zh_CN:` 行分别作为环境的英文、中文描述, + 展示在环境概览中。 diff --git a/docs/source/zh_CN/user_guide/tutorial/building_envs/index.md b/docs/source/zh_CN/user_guide/tutorial/building_envs/index.md new file mode 100644 index 00000000..076924c3 --- /dev/null +++ b/docs/source/zh_CN/user_guide/tutorial/building_envs/index.md @@ -0,0 +1,85 @@ +# 环境构建总览 + +本节介绍如何在 MotrixLab 中编写自己的强化学习环境。开始之前,先建立三个整体印象。 +## 一个环境由什么组成 + +一个环境通常由两部分组成,分别注册到环境名上: + +- **配置类**:一个 `@configclass` 数据类,声明场景、仿真参数与任务参数, + 通过 `@registry.envcfg("name")` 注册; +- **环境类**:实现任务逻辑,通过 `@registry.env("name")` 注册。 + +注册完成后,`scripts/view.py`、训练器与回放流程都通过同一个注册表按名字创建环境。 + + +## 生命周期总览 + +两种工作流共享同一个向量化环境生命周期,`step(actions)` 的固定时序为: + +```{image} /_static/images/tutorial/env-lifecycle-light.svg +:alt: ArrayEnv 单个控制步流水线:apply_action → physics_step → compute_transition → 截断判定 → auto-reset → compute_observation +:class: only-light +``` + +```{image} /_static/images/tutorial/env-lifecycle-dark.svg +:alt: ArrayEnv 单个控制步流水线:apply_action → physics_step → compute_transition → 截断判定 → auto-reset → compute_observation +:class: only-dark +``` + +所有阶段都由 `ArrayEnv` 基类编排,环境实现只填充钩子,不要重复实现生命周期。语义约定: + +- `terminated` 是任务失败等回合终止条件;`truncated` 是达到回合时长上限的时间截断, + `info["time_outs"]` 标记"截断但未失败"的行; +- done 的环境在每步末尾被自动重置,观测在重置之后重新计算。 + +各阶段在两种工作流中分别由谁实现:DirectEnv 在 +[编写 DirectEnv 环境](direct_env.md#生命周期)中逐钩子展开,ManagerEnv 由 +[编写 ManagerEnv 环境](manager_env.md#生命周期)中的配置组驱动。 + + +## 两种工作流怎么选 + +MotrixLab 提供两种环境工作流,区别在于任务逻辑写在哪里: + +| 维度 | DirectEnv(直接工作流) | ManagerEnv(Manager 工作流) | +| --------- | ----------------------------------- | ----------------------------------------------- | +| 配置基类 | `DirectEnvCfg` | `ManagerBasedEnvCfg` | +| 任务逻辑 | 在环境类钩子中手写 | 在配置中逐项声明,由 manager 编译为 kernel | +| 奖励/终止 | `compute_transition` 中手写数组运算 | `rewards` / `terminations` 配置组 | +| 观测 | `compute_observation` 中手写拼接 | `observations` 配置组(`policy`/`value`) | +| 重置逻辑 | 覆写 `reset(env_ids)` 手写 | `sim_reset` 配置组 + command 重置钩子 | +| 典型场景 | 简单或高度定制的任务 | locomotion 等组合式任务 | + +简单任务的逻辑一目了然,直接工作流最省事;当奖励、观测、终止各项的数量多、需要在多个 +环境间复用时,Manager 工作流的声明式组合更易维护。 + + +## 奖励与终止 + +两种工作流只影响奖励的**写法**,不影响设计思路: + +- DirectEnv 在 `compute_transition` 中手写奖励数组运算;ManagerEnv 在 `rewards` + 配置组逐项声明,最终奖励是各项按 `weight` 的加权和。终止条件同理, + 分别对应 `compute_transition` 中的 `terminated` 掩码与 `terminations` 配置组。 +- 设计上建议:每个奖励项只负责一个明确的目标;用平滑函数(如指数衰减) + 代替硬阈值;把权重放在配置中,便于逐项调试;检查是否存在"刷分"漏洞。 + +具体示例见各内置环境的实现与[编写 ManagerEnv 环境](manager_env.md)的 +奖励与终止一节。 + + +## 阅读路径 + +1. [编写 DirectEnv 环境](direct_env.md):从最小示例理解环境骨架; +2. [编写 ManagerEnv 环境](manager_env.md):需要声明式组合时切换到 Manager 工作流; +3. [SceneCfg:搭建物理场景](scene.md):场景文件与仿真参数的细节; +4. [SimBackend:与仿真器解耦](sim_backend.md):环境与仿真器的边界、backend 选择与接入。 + +```{toctree} +:hidden: + +direct_env +manager_env +scene +sim_backend +``` diff --git a/docs/source/zh_CN/user_guide/tutorial/building_envs/manager_env.md b/docs/source/zh_CN/user_guide/tutorial/building_envs/manager_env.md new file mode 100644 index 00000000..44eef3bd --- /dev/null +++ b/docs/source/zh_CN/user_guide/tutorial/building_envs/manager_env.md @@ -0,0 +1,282 @@ +# 编写 ManagerEnv 环境 + +`ManagerEnv` 是 Manager 工作流前端:动作、命令、重置、观测、奖励与终止不在环境类中手写, +而是作为逐项(term)的声明式配置写入 `ManagerBasedEnvCfg` 的各个配置组,由 manager 编译器 +生成融合的 Numba 任务程序。它适合项数量多、需要组合复用 mdp 项的任务 +(如 locomotion、whole-body tracking)。 + +与 DirectEnv 的选型对比见[环境构建总览](index.md)。本页先用一个最小骨架展示全貌, +再逐组介绍每一项的职责与定义方式。 + +## 最小骨架 + +Manager 工作流不需要子类化 `ManagerEnv`:写好一个 `ManagerBasedEnvCfg` 配置, +再把通用环境类注册到环境名上即可: + +```python +from motrix_env_core import registry +from motrix_env_core.config import configclass +from motrix_env_core.config.scene import SceneCfg +from motrix_env_core.manager import ( + ManagerBasedEnvCfg, + ManagerEnv, +) + +# —— 各配置组的声明见下文,这里先省略 —— + +@registry.envcfg("my-task") +@configclass +class MyTaskEnvCfg(ManagerBasedEnvCfg): + scene: SceneCfg = SceneCfg(file="my_model.xml") + actions: ActionsCfg = ActionsCfg() # 必须至少有一个动作项 + observations: ObservationsCfg = ObservationsCfg() # 必须有 policy 观测组 + rewards: RewardsCfg = RewardsCfg() + terminations: TerminationsCfg = TerminationsCfg() + sim_reset: ManagerResetCfg = ResetCfg() + + +registry.env("my-task")(ManagerEnv) +``` + +配置里的每个组是一个 `@configclass` 数据类,字段名即项名,字段值是项配置对象。 +下文按组介绍它们各自做什么、如何定义。 + +## 生命周期 + +`ManagerEnv` 同样继承 `ArrayEnv`,step 时序与 auto-reset / truncation 语义与 DirectEnv 一致。 +下图把各回调整合进流水线:绿色是 manager kernel(`@dispatch`)执行的部分, +紫色是 manager 的主机侧 Python 回调,蓝色是框架——编排各阶段的宿主与其自身的数据通路: + +```{image} /_static/images/tutorial/manager-env-lifecycle-light.svg +:alt: ManagerEnv 单个控制步流水线,包含 done 行分支:apply_action → physics_step → compute_transition → 截断判定 → done? 将 done 行送入 auto-reset 微流程、存活行送往 compute_observation +:class: only-light +``` + +```{image} /_static/images/tutorial/manager-env-lifecycle-dark.svg +:alt: ManagerEnv 单个控制步流水线,包含 done 行分支:apply_action → physics_step → compute_transition → 截断判定 → done? 将 done 行送入 auto-reset 微流程、存活行送往 compute_observation +:class: only-dark +``` + +与 DirectEnv 的两点区别:各项回调不在环境类里手写,而由配置组声明、manager 编译器 +编译成融合 kernel;kernel 的编译与预热发生在首次 `init_state()`(启动阶段), +step 循环内不再编译。 + +## 动作项(actions) + +**做什么**:每个控制步把策略输出的动作批次变换为执行器控制量并写入仿真器。 +各项的动作空间(一维 `Box`、float32)按声明顺序拼接成整个环境的 `action_space`, +策略动作也按同样切分分发给各项。 + +**如何定义**:一个动作项由两部分组成—— + +1. `ActionCfg` 子类:声明静态参数与执行器路由,`__call__(env, actuators)` 返回运行时 term。 + `actuator_names` 为空元组 `()` 时控制全部执行器,写显式名称时只控制列出的执行器; + 一个执行器只能归属一项。设为 `None` 表示本项不写执行器目标。 +2. `@kernel_data` 装饰的 `ActionTerm` 子类:持有持久状态(`np.ndarray` 字段,按环境行)与 + 共享模型数据(`SharedArray` 字段),实现三个方法:`action_space()` 给出本项的动作空间、 + `process(actions)` 把动作批次变换为路由执行器的控制量、`reset(env_ids)` 清理持久状态。 + +```python +@configclass(kw_only=True) +class MyJointActionCfg(ActionCfg): + actuator_names: tuple[str, ...] = () + scale: float = 0.5 + + def __call__(self, env: ManagerEnv, actuators) -> ActionTerm: + return MyJointAction(scale=np.float32(self.scale), num_envs=env.num_envs, + actuators=actuators) + + +@kernel_data +class MyJointAction(ActionTerm): + current: np.ndarray # 持久状态:上一动作缓存(num_envs 行) + scale: SharedArray # 静态参数:由 __call__ 传入 + + def action_space(self, env, actuators) -> gym.spaces.Box: + ... # 由执行器 ctrl_range 推导 + + def process(self, actions: np.ndarray) -> np.ndarray: + ... # 变换动作,返回本项路由的执行器控制量 + + def reset(self, env_ids: np.ndarray) -> None: + self.current[env_ids] = 0.0 +``` + +## 观测项(observations) + +**做什么**:把仿真状态与命令拼装成策略/价值网络的输入。只允许 `policy` 与 `value` +两个观测组,`policy` 必须存在:`policy` 面向 actor(可加噪),`value` 面向 critic +(通常无噪声、可含特权信息)。组内各项的输出按声明顺序拼接成观测向量。 + +**如何定义**:一个观测 term 是 `ObservationTermCfg` 子类,`__call__(env)` 返回 +`ObsTerm(size, dispatch, *args)`。`size` 是本 term 的输出宽度;dispatch 形如 +`def xxx_obs(ctx, out, *args) -> None`,把本项的观测写入 `out`——`ctx.sim["key"]` +读取 `queries` 声明的数据,噪声等作为标量参数传入、在 dispatch 内部添加。 + +```python +@dispatch +def projected_gravity_obs(ctx: ManagerContext, out: np.ndarray, + noise_amplitude: np.float32) -> None: + ... # 计算并写入 out(宽度 3) + add_uniform_noise(out, noise_amplitude, ctx.rand.state) + + +@configclass(kw_only=True) +class ProjectedGravityObsCfg(ObservationTermCfg): + noise: UniformNoiseCfg = UniformNoiseCfg() + + def __call__(self, env: ManagerEnv) -> ObsTerm: + return ObsTerm(3, projected_gravity_obs, np.float32(self.noise.amplitude)) +``` + +观测依赖的仿真数据通过环境 `queries` 组或项的 `required_sim_queries()` 声明, +见下文「查询(queries)」一节。 + +## 奖励项(rewards) + +**做什么**:每个控制步对当前状态给出一个标量评价,总奖励是各项按 `weight` 的加权和: +正 `weight` 是奖励,负 `weight` 是惩罚。 + +**如何定义**:一个奖励 term 是 `RewardTermCfg` 子类(自带 `weight: float` 字段), +`__call__(env)` 返回 `RewardTerm(dispatch, *args)`。dispatch 形如 +`def xxx_reward(ctx, *args) -> float`,从 `ctx.sim["key"]` 读取 `queries` 声明的数据。 + +```python +@dispatch +def base_height_reward(ctx: ManagerContext, target_z: np.float32, + sigma: np.float32) -> float: + error = ctx.sim["robot_base_pos"][2] - target_z + return math.exp(-(error * error) / (sigma * sigma)) + + +@configclass(kw_only=True) +class BaseHeightRewardCfg(RewardTermCfg): + target_z: float + sigma: float + + def __call__(self, env: ManagerEnv) -> RewardTerm: + return RewardTerm(base_height_reward, + np.float32(self.target_z), np.float32(self.sigma)) +``` + +## 终止项(terminations) + +**做什么**:每个控制步对每个环境输出一个布尔值,任一项为真即回合 `terminated` +(区别于超时 `truncated`)。 + +**如何定义**:一个终止 term 是 `TerminationTermCfg` 子类,`__call__(env)` 返回 +`TerminationTerm(dispatch, *args, metric_names=(...))`。dispatch 形如 +`def xxx_termination(ctx, *args) -> bool`;`metric_names` 可选,把 dispatch 写入 +`ctx.metrics` 的每环境量登记为可读指标。 + +```python +@dispatch +def bad_orientation_termination(ctx: ManagerContext, threshold: np.float32) -> bool: + ... + return tilt_sq > threshold * threshold + + +@configclass(kw_only=True) +class BadOrientationTerminationCfg(TerminationTermCfg): + threshold: float + + def __call__(self, env: ManagerEnv) -> TerminationTerm: + return TerminationTerm(bad_orientation_termination, + np.float32(self.threshold), + metric_names=("base_tilt",)) +``` + +## 重置项(sim_reset) + +**做什么**:回合重置(以及命令项请求的中途重算,见命令一节)时,改写选中环境行的 +仿真器状态——初始位姿、速度、关节角等,由 reset kernel 按声明顺序执行。 + +**如何定义**:一个重置 term 是 `ResetTermCfg` 子类,`__call__(env)` 返回 +`ResetTerm(dispatch, *args, writes={...})`。`writes` 声明本 term 要写哪些仿真状态 +(`BodyPositionWrite`、`JointVelocityWrite` 等,键名即 dispatch 里的写入通道); +dispatch 计算写入值。 + +```python +@dispatch +def _reset_body_pos(ctx: ManagerContext, sim_writes, spawn, noise) -> None: + position = sim_writes["position"] + ... # 写入带随机化的初始位置 + + +@configclass(kw_only=True) +class BodyPosResetCfg(ResetTermCfg): + spawn: tuple[float, float, float] + noise: tuple[float, float, float] = (0.02, 0.02, 0.005) + + def __call__(self, env: ManagerEnv) -> ResetTerm: + return ResetTerm( + _reset_body_pos, + tuple(np.asarray(self.spawn, dtype=np.float32)), + tuple(np.asarray(self.noise, dtype=np.float32)), + writes={"position": BodyPositionWrite((base_link,))}, + ) +``` + +## 命令项(commands) + +**做什么**:维护命令状态并在每个周期推进——速度指令采样、curriculum、参考动作推进等。 +命令项还产出每环境指标:dispatch 内写入 `ctx.metrics["name"]`,即可通过 +`env.state.metrics` / `env.metrics` 读取。大多数任务直接复用内置命令即可, +不需要自定义。 + +**如何定义**:一个命令 term 是 `CommandCfg` 子类,`__call__(env)` 返回 +`@kernel_data` 的 `CommandTerm` 子类,实现四个回调: + +| 回调 | 执行位置 | 职责 | +| ------------------------ | ----------------- | ------------------------------------------ | +| `update(ctx)` | evaluate kernel | 更新本周期派生命令数据 | +| `advance(ctx)` | transition kernel | 推进持久命令状态;置 `ctx.sim_reset_requested` 可请求该行在步骤末尾重算仿真状态 | +| `reset(ctx)` | 主机侧 | 为重置的环境行准备数据(含终止统计) | +| `reset_env(ctx)` | reset kernel | 重置持久命令状态 | + +## 查询(queries) + +**做什么**:声明各项需要的仿真数据与结构信息,由框架编译成批量读取程序。 +`queries.data` 是仿真数据查询(`JointPositionQuery`、`LinkPositionQuery` 等), +`queries.model` 是模型查询(`ActuatorKpQuery`、`BodyJointPositionLimitsQuery` 等)。 + +**如何声明**:通常在配置的 `__post_init__` 中根据机器人配置填充: + +```python +def __post_init__(self) -> None: + joint_names = ... # 由 scene.objs.robot 解析 + base_link = ... # 基座 link 名 + self.queries.data["robot_dof_pos"] = JointPositionQuery(joints=joint_names) + self.queries.model["robot_joint_position_limits"] = BodyJointPositionLimitsQuery(body=base_link) +``` + +观测项也可以通过 `required_sim_queries(env_cfg)` 声明自己需要的查询; +所有声明按 key 合并,同名 key 必须完全一致,否则配置失败。查询的编译与执行机制见 +[SimBackend:与仿真器解耦](sim_backend.md)。 + +## 注册与变体 + +- **配置先于环境类注册**;registry 根据类继承关系推断前端类型 + (`ManagerEnv` 与 `DirectEnv` 同属 `"np"` 数据后端),不要手动指定。 +- 配置 docstring 首行与 `zh_CN:` 行分别作为环境的英文、中文描述。 +- 同一环境的多个 preset 通过配置 factory 派生:子变体只覆写与共享配置的差异项 + (例如调低某个奖励权重)。 + +以内置 `microduck-ball-balance` 为例(源码见 +`motrix_envs/src/motrix_envs/locomotion/ball_balance/microduck.py`): + +```python +@registry.envcfg("microduck-ball-balance") +def make_microduck_ball_balance_cfg() -> MicroduckBallBalanceEnvCfg: + """Balance on top of a basketball with Microduck. + + zh_CN: 让 Microduck 双脚站在篮球上并保持平衡。 + """ + return MicroduckBallBalanceEnvCfg() + + +registry.env("microduck-ball-balance")(ManagerEnv) +``` + +注册完成后,`python scripts/view.py env=` 可以预览,按 +[Task 配置](../training/task_config.md)创建训练 Task 即可开始训练。 diff --git a/docs/source/zh_CN/user_guide/tutorial/physics_environment.md b/docs/source/zh_CN/user_guide/tutorial/building_envs/scene.md similarity index 96% rename from docs/source/zh_CN/user_guide/tutorial/physics_environment.md rename to docs/source/zh_CN/user_guide/tutorial/building_envs/scene.md index 9a17c10a..4db41c4b 100644 --- a/docs/source/zh_CN/user_guide/tutorial/physics_environment.md +++ b/docs/source/zh_CN/user_guide/tutorial/building_envs/scene.md @@ -1,6 +1,6 @@ -# 物理环境配置 +# SceneCfg:搭建物理场景 -物理环境配置定义了强化学习训练中的仿真参数和场景设置。 +物理场景搭建定义了强化学习训练中的仿真参数和场景设置。 MotrixLab 使用了[MotrixSim](https://motrixsim.readthedocs.io/zh-cn/latest/user_guide/index.html)作为物理仿真后端。 ## 支持的文件格式 diff --git a/docs/source/zh_CN/user_guide/tutorial/building_envs/sim_backend.md b/docs/source/zh_CN/user_guide/tutorial/building_envs/sim_backend.md new file mode 100644 index 00000000..e1eb0a53 --- /dev/null +++ b/docs/source/zh_CN/user_guide/tutorial/building_envs/sim_backend.md @@ -0,0 +1,29 @@ +# SimBackend:与仿真器解耦 + +`DirectEnv` 与 `ManagerEnv` 的所有仿真交互都通过 backend 中立的 `SimBackend` 接口完成 +(即环境持有的 `self.sim`)。具体仿真器类型不会跨越 core 边界——同一份环境代码可以运行在 +任何实现了该接口的后端上。 + +## 编译三类程序 + +环境在构造阶段对 `self.sim` 编译三类程序,之后每步只做轻量的 buffer 写入和 `execute()`: + +- `compile_model(queries)`:模型查询,即不随仿真变化的结构信息(执行器列表、关节限位等)。 +- `compile_reads(queries)`:仿真数据读取程序(关节位置、速度、link 位姿等), + 每次 `execute(env_ids)` 批量刷新,读取结果按声明时的 key 索引。 +- `compile_writes({name: SimWrite})`:写入程序(由 `write_compiler` 编译器承载)。控制目标用 `CtrlTargetsWrite`; + 重置写入(初始位姿等)传入 `reset=True`。写入时先填 `buffer(name)`,再 `execute(env_ids)`。 + +两类工作流的分工:`DirectEnv` 在自己的构造函数中编译这些程序 +(见[编写 DirectEnv 环境](direct_env.md)的最小示例);`ManagerEnv` 则由配置的 +`queries` 组与各项自动生成,用户通常无需直接操作 +(见[编写 ManagerEnv 环境](manager_env.md)的查询一节)。 + +## backend 的选择与注册 + +- backend 是构造级别的字符串选择:创建环境时 `backend=None` 使用注册的默认后端 + (当前为 `motrixsim`),也可显式传入后端名,或通过 + `registry.make(..., sim="motrixsim")` 指定。 +- backend 通过 `motrix_env.sim_backends` entry-point 组懒发现: + 第三方仿真器实现一个 `SimBackend` 工厂并注册到该组即可接入,导入开销为零。 +- `motrix_env_mujoco` 仅用于 MuJoCo scene 编译,不作为训练后端。 diff --git a/docs/source/zh_CN/user_guide/tutorial/index.md b/docs/source/zh_CN/user_guide/tutorial/index.md new file mode 100644 index 00000000..fb3d0833 --- /dev/null +++ b/docs/source/zh_CN/user_guide/tutorial/index.md @@ -0,0 +1,13 @@ +# 教程总览 + +本节按照"先建立整体图景,再逐步深入细节"的顺序组织,左侧目录即完整的阅读结构: + +- **基础框架**:MotrixLab 的分层架构与运行流程,理解环境、RL provider、Hydra 配置与 Trainer 如何协作。 +- **构建环境**:先写环境(DirectEnv、ManagerEnv 两篇指南),再引入两个横切概念 + (SceneCfg 场景与仿真配置、SimBackend 与仿真器解耦);入口是环境构建总览 + (环境组成、生命周期、工作流选型)。 +- **训练与结果**:从创建 Task 配置、执行训练到分析 `runs/` 产物与 checkpoint。 +- **进阶主题**:按需深入的独立能力——ONNX 导出、真机部署、指令输入架构与自定义训练后端。 + +如果还没有跑通第一个训练,建议先阅读[入门指南](../getting_started/installation.md); +如果只想了解环境实现,直接从[环境构建总览](building_envs/index.md)开始。 diff --git a/docs/source/zh_CN/user_guide/tutorial/rewards.md b/docs/source/zh_CN/user_guide/tutorial/rewards.md deleted file mode 100644 index fb7cb833..00000000 --- a/docs/source/zh_CN/user_guide/tutorial/rewards.md +++ /dev/null @@ -1,51 +0,0 @@ -# 奖励函数设计 - -奖励函数告诉智能体什么样的行为是期望的,是强化学习环境设计中的核心部分。 - -## 奖励函数在训练循环中的位置 - -在 MotrixLab 的 `DirectEnv` 工作流中,奖励计算发生在 `step` 函数的 `compute_transition` 阶段: - -```python -# DirectEnv.step() 的执行流程 -def step(self, actions: np.ndarray) -> ArrayEnvState: - # 1. 准备阶段:清空奖励和状态 - self._prev_physics_step() # reward = 0.0, terminated = False, truncated = False - - # 2. 应用动作 - self._state = self.apply_action(actions, self._state) - - # 3. 物理仿真 - self.physics_step() # 执行物理仿真 - - # 4. 计算 transition ← 奖励函数在这里计算 - self._state = self.compute_transition(self._state) - - # 5. 后续处理 - self._update_truncate() # 检查时间截断 - self._reset_done_envs() # 重置完成的环境 - self._state = self.compute_observation(self._state) # 计算重置后的观察值 - - return self._state -``` - -您需要在子类的 `compute_transition` 方法中实现奖励计算逻辑,具体奖励函数设计思路请参考训练示例。 - -### 奖励组件设计原则 - -1. **分离关注点**:每个奖励函数负责一个特定的目标 -2. **权重配置**:通过配置文件管理不同组件的权重 -3. **归一化**:保持奖励值在合理的范围内 -4. **平滑性**:避免硬性阈值,使用指数函数等平滑过渡 - -这种方法使得奖励函数模块化,便于调试和调整各个组件的权重。 - -## 设计原则 - -1. **明确的目标导向**:奖励函数应该直接反映任务目标 -2. **合理的奖励范围**:避免过大或过小的奖励值,保持训练稳定 -3. **平衡探索与利用**:适当奖励接近目标的行为,避免稀疏奖励 -4. **避免奖励漏洞**:检查智能体是否可能通过不期望的方式获得高奖励 -5. **调试友好**:在开发阶段输出奖励分解信息,便于调优 - -通过在 `compute_transition` 方法中正确实现奖励计算,您可以为各种机器人任务设计有效的学习信号。 diff --git a/docs/source/zh_CN/user_guide/tutorial/training/index.md b/docs/source/zh_CN/user_guide/tutorial/training/index.md new file mode 100644 index 00000000..7629fcf6 --- /dev/null +++ b/docs/source/zh_CN/user_guide/tutorial/training/index.md @@ -0,0 +1,28 @@ +# 训练总览 + +环境注册完成后,训练一个强化学习任务只需要三步: + +```text +创建 Task 配置 指定环境、RL 框架、算法与超参数 + ↓ +python scripts/train.py task=... Hydra 组合配置并启动训练 + ↓ +runs/{env}/... 保存 metadata、日志与 checkpoint +``` + +各页面按使用顺序展开: + +- [Task 配置与命令行参数覆盖](task_config.md):创建 Task 文件、调整运行与算法参数、 + 使用 `key=value` 临时覆盖; +- [执行训练与分析结果](training_and_result.md):启动训练、读取 TensorBoard 日志、 + 用 play 回放策略; +- [训练产物:runs 目录与 checkpoint](runs_and_checkpoints.md):`runs/` 的目录结构、 + best policy 的选择逻辑与续训方式。 + +```{toctree} +:hidden: + +task_config +training_and_result +runs_and_checkpoints +``` diff --git a/docs/source/zh_CN/user_guide/tutorial/runs_and_checkpoints.md b/docs/source/zh_CN/user_guide/tutorial/training/runs_and_checkpoints.md similarity index 100% rename from docs/source/zh_CN/user_guide/tutorial/runs_and_checkpoints.md rename to docs/source/zh_CN/user_guide/tutorial/training/runs_and_checkpoints.md diff --git a/docs/source/zh_CN/user_guide/tutorial/training_environment_config.md b/docs/source/zh_CN/user_guide/tutorial/training/task_config.md similarity index 98% rename from docs/source/zh_CN/user_guide/tutorial/training_environment_config.md rename to docs/source/zh_CN/user_guide/tutorial/training/task_config.md index 0fc4940a..f466d9df 100644 --- a/docs/source/zh_CN/user_guide/tutorial/training_environment_config.md +++ b/docs/source/zh_CN/user_guide/tutorial/training/task_config.md @@ -146,7 +146,7 @@ algo: timesteps: 20000 ``` -完整字段和逐字段注释见 [`configs/algo_base/skrl.ppo.yaml`](../../../../configs/algo_base/skrl.ppo.yaml)。 +完整字段和逐字段注释见 [`configs/algo_base/skrl.ppo.yaml`](../../../../../configs/algo_base/skrl.ppo.yaml)。 ### RSLRL PPO @@ -167,7 +167,7 @@ algo: entropy_coef: 0.005 ``` -完整字段和逐字段注释见 [`configs/algo_base/rslrl.ppo.yaml`](../../../../configs/algo_base/rslrl.ppo.yaml)。 +完整字段和逐字段注释见 [`configs/algo_base/rslrl.ppo.yaml`](../../../../../configs/algo_base/rslrl.ppo.yaml)。 ### Motrix FastSAC @@ -182,7 +182,7 @@ FastSAC 使用同一个算法身份支持两种执行拓扑,主要字段和配 - `weight_publish_interval`、`weight_poll_interval`:策略权重同步频率。 - `max_ingest_per_iter`、`idle_sleep_s`:Learner 数据摄取和空闲退避参数。 -完整字段和逐字段注释见 [`configs/algo_base/motrix.fastsac.yaml`](../../../../configs/algo_base/motrix.fastsac.yaml)。 +完整字段和逐字段注释见 [`configs/algo_base/motrix.fastsac.yaml`](../../../../../configs/algo_base/motrix.fastsac.yaml)。 ## 使用 CLI 临时覆盖参数 diff --git a/docs/source/zh_CN/user_guide/tutorial/training_and_result.md b/docs/source/zh_CN/user_guide/tutorial/training/training_and_result.md similarity index 98% rename from docs/source/zh_CN/user_guide/tutorial/training_and_result.md rename to docs/source/zh_CN/user_guide/tutorial/training/training_and_result.md index 8929b1e6..dd611ac7 100644 --- a/docs/source/zh_CN/user_guide/tutorial/training_and_result.md +++ b/docs/source/zh_CN/user_guide/tutorial/training/training_and_result.md @@ -26,7 +26,7 @@ python scripts/train.py task=g1-walk-flat/motrix.fastsac algo.asynchronous=false | `rslrl.ppo` | `torch` | RSLRL PPO | | `motrix.fastsac` | `torch` | FastSAC;`algo.asynchronous` 选择同步或异步拓扑 | -运行 `python scripts/train.py --help` 可以查看当前代码中全部可选 Task。Task 文件结构和覆盖规则见 [Task 配置与命令行参数覆盖](training_environment_config.md)。 +运行 `python scripts/train.py --help` 可以查看当前代码中全部可选 Task。Task 文件结构和覆盖规则见 [Task 配置与命令行参数覆盖](task_config.md)。 ### 选择训练后端与仿真后端 diff --git a/motrix_env_core/src/motrix_env_core/numba/manager/env.py b/motrix_env_core/src/motrix_env_core/numba/manager/env.py index b33b7033..c622c351 100644 --- a/motrix_env_core/src/motrix_env_core/numba/manager/env.py +++ b/motrix_env_core/src/motrix_env_core/numba/manager/env.py @@ -425,7 +425,7 @@ def __init__(self, cfg: EnvCfgType, num_envs: int = 1, backend: str | None = Non if not self._action_cfgs: raise ValueError("Manager environment config requires at least one action config.") self._action_actuators = self._resolve_action_actuators() - self._action_writes = self.sim.write_compiler.compile( + self._action_writes = self.sim.compile_writes( { name: CtrlTargetsWrite(None if action_cfg.actuator_names == () else action_cfg.actuator_names) for name, action_cfg in self._action_cfgs.items() diff --git a/motrix_env_core/src/motrix_env_core/numba/manager/sim_reset.py b/motrix_env_core/src/motrix_env_core/numba/manager/sim_reset.py index 515ca5a3..7499180e 100644 --- a/motrix_env_core/src/motrix_env_core/numba/manager/sim_reset.py +++ b/motrix_env_core/src/motrix_env_core/numba/manager/sim_reset.py @@ -113,7 +113,7 @@ def create( for term_name, term_writes in writes.items() for output_name, write in term_writes.items() } - program = sim.write_compiler.compile(flat_writes, reset=True) + program = sim.compile_writes(flat_writes, reset=True) buffers = tuple( program.buffer(f"{term_name}.{output_name}") for term_name, term_writes in writes.items() diff --git a/motrix_env_core/src/motrix_env_core/sim/backend.py b/motrix_env_core/src/motrix_env_core/sim/backend.py index 2146e014..f950a07e 100644 --- a/motrix_env_core/src/motrix_env_core/sim/backend.py +++ b/motrix_env_core/src/motrix_env_core/sim/backend.py @@ -24,7 +24,7 @@ from motrix_env_core.config.scene import SceneCfg, SystemCameraCfg from motrix_env_core.sim.model import ModelQuery, SimModelQueryCompiler from motrix_env_core.sim.read import PhysicsReadProgram -from motrix_env_core.sim.write import SimWriteCompiler +from motrix_env_core.sim.write import SimWrite, SimWriteCompiler, WriteProgram class ActuatorType(str, Enum): @@ -230,6 +230,22 @@ def compile_model(self, queries: Mapping[str, ModelQuery]) -> SimModel: """Lower static model-query declarations into one backend-neutral model.""" return self.model_query_compiler.compile(queries) + def compile_writes( + self, + writes: Mapping[str, SimWrite], + *, + reset: bool = False, + forward_kinematics: bool = True, + ) -> WriteProgram: + """Lower named write declarations into one executable write program. + + Convenience twin of :meth:`compile_reads` and :meth:`compile_model` + for the write path: it forwards to :attr:`write_compiler`, so + backends without live write support report the gap when the property + is accessed. + """ + return self.write_compiler.compile(writes, reset=reset, forward_kinematics=forward_kinematics) + @property def write_compiler(self) -> SimWriteCompiler: """Return the compiler bound to this backend's model and live data. diff --git a/motrix_env_core/tests/test_direct_env_sim_backend.py b/motrix_env_core/tests/test_direct_env_sim_backend.py index 1fe21363..b4950d8a 100644 --- a/motrix_env_core/tests/test_direct_env_sim_backend.py +++ b/motrix_env_core/tests/test_direct_env_sim_backend.py @@ -191,8 +191,8 @@ def __init__(self, cfg: _FakeDirectCfg, num_envs: int, backend: str | None = Non "actuator_ctrls": ActuatorCtrlQuery(), } ) - self._ctrl_writes = self.sim.write_compiler.compile({"ctrl": CtrlTargetsWrite()}) - self._reset_program = self.sim.write_compiler.compile( + self._ctrl_writes = self.sim.compile_writes({"ctrl": CtrlTargetsWrite()}) + self._reset_program = self.sim.compile_writes( {"state_position": DofPositionWrite(), "state_velocity": DofVelocityWrite()}, reset=True ) self._action_space = gym.spaces.Box(-1.0, 1.0, (2,), dtype=np.float32)