Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
41 commits
Select commit Hold shift + click to select a range
99dc4a3
feat(manager)!: replace the Numba manager path with a generic Torch e…
wlgys8 Nov 5, 2026
f45a4d2
chore(bench): drop the concluded batch-mdp microbenchmark
wlgys8 Nov 5, 2026
97327f0
fix(ci): drop the removed tbb/numba extras and markers from CI
wlgys8 Nov 5, 2026
a41325f
refactor(core)!: remove the unused shared-primitive namespace
wlgys8 Nov 5, 2026
68f97b3
refactor(manager)!: drop the @dispatch requirement from manager term …
wlgys8 Nov 5, 2026
8942eb1
refactor(mdp): consume the joint-default argument as the tensor view …
wlgys8 Nov 5, 2026
93e5a02
refactor(terrain): tidy heightfield_lookup placement, typing, and vie…
wlgys8 Nov 5, 2026
7e97151
refactor(manager): remove the unused MdpContext/ContextMap protocols
wlgys8 Nov 5, 2026
dbb1429
refactor(wbt): drop the batch_math module alias for quaternion_array
wlgys8 Nov 5, 2026
1ebb9a1
refactor(manager): type the task program as TorchTaskProgram
wlgys8 Nov 5, 2026
3b7e586
refactor(manager): type SimResetRuntime.apply with TorchTaskProgram
wlgys8 Nov 5, 2026
249c6d6
refactor(manager): collapse term-arg leaf checks into one whitelist
wlgys8 Nov 5, 2026
71cf72c
refactor(math): inline the array-namespace shim into quaternion_array
wlgys8 Nov 5, 2026
fae9672
refactor(manager): inline the manager binding shim into env
wlgys8 Nov 5, 2026
572f739
refactor(ball): build reset noise tensors directly from tuples
wlgys8 Nov 5, 2026
f18d372
refactor: drop the last numpy stepping stones in reset/reward term args
wlgys8 Nov 5, 2026
c05ea6f
refactor(torch): type TorchManagerRuntime with ManagerEnv and Resolve…
wlgys8 Nov 5, 2026
411f425
refactor(torch): make the compiled kernel the single implementation p…
wlgys8 Nov 5, 2026
d100cea
refactor(torch): annotate the runtime's kernel and helper signatures
wlgys8 Nov 5, 2026
90b3f60
refactor(torch): drop the unused outputs parameter from kernel builders
wlgys8 Nov 5, 2026
c7757ec
refactor(manager)!: resolve the manager data path as trees, not a fla…
wlgys8 Nov 5, 2026
973c42d
refactor(manager): replace PreparedInvocation with per-kind typed ter…
wlgys8 Nov 5, 2026
3dfd50b
refactor(torch): derive the reset clone surface from the write contract
wlgys8 Nov 5, 2026
7db922c
refactor(manager)!: delete SharedArray and the kernel-data scope mach…
wlgys8 Nov 6, 2026
ed79068
refactor(manager)!: drop @kernel_data and own carrier fields as torch…
wlgys8 Nov 6, 2026
8f7425f
feat(manager)!: exchange env state with the RL layer as torch tensors
wlgys8 Nov 6, 2026
82b9d72
refactor(manager): resolve sim leaves as zero-copy tensor views
wlgys8 Nov 6, 2026
1b48ffc
refactor(manager)!: pass carrier term arguments directly and drop Tre…
wlgys8 Nov 6, 2026
1a9f7a2
refactor(manager)!: pass array term arguments through and drop LeafArg
wlgys8 Nov 6, 2026
a12e76c
refactor(core)!: parameterize the env state by array type
wlgys8 Nov 6, 2026
37753f8
refactor(kernel_data)!: collapse the dual tree schemas into one walk
wlgys8 Nov 6, 2026
8327bee
feat(manager)!: exchange step actions as torch tensors
wlgys8 Nov 6, 2026
8dca412
refactor(core)!: implement the state allocation hooks via the array n…
wlgys8 Nov 6, 2026
033fecd
refactor(kernel_data)!: replace the flatten ABI with tree visitors
wlgys8 Nov 6, 2026
e20b315
refactor(kernel_data)!: consume cfg-provided carriers by reference
wlgys8 Nov 6, 2026
51e5f85
refactor(manager)!: remove the unused warmup execution path
wlgys8 Nov 6, 2026
a73d0c3
feat(manager)!: take reset env_ids as torch tensors
wlgys8 Nov 6, 2026
3b0ee91
refactor(manager)!: delete the program fingerprint chain
wlgys8 Nov 6, 2026
82ed40c
feat(torch_manager): optional barrier-removal pass for Inductor CPU k…
wlgys8 Nov 6, 2026
a2dfea8
refactor(kernel_data)!: replace the schema tree with a direct leaf walk
wlgys8 Nov 6, 2026
3ef701c
fix(kernel_data): drop the cache on type-hint resolution
wlgys8 Nov 6, 2026
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
51 changes: 8 additions & 43 deletions .github/workflows/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -25,10 +25,9 @@ Python 3.10 uses `--all`; the other versions omit the Unitree hardware plugin an

PR and push runs have two independent, parallel suites per Python version:

- **fast**: `pytest -m "not slow and not integration and not numba"`.
- **numba**: `pytest -m "not integration and (slow or numba)"`. This includes the two expensive
- **fast**: `pytest -m "not slow and not integration"`.
- **slow**: `pytest -m "not integration and (slow)"`. This includes the expensive
WBT kernel regressions on every PR, so failures are not deferred until the next day.
The entire `test_wbt_numba.py` module stays together: excluding only the first two expensive
tests shifts first-use compilation into later tests rather than making the fast suite cheap.

These selections partition all non-integration tests, including future `slow` tests. They run on
Expand All @@ -40,64 +39,30 @@ with these matrix checks.
Nightly and manual dispatch run the **full** suite, including
`test/test_all_envs.py::test_all_demos`. The full-environment subprocess smoke test is excluded
from PR/push jobs. Each job reports the 20 slowest phases and uploads JUnit results, including on
failure. Numba caches have per-suite write keys to avoid parallel jobs competing for one cache.
failure.

Local commands (activate `.venv` first):

```bash
# Quick feedback; excludes expensive kernel regressions and full-environment smoke.
python -m pytest -m "not slow and not integration and not numba"
python -m pytest -m "not slow and not integration"
# Dedicated compilation/runtime regressions.
python -m pytest -m "not integration and (slow or numba)"
python -m pytest -m "not integration and (slow)"
# Only the full-environment integration tests.
python -m pytest -m integration
# Full coverage remains the default; no tests are silently excluded by addopts.
python -m pytest
```

Markers are registered in the root `pyproject.toml`: `slow` describes cost, `numba` identifies
specialized kernel regressions, and `integration` identifies tests reserved for full runs.
Not every test importing Numba needs the `numba` marker: inexpensive core contract tests remain
in the fast suite.
Markers are registered in the root `pyproject.toml`: `slow` describes cost and `integration`
identifies tests reserved for full runs.

For new Numba environments, reuse small synthetic-manager tests for shared runtime contracts,
For new manager environments, reuse small synthetic-manager tests for shared runtime contracts,
keep task-specific smoke tests minimal, and retain a small set of representative real-environment
regressions. The existing nightly smoke test automatically covers all registered environments.
Do not duplicate a full expensive rollout for every preset. If the dedicated job grows too large,
measure `--durations` and shard by environment family; do not silently drop kernel coverage from PRs.

#### Diagnosing Numba startup costs

Compare a fresh `NUMBA_CACHE_DIR` with a second process using the same directory, then compare
another environment in the same process. `NUMBA_DEBUG_CACHE=1` reports actual specialization
cache loads/saves. Manager INFO logs time each evaluate/observe/reset specialization; DEBUG
logs include signatures. Dispatcher preparation is not the actual compilation phase.

Precompilation uses the real input types, and warmup copies preserve strides and readonly flags,
so warmup and reset/step share signatures. Function fingerprints use explicit code fields rather
than marshal serialization, which can change after Numba inspects an unchanged code object.
Keep these invariants intact when extending the compiler: a warmup must not introduce new
signatures, and constructing the same configuration after execution must reuse the same plan.

A local G1 WBT play-mode benchmark (`num_envs=2`, init and one step; interpreter imports excluded)
measured about 29 seconds with an empty cache, 0.27 seconds in a second process with disk cache,
and 0.07 seconds for another environment using the in-process cache after these fixes.
The original cold benchmark took about 57 seconds and compiled each kernel twice.
These are diagnostic measurements, not CI timing thresholds; hardware and cache state vary.

### File: `codeql.yml`

Runs GitHub CodeQL analysis for Python on pushes, pull requests, and a weekly
scheduled scan. Review results under the repository's **Security** tab.

### File: `branch-policy.yml`

Checks pull request metadata without checking out contributor code. Normal PRs
must target `main`; a post-patch back-merge may use upstream `stable → main`.
Only a release PR from the upstream `main` branch or a `hotfix/*` stable patch
PR may target `stable`. Add `Branch policy / validate` as a required status check on both
protected branches.

## Docker Image Build Workflow

### File: `docker-build.yml`
Expand Down
25 changes: 4 additions & 21 deletions .github/workflows/tests.yml
Original file line number Diff line number Diff line change
Expand Up @@ -25,25 +25,8 @@ jobs:
fail-fast: false
matrix:
python: ["3.10", "3.11", "3.12", "3.13"]
suite: ${{ fromJSON((github.event_name == 'schedule' || github.event_name == 'workflow_dispatch') && '["full"]' || '["fast","numba"]') }}
suite: ${{ fromJSON((github.event_name == 'schedule' || github.event_name == 'workflow_dispatch') && '["full"]' || '["fast","slow"]') }}
steps:
- name: Resolve numba cache dir
run: |
# Redirect all Numba disk caches (manager fused kernels and term
# specializations) to a single path so they can be cached between
# runs. Kernel plan keys are content hashes, so stale cache entries
# are simply unused and never produce wrong results.
# Numba does not expand "~" in NUMBA_CACHE_DIR, so resolve the path
# via $HOME here; the cache action below reads the same variable.
echo "NUMBA_CACHE_DIR=$HOME/.cache/motrixlab/numba-manager" >> "$GITHUB_ENV"
- name: Restore numba kernel cache
uses: actions/cache@55cc8345863c7cc4c66a329aec7e433d2d1c52a9 # v6.1.0
with:
path: ~/.cache/motrixlab
key: numba-kernel-${{ runner.os }}-py${{ matrix.python }}-${{ matrix.suite }}-${{ github.sha }}
restore-keys: |
numba-kernel-${{ runner.os }}-py${{ matrix.python }}-${{ matrix.suite }}-
numba-kernel-${{ runner.os }}-py${{ matrix.python }}-
- name: Checkout repository
uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
with:
Expand All @@ -64,16 +47,16 @@ jobs:
sh install.sh --all --python 3.10
else
# Unitree/DDS runtime support remains limited to Python 3.10.
uv sync --frozen --python '${{ matrix.python }}' --extra cuda --extra rslrl --extra skrl-jax --extra tbb --group test
uv sync --frozen --python '${{ matrix.python }}' --extra cuda --extra rslrl --extra skrl-jax --group test
fi
.venv/bin/python -c 'import sys; assert ".".join(map(str, sys.version_info[:2])) == "${{ matrix.python }}"'

- name: Run tests
run: |
selection=''
case '${{ matrix.suite }}' in
fast) selection='not slow and not integration and not numba' ;;
numba) selection='not integration and (slow or numba)' ;;
fast) selection='not slow and not integration' ;;
slow) selection='not integration and (slow)' ;;
esac
set --
if [ '${{ matrix.python }}' != '3.10' ]; then
Expand Down
19 changes: 11 additions & 8 deletions AGENTS.md
Original file line number Diff line number Diff line change
Expand Up @@ -76,11 +76,14 @@ README 从宏观上说明 package:它提供什么能力、如何安装与使
框架级契约(Runtime/ControlSession/Controller 语义、rollout result 口径等)只在一处维护——对应框架 package 的
README 或教程;插件与任务 package 的 README 不重复这些内容,只描述本 package 特有的部分。

### Manager Kernel 入口方法
### Manager term 入口方法

Manager fused-kernel 入口方法必须用 `@dispatch` 装饰,绝不使用 `@njit`——Numba 编译、缓存和内联由 manager
编译器统一管理。`@njit` 只用于被编译 kernel 调用的独立 Numba 兼容 helper。哪些方法是 kernel 入口、哪些是普通
Python 方法,由 manager 运行时契约定义(`wiki/design/manager/runtime.md`)。
Manager 数值 term 是普通 Python 函数,契约由签名注解表达并由 resolver 校验(首参数注解
`ManagerContext`,observation 的 `out` 注解 `torch.Tensor`,reset 的 `sim_writes` 注解
`Map[torch.Tensor]`);数值执行由 Torch 执行器按整批张量调度。`@dispatch` 仅用于标记 kernel_data
类上需要进入布局方法代理表的方法(如 `ActionState.current/previous`、CommandTerm 的
`update`/`advance`/`reset_env` hooks,后者为 `@staticmethod`,签名为 `(receiver, ctx)`)。运行时
契约见 `wiki/design/manager/runtime.md`。

### 环境配置与注册规范

Expand Down Expand Up @@ -238,10 +241,10 @@ tensorboard --logdir runs/{env-name}
### 测试

```bash
# 与 PR CI 一致的快速测试(不执行集成、慢速或 Numba 编译测试)
python -m pytest -m "not slow and not integration and not numba"
# CPU 编译回归、跨模块/进程集成分别执行
python -m pytest -m "not integration and (slow or numba)"
# 与 PR CI 一致的快速测试(不执行集成、慢速测试)
python -m pytest -m "not slow and not integration"
# 慢速回归、跨模块/进程集成分别执行
python -m pytest -m "not integration and slow"
python -m pytest -m integration
```

Expand Down
7 changes: 0 additions & 7 deletions bench/bench_physics.py
Original file line number Diff line number Diff line change
Expand Up @@ -81,13 +81,6 @@ def _device_context() -> dict:
"num_numa_nodes": _detect_num_nodes(),
"binding": _current_binding(),
}
try:
import numba as _numba

ctx["numba_threads"] = _numba.get_num_threads()
ctx["numba_layer"] = _numba.threading_layer()
except Exception:
pass
return ctx


Expand Down
Loading
Loading