Repository navigation
Conversation
…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.
wlgys8
force-pushed
the
feat/backend-neutral-manager
branch
from
November 5, 2026 15:18
4d2597f to
97327f0
Compare
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).
wlgys8
force-pushed
the
feat/backend-neutral-manager
branch
from
November 5, 2026 15:34
8a23fa9 to
68f97b3
Compare
…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.
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.
…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.
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
摘要
移除 ManagerEnv 的 Numba fused-kernel 执行路径，替换为通用批量 Torch 执行器；ManagerEnv host 生命周期与 term 配置 API 保持不变。
主要变更
motrix_env_core/torch_manager/runtime.py）：复用 resolver 元数据（read plan / layout / invocations），按原 fused-kernel 语句顺序批量调度全部 reward/termination/observation/reset term 与命令钩子；per-lane SplitMix64 RNG 逐位保持。默认torch.compile（evaluate/observe 静态图、reset 静态全批掩码图），compile=False可回退 eager 调试。update/advance/reset_env(receiver, ctx)，参数类型完整标注。motrix_env_core/manager/；通用四元数/地形数学归位 core（math/quaternion_array.py、heightfield_lookup）；声明式 definition lowering 死链（build_definition全家族）删除。性能（g1-wbt-dance，2048 envs）
EPYC 残余 ~1ms 归因：物理后冷 kernel 惩罚（已直接测量）+ reset 掩码设计 3x lane 计算；关键修复：Inductor
cpp.min_chunk_size4096→256（默认值会把 manager 规模循环串行化为 omp-single）。验收
python -m pytest motrix_env_core/tests motrix_envs/tests tests -m "not slow and not integration"→ 763 passed🤖 Generated with Claude Code