feat(manager)!: replace the Numba manager path with a generic Torch executor by wlgys8 · Pull Request #157 · Motphys/MotrixLab

wlgys8 · 2026-11-05T12:53:12Z

摘要

移除 ManagerEnv 的 Numba fused-kernel 执行路径，替换为通用批量 Torch 执行器；ManagerEnv 全链路（数据所有权、执行、env↔算法层交换）统一为 torch，term 配置 API 形式保持不变。

主要变更

TorchManagerRuntime（motrix_env_core/torch_manager/runtime.py）：复用 resolver 产出的执行元数据，按原 fused-kernel 语句顺序批量调度全部 term 与命令钩子；per-lane SplitMix64 RNG 逐位保持。默认 torch.compile，compile=False 可回退 eager 调试。
全部 manager MDP term 批量化并全类型标注：core mdp、humanoid walk、WBT motion command（含 sequential/non-sequential/adaptive sampling）、ball balance。
Numba 完全移除：编译器/codegen/缓存、per-lane dispatch、TBB 预加载、numba/tbb 依赖（pyproject + uv.lock）；src/tests 零 numba import。
数据层 torch 化：kernel_data 载体字段（零拷贝视图 over sim arena）、sim 叶子、输出缓冲（obs/reward/terminated/双缓冲/final-obs）、step 输入（动作）全 torch；ArrayEnvState 泛型化（EnvObs[T]，Manager 钉 torch、Direct 钉 numpy）；env.reset 收 torch 张量。公开 API 单一货币：ManagerEnv 进出皆张量。
结构收敛（每轮均带逐位数值 + graph-break + bench 三重验收）：resolver 树形数据路径（flat ABI 删除）；PreparedInvocation 分型化；TreeArg/LeafArg/SharedArray/scope/program-fingerprint/warmup-clone 等中间层逐一删除或简化；kernel_data 从 Numba ABI 机器收缩为五个小文件。
多核优化：Inductor cpp.min_chunk_size 4096→256（根因修复：默认值将 manager 规模循环串行化为 omp-single）；observe 按 group 拆分编译图（I-cache 热度，A/B 复测确认在高线程下优势扩大）；掩码静态 reset 图；MOTRIX_TORCH_MERGED_OMP=1 实验：对 Inductor 生成码做循环间依赖分析，无数据冒险处加 nowait 消 barrier（EPYC +6–8%，128 线程不设上限时改善 26%）。
文档：wiki 设计文档全面更新（torch-executor.md 含性能特性节）。

性能

机器	本分支（最优配置）	Numba（main 基线，同机同参）
32 核桌面	15.7–16.0 ms/step（~130k steps/s）	16.98 ms/step
192 核 EPYC	16 线程 12.0 ms（170k steps/s）	16 线程 12.99 / ~180 线程 12.46 ms/step

已知残余：EPYC 不设图内线程上限时多线程扩展性差于 Numba（192 线程 ~30 vs ~12.5ms）——归因于 Inductor CPU codegen 的数据并行模型（每缓冲一个循环 + 隐式 barrier + 每图 fork/join）对比 Numba 的 lane 并行（单 prange、零中间物化）。默认图内 8 线程上限下 192 核不崩（15.3ms）且 16 线程已优于 Numba。消 barrier 实验开关（MOTRIX_TORCH_MERGED_OMP）部分缓解。

Torch 路线的技术权衡（已评估）：

✅ 优点：单一数据类型（torch tensor 端到端）、单一依赖链（无 numba/tbb）、term 源码后端无关（batch 写法让并行策略成为编译器自由度，未来 torch 版本改进 CPU codegen 可零改动受益）、调试路径统一（eager fallback 语义相同）。
⚠️ 权衡：高核数不设上限时多线程扩展不如手写融合 kernel（Inductor 保守 barrier 策略的结构性代价）；torch.compile 首次编译秒级成本（Inductor 磁盘缓存跨进程复用缓解）；依赖 PyTorch 版本行为（2.7 验证，升级需重新校准——batch term 源码本身不依赖版本特定行为）。
📋 后续方向：给 PyTorch 上游提 CPU codegen 改进需求（barrier 消除 / lane 调度，merged_OMP 的依赖分析可作证据）；torch 新版本 CPU codegen 评估（独立探针项目进行中）；自动线程配置（torch.set_num_threads 按拓扑选择）。

验收

数值 parity（迁移期 vs Numba 同 seed）：core bit-exact；walk 2.4e-7；WBT 2.4e-7（含全部 advance/reset 分支）；ball balance 3e-8；RNG 逐位一致
训练效果对比（g1-wbt-dance/FastSAC/2048 envs/seed=1/50k iters，多轮）：最终平滑 return 46.7–47.3 vs main 46.4，统计等价（该对比曾抓出动作 echo 观测的 gather bug，含回归测试）
测试：python -m pytest motrix_env_core/tests motrix_envs/tests tests -m "not slow and not integration" → 737+ passed
冒烟：全部 manager 注册环境（init + steps + torch 部分复位）通过

…xecutor Remove the Numba fused-kernel execution path from ManagerEnv entirely and make a generic batched Torch executor the only numerical backend, keeping the ManagerEnv host lifecycle and term-configuration API unchanged. - Add TorchManagerRuntime (motrix_env_core/torch_manager/runtime.py): it consumes the resolver metadata (read plan, layout, invocations) and schedules every reward/termination/observation/reset term and command hook in the exact fused-kernel statement order with one batched tensor call per term; per-lane SplitMix64 RNG semantics are preserved bit for bit, and reset kernels run on lane-gathered contexts with compact sim-reset writes. - Rewrite all manager MDP terms (core mdp, humanoid walk, WBT motion command incl. sequential/non-sequential/adaptive sampling, ball balance) as batched tensor implementations with fully typed signatures; command hooks are static update/advance/reset_env(receiver, ctx). - Compile each kernel's term schedule with torch.compile by default (Inductor; RNG draw order preserved). Zero-dim scalar leaves resolve to Python scalars so config branches specialize at trace time and the graphs stay unbroken; Inductor cpp.min_chunk_size is lowered to 256 (env MOTRIX_TORCH_MIN_CHUNK) because the default 4096 serialized all manager-sized loops into single-threaded omp-single regions. Observe compiles one graph per observation group over per-term contiguous buffers with a single fused assembly copy; reset compiles a static full-batch masked graph (RNG state gated by where(mask, ...) so selected lanes keep their exact SplitMix64 sequence); output tensors are graph arguments so the double-buffered slots share compiled graphs. Eager debugging via TorchManagerRuntime(..., compile=False). g1-wbt-dance 2048 envs step median (vs Numba main): 32-core desktop 16.8-17.0 vs 16.98 ms, 192-core EPYC 14.0-14.5 vs 12.99 ms (evaluate 1.66 vs 1.78 ms; remaining gap is the measured after-physics cold-kernel penalty and reset's masked 3x lane work). - Extract the backend-neutral manager resolver (manager/resolver.py) from the Numba compiler and delete the Numba path: compiler codegen and caches, NumbaTaskProgram, per-lane dispatch twins, literally helpers, TBB preload, the WBT torch pilot runtime, and the numba/tbb extras from pyproject and uv.lock. No src or test file imports numba anymore. - Rehome the surviving manager infrastructure out of the stale numba/ namespace under motrix_env_core/manager/ (term base classes, kernel_data layout package, KernelOutputs, fingerprint helper; BuildContext and ManagerContext merged into manager/context.py; the manager facade keeps its public exports). Shared batch quaternion math moved to motrix_env_core/math/quaternion_array.py and namespace.py. - Annotate the term protocol with its runtime types: reset sim_writes buffers are Map[torch.Tensor] (zero-copy views over the numpy write-program buffers), observation outputs are torch.Tensor slices, and reset noise amplitudes travel as prebuilt torch tensors instead of scalar tuples. - Update tests (lifecycle contracts migrated to the torch executor; numba-only suites removed) and wiki design docs (torch-executor.md, runtime.md, kernel-abi.md, plan) for the single-executor design. Validated with same-seed parity against the previous Numba path (core bit-exact; walk 2.4e-7; WBT 2.4e-7 across all advance/reset branches; ball balance 3e-8; RNG bit-identical), 755 fast-suite tests, and a registry smoke of all 15 manager environments (init + steps + partial reset).

Its experiment (one batch source across NumPy/Torch to decide the unified-executor direction) concluded with the torch executor landing; the shared-math contract stays protected by test_batch_mdp.py.

The torch executor PR removed the tbb extra, the numba pytest marker, and the numba kernel cache; align tests.yml (fast/slow suites, no cache step), the workflow README, and one unused-local ruff error.

The eager NumPy/Torch primitive dispatch (array/primitives.py, mx, array_api Array wrapper) served the deleted dual-lowering architecture; production manager terms call torch ops directly. Drop it with its tests and facade exports; align the CI workflow README.

…functions The term-entry marker was redundant with the resolver's signature contract checks (ManagerContext / torch.Tensor / Map[torch.Tensor] annotations, positional arity). @dispatch remains solely as the marker for kernel_data class methods that join the layout method proxy table and fingerprint (ActionState views, CommandTerm hooks).

…it is The runtime already hands ndarray term arguments over as cached tensor views; rebuilding them through np.asarray/torch.as_tensor per step was a defensive leftover and cost two needless conversions.

…w use

Leftover typing surface from the deleted dual-lowering architecture; the context contract is expressed by ManagerContext (resolver-validated signature annotation) and the executor's runtime contexts.

The Any annotation was a dual-executor migration leftover; the Numba program is gone and TorchTaskProgram is the only implementation. Its kernel fields are now typed as callables with their full signatures.

The single-attribute structural Protocol was another dual-executor leftover; TorchTaskProgram is the only implementation.

The ndarray/tensor/query branches were identity pass-throughs next to the scalar branch; one accepted-leaves tuple expresses the same contract with less ceremony.

The namespace module predates inlining the batch cores and had exactly one consumer left; the Dynamo-safe isinstance dispatch now lives beside its only caller.

binding.py collapsed to a one-line ManagerProgramResolver delegation after the resolver extraction; call it directly from _build_task_program.

torch.tensor accepts sequences; routing them through numpy was a migration habit from the numba scalar-tuple literals.

…dManagerProgram

…er graph The eager fallback duplicated the evaluate/observe schedules verbatim and reset kept an entirely different legacy gather implementation, so debug mode exercised different code than production. Builders now return the raw closure and apply torch.compile only when enabled; the eager path runs the same closure uncompiled.

…t kernel ABI The flat input layer existed for the deleted Numba lane ABI and lane gathering; the torch executor consumes the manager context and term arguments as resolved trees (path -> leaf value + scope). Removes KernelInputSource, input offsets, expression strings and their regex parser, scalar/array buffer wrappers, and the kernel_data lowered-proxy machinery; the resolver emits typed argument descriptors and a tree-structural program fingerprint of equivalent invalidation strength. Scalar-leaf Python specialization, RNG order, per-group observe graphs, masked reset, warmup clone isolation, and performance are preserved.

…m calls The generic invocation record and its kind-string dispatch become structural: the resolver emits CommandUpdateHook/CommandAdvanceHook/ CommandResetHook, RewardTermCall/TerminationTermCall (with resolver- assigned buffer indices), ObservationTermCall, and SimResetTermCall; ResolvedManagerProgram carries the evaluate phases (updates, rewards, terminations, advances), reset hooks, sim-reset terms, and per-group observations directly. Execution order, indices, arg resolution, and fingerprint invalidation are unchanged.

The masked reset cloned every PER_ENV-scoped leaf; it now clones exactly the leaves a reset term may write — each reset hook's receiver command tree, the RNG state, and metrics (lane-shaped tensors only) — and passes sim views, action state, clip corpora, and shared parameters through untouched. Reset no longer reads KernelDataScope. Verified: bit-identical compiled-vs-eager obs/terminated/rand, zero graph breaks, bench median 15.70 ms, and a full 50k-iteration FastSAC training run matching the baseline curve (final smoothed return 46.7 vs 47.3, seed-variance band).

…inery The reset write-contract change removed the last runtime consumer of leaf scopes; drop the annotation (44 sites back to plain ndarray), the KernelDataScope enum, scope fields in layouts and resolved leaves, the scope-based batch-dim validation, and the fingerprint scope hash. Existing compiled-plan caches invalidate once by design.

… tensors Carriers become plain frozen dataclasses recognized structurally (dataclass minus arrays/tensors/queries); their array fields flip from numpy to zero-copy torch views over the existing arenas — the SimBackend read program still rewrites the numpy storage in place, rand.state is an int64 view over the uint64 bits, and ArrayEnvState outputs stay numpy. Host paths (action queue, command hooks, metrics views, warmup clone) operate through the views with identical semantics. Verified: RNG state, obs, and reward streams bitwise identical to the previous commit over 21 steps + reset; per-step compiled-vs-eager parity bitwise on obs/terminated; zero graph breaks (wbt/ball); bench median 15.7-15.9 ms; full 50k-iteration FastSAC training matches the baseline curve (final smoothed return 47.16 vs 47.33, seed-variance band).

ManagerEnv now owns its outputs as torch storage — observation double-buffer and final-obs slots, reward, terminated, reward-term and termination-mask buffers — and the compiled kernels write tensors directly with no numpy boundary wrapping. ArrayEnvState is dual-typed (producer-selected: Manager emits tensors, DirectEnv keeps numpy); the lifecycle is duck-typed with allocation hooks instead of a state subclass. Provider adapters pass tensors through zero-copy (fastsac, skrl torch, rslrl), the skrl-jax path takes np.asarray views, and metrics keep the numpy contract via views. Verified: obs and RNG streams bitwise identical to the previous commit; zero graph breaks; bench median 15.8-16.0 ms; 50k-iteration FastSAC training final smoothed return 46.75 (baseline band 46-48.5); skrl JAX manager-env smoke clean; 1130 fast + 4 slow tests green.

The resolved context tree now carries torch.from_numpy views over the read-program arena at resolution time, so upper layers see one array type end-to-end; the read program still rewrites the arena in place (SimBackend untouched) and metrics keep the numpy contract.

…eArg Kernel-data carriers passed as term arguments are plain frozen dataclasses with tensor and Python-scalar fields; Dynamo traces their attribute access directly, so the namespace wrapper and per-argument tree registration are gone. The resolver fingerprints carrier contents structurally (field name/type, tensor dtype/shape/content hash, scalar values) — closing a pre-existing hole where editing carrier content did not invalidate compiled plans. The tree machinery now serves only the ManagerContext Map surface.

Term arguments arrive as prebuilt tensors or Python scalars at cfg construction time, so the runtime leaf materialization branch is gone: PreparedArg shrinks to ConstantArg | SimArg. Tensor argument contents now feed the program fingerprint (the old LeafArg path never did), closing another cache-invalidation hole.

EnvObs[T]/FinalObs[T]/ArrayEnvState[T] flow a constrained TypeVar (np.ndarray | torch.Tensor); DirectEnv pins numpy and ManagerEnv pins torch, replacing the loose unions with per-frontend precision. Metrics stay numpy-contract and are not parameterized. NpObs renamed EnvObs. Annotation-only: runtime behavior identical.

TreeClassDef is the single schema: registration zips its ordered leaves with flattened values to build the resolved path->leaf map directly. The lowering layer (record/map/leaf layouts and their flat-ABI slot_index), the Annotated-metadata machinery left over from the deleted SharedArray, and the per-type tree-def cache are gone. Schema fingerprints keep structural + leaf dtype/shape identity plus dispatch method code; one-time compiled-cache invalidation by design.

ManagerEnv.step receives torch actions end to end: the action-term pipeline (queue write, delay gather, target mapping) is pure torch and providers pass tensors straight through, removing the per-step torch->numpy conversion; the backend-facing ctrl-write buffers keep their numpy ABI behind one bridge copy. DirectEnv keeps ndarray actions; step signatures flow the array TypeVar. Action stream verified bitwise identical to the previous commit.

…amespace The four _zeros_* overrides collapse into base-class implementations over an array_api_compat namespace class attribute: numpy frontends inherit the default, ManagerEnv pins the torch namespace.

The leaves-tuple intermediate was a dead kernel-ABI artifact. tree_def_for builds the schema without walking values, iter_leaves yields (path, value) pairs in declaration order for registration and validation, and clone_tree is a direct recursive clone for canonicalize. Resolved trees, fingerprints, and clone semantics are unchanged.

cfg.__call__ allocates every per-env mutable buffer fresh and cfg-produced data is read-only by contract, so the construction-time deep clone only duplicated freshly allocated storage and each env's motion corpus. canonicalize_kernel_data is now validation-only and clone_tree is gone; warmup snapshot isolation is unchanged. A test protects multi-instance state independence, the invariant the clone used to backstop.

env.warmup had no production caller (only its own test): init_state already executes one observe pass (compiling it) and evaluate compiles at the first transition inside the trainer's own warmup steps, with Inductor's disk cache amortizing later runs. Drop warmup, the snapshot clone machinery (clone_resolved_trees/clone_kernel_value), and the ManagerWarmupResult diagnostic record.

ManagerEnv's public surface is now single-currency: reset accepts an integer torch tensor (converted once to the internal numpy control plane) and FinalObs.env_ids is a tensor; ndarray input raises with a clear message. DirectEnv keeps its ndarray currency.

ResolvedManagerProgram.fingerprint had no production consumer left: the Numba disk/process kernel caches are gone and the torch runtime caches by object identity, with torch.compile guards handling respecialization. Drop the field, the _layout_identity collection (sim inputs, tree defs, metrics, dispatch-method code), the carrier and tree-leaf fingerprint helpers, and the now-orphaned manager/fingerprint module. Tree-schema fingerprints in kernel_data stay (test-consumed schema identity).

…ernels Inductor 2.7 already emits one #pragma omp parallel region per compiled graph — the per-loop cost is the implicit barrier of each #pragma omp for, not fork/join (four g1-wbt-dance kernels carry 30/23/13/11 loops). MOTRIX_TORCH_MERGED_OMP=1 (default off) post-processes the generated C++ at CppCodeCache.load_async: per-loop pointer read/write sets (with in_out_ptr alias normalization) drive a dependency scan that keeps only barriers guarding cross-loop RAW/WAR hazards; the rest become nowait (25/30, 21/23, 11/13, 7/11 loops across the four kernels). Validation: obs and the torch RNG stream are bitwise identical on/off over 21 steps + full reset + 21 steps; graph_breaks 0; fast suites pass flag-off. EPYC: ~5-8% faster under the default 8-thread graph cap (16t 12.0->11.3, 192t 15.3->14.4 median); with uncapped MOTRIX_TORCH_COMPILE_THREADS=$T it cuts the high-thread collapse 128t 27.6->20.4ms (-26%), 192t 33.0->29.8ms (-10%). Restoring flat Numba-like scaling at 192 threads still needs cross-graph region fusion or persistent-lane kernels (future work).

iter_leaves walks the value live — cached type hints, declaration and map insertion order, cycle detection — yielding (path, annotation, value); the TreeClassDef schema classes, the schema builders, and their fingerprint hashing (consumerless since the program fingerprint chain died) are gone. Leaf order, structural TypeErrors, and validation semantics are unchanged.

functools.cache on get_type_hints returns a shared mutable dict; in the multi-process training pipeline this produced corrupted annotation lookups and collapsed FastSAC learning (return 19.5 without cache vs 3.1 with). The walk is construction-time only, so the cache bought nothing.

wlgys8 added 3 commits November 5, 2026 20:51

chore(bench): drop the concluded batch-mdp microbenchmark

f45a4d2

Its experiment (one batch source across NumPy/Torch to decide the unified-executor direction) concluded with the torch executor landing; the shared-math contract stays protected by test_batch_mdp.py.

fix(ci): drop the removed tbb/numba extras and markers from CI

97327f0

The torch executor PR removed the tbb extra, the numba pytest marker, and the numba kernel cache; align tests.yml (fast/slow suites, no cache step), the workflow README, and one unused-local ruff error.

wlgys8 force-pushed the feat/backend-neutral-manager branch from 4d2597f to 97327f0 Compare November 5, 2026 15:18

wlgys8 added 2 commits November 5, 2026 23:23

wlgys8 force-pushed the feat/backend-neutral-manager branch from 8a23fa9 to 68f97b3 Compare November 5, 2026 15:34

wlgys8 added 23 commits November 5, 2026 23:37

refactor(mdp): consume the joint-default argument as the tensor view …

8942eb1

…it is The runtime already hands ndarray term arguments over as cached tensor views; rebuilding them through np.asarray/torch.as_tensor per step was a defensive leftover and cost two needless conversions.

refactor(terrain): tidy heightfield_lookup placement, typing, and vie…

93e5a02

…w use

refactor(manager): remove the unused MdpContext/ContextMap protocols

7e97151

Leftover typing surface from the deleted dual-lowering architecture; the context contract is expressed by ManagerContext (resolver-validated signature annotation) and the executor's runtime contexts.

refactor(wbt): drop the batch_math module alias for quaternion_array

dbb1429

refactor(manager): type the task program as TorchTaskProgram

1ebb9a1

The Any annotation was a dual-executor migration leftover; the Numba program is gone and TorchTaskProgram is the only implementation. Its kernel fields are now typed as callables with their full signatures.

refactor(manager): type SimResetRuntime.apply with TorchTaskProgram

3b7e586

The single-attribute structural Protocol was another dual-executor leftover; TorchTaskProgram is the only implementation.

refactor(manager): collapse term-arg leaf checks into one whitelist

249c6d6

The ndarray/tensor/query branches were identity pass-throughs next to the scalar branch; one accepted-leaves tuple expresses the same contract with less ceremony.

refactor(math): inline the array-namespace shim into quaternion_array

71cf72c

The namespace module predates inlining the batch cores and had exactly one consumer left; the Dynamo-safe isinstance dispatch now lives beside its only caller.

refactor(manager): inline the manager binding shim into env

fae9672

binding.py collapsed to a one-line ManagerProgramResolver delegation after the resolver extraction; call it directly from _build_task_program.

refactor(ball): build reset noise tensors directly from tuples

572f739

torch.tensor accepts sequences; routing them through numpy was a migration habit from the numba scalar-tuple literals.

refactor: drop the last numpy stepping stones in reset/reward term args

f18d372

refactor(torch): type TorchManagerRuntime with ManagerEnv and Resolve…

c05ea6f

…dManagerProgram

refactor(torch): annotate the runtime's kernel and helper signatures

d100cea

refactor(torch): drop the unused outputs parameter from kernel builders

90b3f60

wlgys8 added 12 commits November 6, 2026 16:33

refactor(core)!: implement the state allocation hooks via the array n…

8dca412

…amespace The four _zeros_* overrides collapse into base-class implementations over an array_api_compat namespace class attribute: numpy frontends inherit the default, ManagerEnv pins the torch namespace.

wlgys8 force-pushed the feat/backend-neutral-manager branch from 35e28ec to a2dfea8 Compare November 6, 2026 16:16

Navigation Menu

Repository navigation

feat(manager)!: replace the Numba manager path with a generic Torch executor - #157

feat(manager)!: replace the Numba manager path with a generic Torch executor#157
wlgys8 wants to merge 41 commits into
mainfrom
feat/backend-neutral-manager

wlgys8 commented Nov 5, 2026 •

edited

Loading

Uh oh!

This branch has not been deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Conversation

wlgys8 commented Nov 5, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

摘要

主要变更

性能

验收

Uh oh!

This branch has not been deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

wlgys8 commented Nov 5, 2026 •

edited

Loading