From 99dc4a3d9a0330dc2e2bedc455c217bbdf758f09 Mon Sep 17 00:00:00 2001 From: zilch <147668916@qq.com> Date: Thu, 5 Nov 2026 18:18:06 +0800 Subject: [PATCH 01/41] feat(manager)!: replace the Numba manager path with a generic Torch executor 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). --- AGENTS.md | 17 +- bench/bench_batch_mdp.py | 103 ++ bench/bench_physics.py | 7 - bench/bench_wbt_kernel.py | 253 ---- motrix_env_core/pyproject.toml | 16 +- .../src/motrix_env_core/__init__.py | 2 + .../src/motrix_env_core/array/__init__.py | 4 + .../src/motrix_env_core/array/mx.py | 14 + .../src/motrix_env_core/array/primitives.py | 74 ++ .../src/motrix_env_core/array_api.py | 127 +- .../src/motrix_env_core/config/sim_reset.py | 2 +- .../src/motrix_env_core/manager/__init__.py | 34 +- .../{numba => }/manager/actions.py | 8 +- .../{numba => }/manager/commands.py | 8 +- .../{numba => }/manager/context.py | 31 +- .../{numba => }/manager/dispatch.py | 7 +- .../{numba => }/manager/env.py | 186 +-- .../{numba => manager}/fingerprint.py | 0 .../{numba => manager}/kernel.py | 31 +- .../kernel_data/__init__.py | 8 +- .../kernel_data/canonical.py | 4 +- .../kernel_data/lowering.py | 10 +- .../{numba => manager}/kernel_data/map.py | 23 - .../manager/kernel_data/methods.py | 27 + .../{numba => manager}/kernel_data/tree.py | 2 +- .../motrix_env_core/manager/math/__init__.py | 8 - .../manager/math/quaternion.py | 15 - .../{numba => }/manager/metrics.py | 0 .../{numba => }/manager/observations.py | 6 +- .../{numba => }/manager/rand.py | 41 +- .../compiler.py => manager/resolver.py} | 692 +++++------ .../{numba => }/manager/rewards.py | 17 +- .../{numba => }/manager/sim_reset.py | 14 +- .../{numba => }/manager/terminations.py | 10 +- .../{numba => }/manager/terms.py | 9 +- .../src/motrix_env_core/math/__init__.py | 4 +- .../src/motrix_env_core/math/namespace.py | 13 + .../motrix_env_core/math/quaternion_array.py | 126 ++ .../src/motrix_env_core/mdp/noise.py | 36 - .../src/motrix_env_core/mdp/observations.py | 112 +- .../src/motrix_env_core/mdp/rewards.py | 47 +- .../src/motrix_env_core/mdp/state.py | 2 +- .../src/motrix_env_core/mdp/terminations.py | 33 +- .../src/motrix_env_core/mdp/terrain.py | 76 +- .../src/motrix_env_core/numba/__init__.py | 18 - .../numba/kernel_data/methods.py | 51 - .../motrix_env_core/numba/manager/__init__.py | 4 - .../numba/manager/compiler/__init__.py | 19 - .../numba/manager/compiler/cache.py | 106 -- .../numba/manager/compiler/codegen.py | 202 ---- .../numba/manager/compiler/fingerprint.py | 18 - .../numba/manager/compiler/plan.py | 112 -- .../numba/manager/compiler/program.py | 51 - .../motrix_env_core/numba/math/__init__.py | 4 - .../motrix_env_core/numba/math/quaternion.py | 175 --- .../src/motrix_env_core/numba/program.py | 18 - .../src/motrix_env_core/numba/threading.py | 59 - .../src/motrix_env_core/registry.py | 2 +- .../motrix_env_core/torch_manager/__init__.py | 8 + .../motrix_env_core/torch_manager/actions.py | 63 + .../motrix_env_core/torch_manager/binding.py | 20 + .../src/motrix_env_core/torch_manager/rand.py | 49 + .../motrix_env_core/torch_manager/runtime.py | 868 ++++++++++++++ motrix_env_core/tests/test_action_state.py | 4 +- motrix_env_core/tests/test_array_env.py | 2 +- motrix_env_core/tests/test_array_wrapper.py | 25 + motrix_env_core/tests/test_batch_mdp.py | 47 + .../tests/test_function_fingerprint.py | 113 -- motrix_env_core/tests/test_kernel_clone.py | 7 +- motrix_env_core/tests/test_kernel_data.py | 72 +- .../tests/test_kernel_data_methods.py | 275 ----- ...t_numba_manager.py => test_manager_env.py} | 421 +------ ...nv.py => test_manager_env_config_guard.py} | 2 +- .../tests/test_manager_sim_backend.py | 38 +- motrix_env_core/tests/test_mdp_noise.py | 21 - motrix_env_core/tests/test_numba_math.py | 144 --- motrix_env_core/tests/test_numba_threading.py | 10 - motrix_env_core/tests/test_registry.py | 2 +- .../tests/test_torch_manager_actions.py | 139 +++ .../tests/test_torch_manager_rand.py | 61 + motrix_env_core/tests/typecheck/numba_env.py | 18 - .../ball_balance/mdp/observations.py | 42 +- .../locomotion/ball_balance/mdp/reset.py | 125 +- .../locomotion/ball_balance/mdp/rewards.py | 56 +- .../ball_balance/mdp/terminations.py | 38 +- .../humanoid/walk_manager_mdp/command.py | 183 +-- .../humanoid/walk_manager_mdp/observations.py | 15 +- .../humanoid/walk_manager_mdp/reset.py | 150 ++- .../humanoid/walk_manager_mdp/rewards.py | 149 +-- .../motrix_envs/locomotion/wbt/mdp/command.py | 276 ++--- .../locomotion/wbt/mdp/observations.py | 136 +-- .../motrix_envs/locomotion/wbt/mdp/reset.py | 125 +- .../motrix_envs/locomotion/wbt/mdp/rewards.py | 130 +- .../locomotion/wbt/mdp/terminations.py | 111 +- motrix_envs/src/motrix_envs/motion/library.py | 2 +- motrix_envs/src/motrix_envs/motion/tracked.py | 2 +- motrix_envs/tests/test_mdp_obs.py | 68 +- motrix_envs/tests/test_wbt_batch_math.py | 55 + motrix_envs/tests/test_wbt_numba.py | 1054 ----------------- pyproject.toml | 8 - uv.lock | 99 +- wiki/design/index.md | 2 - wiki/design/manager/index.md | 8 +- .../manager/isaaclab-manager-cfg-adapter.md | 26 +- wiki/design/manager/kernel-abi.md | 75 +- wiki/design/manager/runtime.md | 66 +- wiki/design/manager/sim-read-program.md | 19 +- wiki/design/manager/task-authoring.md | 27 +- wiki/design/manager/term-required-queries.md | 157 +-- wiki/design/manager/torch-executor.md | 107 ++ 110 files changed, 3525 insertions(+), 5293 deletions(-) create mode 100644 bench/bench_batch_mdp.py delete mode 100644 bench/bench_wbt_kernel.py create mode 100644 motrix_env_core/src/motrix_env_core/array/mx.py create mode 100644 motrix_env_core/src/motrix_env_core/array/primitives.py rename motrix_env_core/src/motrix_env_core/{numba => }/manager/actions.py (93%) rename motrix_env_core/src/motrix_env_core/{numba => }/manager/commands.py (92%) rename motrix_env_core/src/motrix_env_core/{numba => }/manager/context.py (71%) rename motrix_env_core/src/motrix_env_core/{numba => }/manager/dispatch.py (74%) rename motrix_env_core/src/motrix_env_core/{numba => }/manager/env.py (84%) rename motrix_env_core/src/motrix_env_core/{numba => manager}/fingerprint.py (100%) rename motrix_env_core/src/motrix_env_core/{numba => manager}/kernel.py (86%) rename motrix_env_core/src/motrix_env_core/{numba => manager}/kernel_data/__init__.py (81%) rename motrix_env_core/src/motrix_env_core/{numba => manager}/kernel_data/canonical.py (87%) rename motrix_env_core/src/motrix_env_core/{numba => manager}/kernel_data/lowering.py (97%) rename motrix_env_core/src/motrix_env_core/{numba => manager}/kernel_data/map.py (81%) create mode 100644 motrix_env_core/src/motrix_env_core/manager/kernel_data/methods.py rename motrix_env_core/src/motrix_env_core/{numba => manager}/kernel_data/tree.py (99%) delete mode 100644 motrix_env_core/src/motrix_env_core/manager/math/__init__.py delete mode 100644 motrix_env_core/src/motrix_env_core/manager/math/quaternion.py rename motrix_env_core/src/motrix_env_core/{numba => }/manager/metrics.py (100%) rename motrix_env_core/src/motrix_env_core/{numba => }/manager/observations.py (95%) rename motrix_env_core/src/motrix_env_core/{numba => }/manager/rand.py (58%) rename motrix_env_core/src/motrix_env_core/{numba/manager/compiler/compiler.py => manager/resolver.py} (56%) rename motrix_env_core/src/motrix_env_core/{numba => }/manager/rewards.py (82%) rename motrix_env_core/src/motrix_env_core/{numba => }/manager/sim_reset.py (92%) rename motrix_env_core/src/motrix_env_core/{numba => }/manager/terminations.py (92%) rename motrix_env_core/src/motrix_env_core/{numba => }/manager/terms.py (90%) create mode 100644 motrix_env_core/src/motrix_env_core/math/namespace.py create mode 100644 motrix_env_core/src/motrix_env_core/math/quaternion_array.py delete mode 100644 motrix_env_core/src/motrix_env_core/mdp/noise.py delete mode 100644 motrix_env_core/src/motrix_env_core/numba/__init__.py delete mode 100644 motrix_env_core/src/motrix_env_core/numba/kernel_data/methods.py delete mode 100644 motrix_env_core/src/motrix_env_core/numba/manager/__init__.py delete mode 100644 motrix_env_core/src/motrix_env_core/numba/manager/compiler/__init__.py delete mode 100644 motrix_env_core/src/motrix_env_core/numba/manager/compiler/cache.py delete mode 100644 motrix_env_core/src/motrix_env_core/numba/manager/compiler/codegen.py delete mode 100644 motrix_env_core/src/motrix_env_core/numba/manager/compiler/fingerprint.py delete mode 100644 motrix_env_core/src/motrix_env_core/numba/manager/compiler/plan.py delete mode 100644 motrix_env_core/src/motrix_env_core/numba/manager/compiler/program.py delete mode 100644 motrix_env_core/src/motrix_env_core/numba/math/__init__.py delete mode 100644 motrix_env_core/src/motrix_env_core/numba/math/quaternion.py delete mode 100644 motrix_env_core/src/motrix_env_core/numba/program.py delete mode 100644 motrix_env_core/src/motrix_env_core/numba/threading.py create mode 100644 motrix_env_core/src/motrix_env_core/torch_manager/__init__.py create mode 100644 motrix_env_core/src/motrix_env_core/torch_manager/actions.py create mode 100644 motrix_env_core/src/motrix_env_core/torch_manager/binding.py create mode 100644 motrix_env_core/src/motrix_env_core/torch_manager/rand.py create mode 100644 motrix_env_core/src/motrix_env_core/torch_manager/runtime.py create mode 100644 motrix_env_core/tests/test_array_wrapper.py create mode 100644 motrix_env_core/tests/test_batch_mdp.py delete mode 100644 motrix_env_core/tests/test_function_fingerprint.py delete mode 100644 motrix_env_core/tests/test_kernel_data_methods.py rename motrix_env_core/tests/{test_numba_manager.py => test_manager_env.py} (65%) rename motrix_env_core/tests/{test_numba_env.py => test_manager_env_config_guard.py} (87%) delete mode 100644 motrix_env_core/tests/test_mdp_noise.py delete mode 100644 motrix_env_core/tests/test_numba_math.py delete mode 100644 motrix_env_core/tests/test_numba_threading.py create mode 100644 motrix_env_core/tests/test_torch_manager_actions.py create mode 100644 motrix_env_core/tests/test_torch_manager_rand.py delete mode 100644 motrix_env_core/tests/typecheck/numba_env.py create mode 100644 motrix_envs/tests/test_wbt_batch_math.py delete mode 100644 motrix_envs/tests/test_wbt_numba.py create mode 100644 wiki/design/manager/torch-executor.md diff --git a/AGENTS.md b/AGENTS.md index ed523cb8..30057716 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -76,11 +76,12 @@ 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 入口方法必须用 `@dispatch` 装饰。`@dispatch` 是透明的标记装饰器:resolver 只做签名/标注 +校验,数值执行由 Torch 执行器按整批张量调度,不存在 Numba 编译路径。Command hooks(`update`/`advance`/ +`reset_env`)是类上的 `@staticmethod` + `@dispatch` 函数,签名为 `(receiver, ctx)`。哪些方法是 term 入口、 +哪些是普通 Python 方法,由 manager 运行时契约定义(`wiki/design/manager/runtime.md`)。 ### 环境配置与注册规范 @@ -238,10 +239,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 ``` diff --git a/bench/bench_batch_mdp.py b/bench/bench_batch_mdp.py new file mode 100644 index 00000000..c8c640be --- /dev/null +++ b/bench/bench_batch_mdp.py @@ -0,0 +1,103 @@ +# Copyright Motphys Technology Co., Ltd. 2025, 2026 +# SPDX-License-Identifier: Apache-2.0 + +"""One batch MDP source executed by NumPy and Torch; no simulator required.""" + +import argparse +import json +from time import perf_counter + +import numpy as np +import torch + +from motrix_env_core.array import mx +from motrix_env_core.manager.dispatch import dispatch + + +@dispatch +def batch_rewards(position, default, weights, action, previous, velocity, command): + delta = position - default + pose = mx.sum(delta * delta * weights, axis=-1) + change = action - previous + rate = mx.sum(change * change, axis=-1) + error = velocity - command + tracking = mx.exp(-mx.sum(error * error, axis=-1)) + return pose, rate, tracking + + +def lane_rewards(position, default, weights, action, previous, velocity, command): + """Independent scalar oracle and lane-fused performance reference.""" + count, joints = position.shape + pose = np.empty(count, dtype=position.dtype) + rate = np.empty(count, dtype=position.dtype) + tracking = np.empty(count, dtype=position.dtype) + for lane in range(count): + p = np.float32(0) + r = np.float32(0) + e = np.float32(0) + for joint in range(joints): + delta = position[lane, joint] - default[joint] + p += delta * delta * weights[joint] + change = action[lane, joint] - previous[lane, joint] + r += change * change + for component in range(velocity.shape[1]): + error = velocity[lane, component] - command[lane, component] + e += error * error + pose[lane] = p + rate[lane] = r + tracking[lane] = np.exp(-e) + return pose, rate, tracking + + +def inputs(count, joints, seed=7): + rng = np.random.default_rng(seed) + shapes = ((count, joints), (joints,), (joints,), (count, joints), (count, joints), (count, 3), (count, 3)) + return tuple(rng.normal(size=shape).astype(np.float32) for shape in shapes) + + +def measure(function, args, repeats): + start = perf_counter() + result = function(*args) + first_ms = (perf_counter() - start) * 1000 + for _ in range(10): + function(*args) + samples = [] + for _ in range(7): + start = perf_counter() + for _ in range(repeats): + function(*args) + samples.append((perf_counter() - start) * 1e6 / repeats) + return result, {"first_call_ms": first_ms, "median_us": float(np.median(samples))} + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--num-envs", type=int, default=256) + parser.add_argument("--joints", type=int, default=29) + parser.add_argument("--threads", type=int, default=4) + parser.add_argument("--repeats", type=int, default=200) + options = parser.parse_args() + torch.set_num_threads(options.threads) + arrays = inputs(options.num_envs, options.joints) + tensors = tuple(torch.from_numpy(value) for value in arrays) + expected = lane_rewards(*arrays) + functions = { + "numpy_batch": (batch_rewards, arrays), + "numpy_lane": (lane_rewards, arrays), + "torch_eager_cpu": (batch_rewards, tensors), + "torch_inductor_cpu": (torch.compile(batch_rewards, fullgraph=True), tensors), + } + report = {"num_envs": options.num_envs, "joints": options.joints, "threads": options.threads, "results": {}} + for name, (function, args) in functions.items(): + result, timing = measure(function, args, options.repeats) + for actual, reference in zip(result, expected, strict=True): + if isinstance(actual, torch.Tensor): + actual = actual.detach().numpy() + np.testing.assert_allclose(actual, reference, rtol=3e-5, atol=3e-5) + assert actual.dtype == reference.dtype + report["results"][name] = timing + print(json.dumps(report, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/bench/bench_physics.py b/bench/bench_physics.py index d1028dbc..4bc88046 100644 --- a/bench/bench_physics.py +++ b/bench/bench_physics.py @@ -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 diff --git a/bench/bench_wbt_kernel.py b/bench/bench_wbt_kernel.py deleted file mode 100644 index 94907b6c..00000000 --- a/bench/bench_wbt_kernel.py +++ /dev/null @@ -1,253 +0,0 @@ -# Copyright Motphys Technology Co., Ltd. 2025, 2026 -# SPDX-License-Identifier: Apache-2.0 - -"""Microbenchmark the managed WBT fused kernel and its MotrixSim read plan.""" - -from __future__ import annotations - -import argparse -import json -import os -import platform -import subprocess -import time -from collections.abc import Callable -from dataclasses import fields, is_dataclass -from importlib.metadata import version -from typing import Any - -import numba -import numpy as np - -from motrix_env_core import registry -from motrix_env_core.manager import ManagerEnv -from motrix_env_core.numba.kernel import clone_kernel_value, get_num_threads, get_threading_layer - - -def _launch_kernel(source: np.ndarray[Any, Any], target: np.ndarray[Any, Any]) -> None: - """Small parallel loop used to expose dispatch and barrier overhead.""" - for index in numba.prange(source.shape[0]): - target[index] = source[index] + np.float32(1.0) - - -LAUNCH_KERNEL = numba.njit(cache=True, nogil=True, parallel=True)(_launch_kernel) - - -def _measure( - call: Callable[[int], None], - *, - warmup: int, - steps: int, - setup: Callable[[int], None] | None = None, - setup_delay_seconds: float = 0.0, -) -> np.ndarray[Any, Any]: - for index in range(warmup): - if setup is not None: - setup(index) - if setup_delay_seconds > 0.0: - time.sleep(setup_delay_seconds) - call(index) - - samples_ns = np.empty((steps,), dtype=np.int64) - for index in range(steps): - if setup is not None: - setup(warmup + index) - if setup_delay_seconds > 0.0: - time.sleep(setup_delay_seconds) - started_ns = time.perf_counter_ns() - call(warmup + index) - samples_ns[index] = time.perf_counter_ns() - started_ns - return samples_ns - - -def _array_bytes(value: Any) -> int: - if isinstance(value, np.ndarray): - return value.nbytes - if isinstance(value, tuple): - return sum(_array_bytes(item) for item in value) - if is_dataclass(value): - return sum(_array_bytes(getattr(value, field.name)) for field in fields(value)) - return 0 - - -def _array_checksum(value: Any) -> float: - if isinstance(value, np.ndarray): - return float(np.where(np.isfinite(value), value, 0).sum(dtype=np.float64)) - if isinstance(value, tuple): - return sum(_array_checksum(item) for item in value) - if is_dataclass(value): - return sum(_array_checksum(getattr(value, field.name)) for field in fields(value)) - return 0.0 - - -def _array_nonfinite_count(value: Any) -> int: - if isinstance(value, np.ndarray) and np.issubdtype(value.dtype, np.floating): - return int(np.count_nonzero(~np.isfinite(value))) - if isinstance(value, tuple): - return sum(_array_nonfinite_count(item) for item in value) - if is_dataclass(value): - return sum(_array_nonfinite_count(getattr(value, field.name)) for field in fields(value)) - return 0 - - -def main() -> None: - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument( - "--stage", - choices=("main", "main-after-read", "read", "partial-reset-read", "read-main", "step", "launch"), - required=True, - ) - parser.add_argument("--env", default="g1-wbt-dance", help="Registered ManagerEnv name.") - parser.add_argument("--mode", default="play") - parser.add_argument("--num-envs", type=int, default=4096) - parser.add_argument("--steps", type=int, default=200) - parser.add_argument("--warmup", type=int, default=20) - parser.add_argument("--reset-size", type=int, default=256) - parser.add_argument( - "--setup-delay-ms", - type=float, - default=0.0, - help="Untimed delay between a stage setup operation and the measured call.", - ) - parser.add_argument( - "--input-copies", - type=int, - default=1, - help="Rotate this many cloned real input snapshots to control cache residency.", - ) - parser.add_argument( - "--parallel-chunksize", - type=int, - default=0, - help="Numba parallel chunksize; zero keeps the default static schedule.", - ) - args = parser.parse_args() - if args.input_copies < 1: - parser.error("--input-copies must be at least 1") - if args.setup_delay_ms < 0.0: - parser.error("--setup-delay-ms must not be negative") - if not 1 <= args.reset_size <= args.num_envs: - parser.error("--reset-size must be between 1 and --num-envs") - - numba.set_parallel_chunksize(args.parallel_chunksize) - env = registry.make(args.env, mode=args.mode, num_envs=args.num_envs) - assert isinstance(env, ManagerEnv) - env.init_state() - env.warmup() - env._refresh_sim_reads() - inputs = env._kernel_inputs - input_snapshots = (inputs,) + tuple(clone_kernel_value(inputs) for _ in range(args.input_copies - 1)) - assert env._task_program is not None - assert env._kernel_buffers is not None - assert env._kernel_outputs is not None - - task_kernel = env._task_program.evaluate_kernel - reward_weights = env._task_program.reward_weights - buffers = env._kernel_buffers - outputs = env._kernel_outputs - latest_inputs = [inputs] - launch_source = np.arange(args.num_envs, dtype=np.float32) - launch_target = np.zeros_like(launch_source) - reset_env_ids = np.arange(args.reset_size, dtype=np.int64) - actions = np.zeros((args.num_envs, *env.action_space.shape), dtype=np.float32) - - def call_main(index: int) -> None: - task_kernel(input_snapshots[index % args.input_copies], reward_weights, buffers, outputs) - - def call_main_after_read(_index: int) -> None: - task_kernel(latest_inputs[0], reward_weights, buffers, outputs) - - def setup_read(_index: int) -> None: - env._refresh_sim_reads() - latest_inputs[0] = env._kernel_inputs - - def call_read(index: int) -> None: - setup_read(index) - - def call_partial_reset_read(_index: int) -> None: - env.sim_data.execute(reset_env_ids) - - def call_read_main(index: int) -> None: - setup_read(index) - call_main_after_read(index) - - def call_launch(_index: int) -> None: - LAUNCH_KERNEL(launch_source, launch_target) - - def call_step(_index: int) -> None: - env.step(actions) - - call = { - "main": call_main, - "main-after-read": call_main_after_read, - "read": call_read, - "partial-reset-read": call_partial_reset_read, - "read-main": call_read_main, - "step": call_step, - "launch": call_launch, - }[args.stage] - setup = {"main-after-read": setup_read}.get(args.stage) - - samples_ns = _measure( - call, - warmup=args.warmup, - steps=args.steps, - setup=setup, - setup_delay_seconds=args.setup_delay_ms / 1e3, - ) - percentiles_ns = np.percentile(samples_ns, (10, 50, 90, 99)) - sim_data = env.sim_data - kernel_values = (outputs, buffers) if args.stage in {"main", "main-after-read", "read-main", "step"} else () - checksum_values = (*(sim_data[key] for key in sim_data.keys), *kernel_values, launch_target) - result = { - "stage": args.stage, - "num_envs": args.num_envs, - "steps": args.steps, - "warmup": args.warmup, - "numba_threads": get_num_threads(), - "numba_version": numba.__version__, - "motrixsim_version": version("motrixsim"), - "git_sha": subprocess.check_output(["git", "rev-parse", "HEAD"], text=True).strip(), - "git_dirty": bool(subprocess.check_output(["git", "status", "--short"], text=True).strip()), - "parallel_chunksize": numba.get_parallel_chunksize(), - "threading_layer": get_threading_layer(), - "rayon_threads": os.environ.get("RAYON_NUM_THREADS", "default"), - "omp_wait_policy": os.environ.get("OMP_WAIT_POLICY", "default"), - "gomp_spincount": os.environ.get("GOMP_SPINCOUNT", "default"), - "affinity_cpus": sorted(os.sched_getaffinity(0)), - "hostname": platform.node(), - "pid": os.getpid(), - "input_bytes": _array_bytes(inputs), - "input_copies": args.input_copies, - "input_working_set_bytes": _array_bytes(input_snapshots), - "reset_size": args.reset_size, - "sim_queries": len(sim_data.keys), - "sim_query_layouts": [ - { - "key": key, - "query_type": type(sim_data.query(key)).__name__, - "shape": sim_data[key].shape, - "strides": sim_data[key].strides, - } - for key in sim_data.keys - ], - "arena_bytes": sim_data.arena_bytes, - "alias_count": 0, - "setup_delay_ms": args.setup_delay_ms, - "reward_weights_bytes": _array_bytes(reward_weights), - "buffer_bytes": _array_bytes(buffers), - "output_bytes": _array_bytes(outputs), - "min_ms": float(samples_ns.min() / 1e6), - "p10_ms": float(percentiles_ns[0] / 1e6), - "median_ms": float(percentiles_ns[1] / 1e6), - "p90_ms": float(percentiles_ns[2] / 1e6), - "p99_ms": float(percentiles_ns[3] / 1e6), - "max_ms": float(samples_ns.max() / 1e6), - "checksum": _array_checksum(checksum_values), - "nonfinite_count": _array_nonfinite_count(checksum_values), - } - print(json.dumps(result, sort_keys=True)) - - -if __name__ == "__main__": - main() diff --git a/motrix_env_core/pyproject.toml b/motrix_env_core/pyproject.toml index e1f352d5..0b29828e 100644 --- a/motrix_env_core/pyproject.toml +++ b/motrix_env_core/pyproject.toml @@ -30,8 +30,8 @@ classifiers = [ "Topic :: Scientific/Engineering :: Artificial Intelligence", "Typing :: Typed", ] -# Heavy runtime deps (numba/scipy for manager kernel compilation, imageio for -# video recording) are optional extras so deployment-side installs +# Heavy runtime deps (imageio for video recording) are optional extras so +# deployment-side installs # (motrix-deploy and its backend plugins) only pull the lightweight config, # registry and input layers of this package. gymnasium stays in the base # dependencies because ``motrix_env_core.base`` (on the bare package import @@ -47,11 +47,6 @@ dependencies = [ ] [project.optional-dependencies] -# Manager kernel compilation and numba-accelerated mdp terms. -numba = [ - "numba==0.61.2", - "scipy==1.15.3", # numba requires scipy at compile time for np.linalg in kernels -] # Headless video recording via VideoRecorder. video = [ "imageio>=2.30", @@ -59,13 +54,8 @@ video = [ ] # Everything the training-side workspace needs from this package. all = [ - "motrix-env-core[numba,video]", + "motrix-env-core[video]", ] -# TBB threading layer for numba parallel kernels: on many-core machines the -# default OpenMP layer pays a per-region thread-wakeup storm that dominates -# short kernels. The wheel drops libtbb.so.12 into /lib, so the runtime -# preloads it by full path (see motrix_env_core.numba.threading). -tbb = ["tbb==2021.13.0"] [project.urls] Homepage = "https://github.com/Motphys/MotrixLab" diff --git a/motrix_env_core/src/motrix_env_core/__init__.py b/motrix_env_core/src/motrix_env_core/__init__.py index a2fd0180..787b7b37 100644 --- a/motrix_env_core/src/motrix_env_core/__init__.py +++ b/motrix_env_core/src/motrix_env_core/__init__.py @@ -12,6 +12,7 @@ """ from motrix_env_core import registry +from motrix_env_core.array import Array from motrix_env_core.base import ABEnv, EnvCfg, ObsSpace, SimCfg from motrix_env_core.config import configclass from motrix_env_core.config.scene import SceneCfg @@ -30,6 +31,7 @@ __all__ = [ "ABEnv", + "Array", "BoundedGamePadPlanarVelocityBinding", "CommandBinding", "ConstantPlanarVelocityBinding", diff --git a/motrix_env_core/src/motrix_env_core/array/__init__.py b/motrix_env_core/src/motrix_env_core/array/__init__.py index 30db3877..3949e63f 100644 --- a/motrix_env_core/src/motrix_env_core/array/__init__.py +++ b/motrix_env_core/src/motrix_env_core/array/__init__.py @@ -1,3 +1,7 @@ # Copyright Motphys Technology Co., Ltd. 2025, 2026 # SPDX-License-Identifier: Apache-2.0 +from motrix_env_core.array import mx +from motrix_env_core.array_api import Array + +__all__ = ["Array", "mx"] diff --git a/motrix_env_core/src/motrix_env_core/array/mx.py b/motrix_env_core/src/motrix_env_core/array/mx.py new file mode 100644 index 00000000..3e507941 --- /dev/null +++ b/motrix_env_core/src/motrix_env_core/array/mx.py @@ -0,0 +1,14 @@ +# Copyright Motphys Technology Co., Ltd. 2025, 2026 +# SPDX-License-Identifier: Apache-2.0 + +"""Shared MDP primitive namespace. + +The functions in this module are the only array operations that shared MDP +code should call when it must run eagerly on NumPy and Torch values. Compiler +lowerings may replace these calls in compiled regions; this module itself is +kept small and transparent for host-side tests. +""" + +from motrix_env_core.array.primitives import clip, exp, sum, where, zeros_like + +__all__ = ["clip", "exp", "sum", "where", "zeros_like"] diff --git a/motrix_env_core/src/motrix_env_core/array/primitives.py b/motrix_env_core/src/motrix_env_core/array/primitives.py new file mode 100644 index 00000000..f27f0ad6 --- /dev/null +++ b/motrix_env_core/src/motrix_env_core/array/primitives.py @@ -0,0 +1,74 @@ +# Copyright Motphys Technology Co., Ltd. 2025, 2026 +# SPDX-License-Identifier: Apache-2.0 + +"""Backend-neutral primitive namespace for shared MDP functions. + +Public functions dispatch eagerly to native NumPy or Torch arrays. +""" + +from __future__ import annotations + +import numpy as np +from array_api_compat import is_torch_array + +from motrix_env_core.array_api import Array, NativeArray, unwrap_array + + +def sum(x: Array | NativeArray, axis: int | tuple[int, ...] | None = None, *, keepdims: bool = False) -> NativeArray: + x = unwrap_array(x) + if isinstance(x, np.ndarray): + return np.sum(x, axis=axis, keepdims=keepdims) + if is_torch_array(x): + import torch + + return torch.sum(x, dim=axis, keepdim=keepdims) + raise TypeError(f"Unsupported array type {type(x).__name__}.") + + +def exp(x: Array | NativeArray) -> NativeArray: + x = unwrap_array(x) + if isinstance(x, np.ndarray): + return np.exp(x) + if is_torch_array(x): + import torch + + return torch.exp(x) + raise TypeError(f"Unsupported array type {type(x).__name__}.") + + +def where(condition: Array | NativeArray, x: Array | NativeArray, y: Array | NativeArray) -> NativeArray: + condition = unwrap_array(condition) + x = unwrap_array(x) + y = unwrap_array(y) + if isinstance(condition, np.ndarray): + return np.where(condition, x, y) + if is_torch_array(condition): + import torch + + return torch.where(condition, x, y) + raise TypeError(f"Unsupported array type {type(condition).__name__}.") + + +def clip(x: Array | NativeArray, low: float | int, high: float | int) -> NativeArray: + x = unwrap_array(x) + if isinstance(x, np.ndarray): + return np.clip(x, low, high) + if is_torch_array(x): + import torch + + return torch.clamp(x, min=low, max=high) + raise TypeError(f"Unsupported array type {type(x).__name__}.") + + +def zeros_like(x: Array | NativeArray) -> NativeArray: + x = unwrap_array(x) + if isinstance(x, np.ndarray): + return np.zeros_like(x) + if is_torch_array(x): + import torch + + return torch.zeros_like(x) + raise TypeError(f"Unsupported array type {type(x).__name__}.") + + +__all__ = ["clip", "exp", "sum", "where", "zeros_like"] diff --git a/motrix_env_core/src/motrix_env_core/array_api.py b/motrix_env_core/src/motrix_env_core/array_api.py index 010c081b..9d15bdd8 100644 --- a/motrix_env_core/src/motrix_env_core/array_api.py +++ b/motrix_env_core/src/motrix_env_core/array_api.py @@ -1,41 +1,136 @@ # Copyright Motphys Technology Co., Ltd. 2025, 2026 # SPDX-License-Identifier: Apache-2.0 -"""Backend-neutral array types and operations shared across runtimes.""" +"""Backend-neutral array values and namespace helpers. + +``Array`` is a host-side value object. Compiled Manager regions must receive its +``native`` storage instead of the wrapper. +""" from __future__ import annotations from types import ModuleType -from typing import Any, TypeAlias, TypeGuard, Union +from typing import Any, TypeGuard, Union, overload import numpy as np import numpy.typing as npt from array_api_compat import array_namespace as _array_namespace from array_api_compat import is_numpy_array, is_torch_array -Array: TypeAlias = Union[npt.NDArray[Any], "torch.Tensor"] # noqa: F821 +NativeArray = Union[npt.NDArray[Any], "torch.Tensor"] # noqa: F821 -def array_namespace(*arrays: Array) -> ModuleType: - """Return the NumPy or Torch Array API namespace shared by ``arrays``.""" +class Array: + """A small host-side wrapper around a NumPy array or Torch tensor.""" + + __slots__ = ("_value",) + + def __init__(self, value: NativeArray | Array) -> None: + if isinstance(value, Array): + value = value._value + if not is_array(value): + raise TypeError(f"Array only supports NumPy arrays and Torch tensors, got {type(value).__name__}.") + self._value = value + + @property + def native(self) -> NativeArray: + """Return the backing native array without copying.""" + return self._value + + @property + def shape(self) -> tuple[int, ...]: + return tuple(self._value.shape) + + @property + def dtype(self) -> Any: + return self._value.dtype + + @property + def device(self) -> Any: + return getattr(self._value, "device", None) + + def numpy(self) -> np.ndarray: + if isinstance(self._value, np.ndarray): + return self._value + return self._value.detach().cpu().numpy() + + def torch(self, *, device: Any = None): + import torch + + if isinstance(self._value, torch.Tensor): + return self._value if device is None else self._value.to(device) + return torch.as_tensor(self._value, device=device) + + def __array_namespace__(self, *, api_version: str | None = None) -> ModuleType: + del api_version + return array_namespace(self._value) + + def __getitem__(self, index: Any) -> Array: + return Array(self._value[index]) + + def _binary(self, other: Any, operation) -> Array: + other = unwrap_array(other) + return Array(operation(self._value, other)) + + def __add__(self, other: Any) -> Array: + return self._binary(other, lambda left, right: left + right) + + def __radd__(self, other: Any) -> Array: + return self._binary(other, lambda left, right: right + left) + + def __sub__(self, other: Any) -> Array: + return self._binary(other, lambda left, right: left - right) - for array in arrays: + def __rsub__(self, other: Any) -> Array: + return self._binary(other, lambda left, right: right - left) + + def __mul__(self, other: Any) -> Array: + return self._binary(other, lambda left, right: left * right) + + def __rmul__(self, other: Any) -> Array: + return self._binary(other, lambda left, right: right * left) + + def __truediv__(self, other: Any) -> Array: + return self._binary(other, lambda left, right: left / right) + + def __rtruediv__(self, other: Any) -> Array: + return self._binary(other, lambda left, right: right / left) + + def __neg__(self) -> Array: + return Array(-self._value) + + def __repr__(self) -> str: + return f"Array({self._value!r})" + + +@overload +def unwrap_array(value: Array) -> NativeArray: ... + + +@overload +def unwrap_array(value: NativeArray) -> NativeArray: ... + + +def unwrap_array(value: Array | NativeArray) -> NativeArray: + return value.native if isinstance(value, Array) else value + + +def array_namespace(*arrays: Array | NativeArray) -> ModuleType: + """Return the NumPy or Torch Array API namespace shared by ``arrays``.""" + native = tuple(unwrap_array(array) for array in arrays) + for array in native: if not is_array(array): raise TypeError( f"array_namespace only supports NumPy arrays and Torch tensors, got {type(array).__name__}." ) + return _array_namespace(*native) - return _array_namespace(*arrays) - - -def is_array(value: object) -> TypeGuard[Array]: - """Return whether ``value`` is a supported NumPy array or Torch tensor.""" +def is_array(value: object) -> TypeGuard[NativeArray | Array]: + """Return whether ``value`` is a supported array or host wrapper.""" + if isinstance(value, Array): + return True return (is_numpy_array(value) and isinstance(value, np.ndarray)) or is_torch_array(value) -__all__ = [ - "Array", - "array_namespace", - "is_array", -] +__all__ = ["Array", "NativeArray", "array_namespace", "is_array", "unwrap_array"] diff --git a/motrix_env_core/src/motrix_env_core/config/sim_reset.py b/motrix_env_core/src/motrix_env_core/config/sim_reset.py index 176291ad..e2f508bd 100644 --- a/motrix_env_core/src/motrix_env_core/config/sim_reset.py +++ b/motrix_env_core/src/motrix_env_core/config/sim_reset.py @@ -10,7 +10,7 @@ from motrix_env_core.config.decorate import configclass if TYPE_CHECKING: - from motrix_env_core.numba.manager.sim_reset import ResetTerm + from motrix_env_core.manager.sim_reset import ResetTerm @configclass(kw_only=True) diff --git a/motrix_env_core/src/motrix_env_core/manager/__init__.py b/motrix_env_core/src/motrix_env_core/manager/__init__.py index b20208b7..2a4c4f3e 100644 --- a/motrix_env_core/src/motrix_env_core/manager/__init__.py +++ b/motrix_env_core/src/motrix_env_core/manager/__init__.py @@ -3,45 +3,48 @@ """Public API for manager-based environments and terms.""" -from numba import njit - +from motrix_env_core.array import Array, mx from motrix_env_core.config.sim_reset import ManagerResetCfg, ResetTermCfg -from motrix_env_core.numba.kernel import ManagerWarmupResult -from motrix_env_core.numba.kernel_data import SharedArray, kernel_data -from motrix_env_core.numba.manager.actions import ( +from motrix_env_core.manager.actions import ( ActionCfg, ActionState, ActionTerm, ManagerActionsCfg, ) -from motrix_env_core.numba.manager.commands import CommandCfg, CommandTerm, ManagerCommandsCfg -from motrix_env_core.numba.manager.context import ManagerContext -from motrix_env_core.numba.manager.dispatch import dispatch -from motrix_env_core.numba.manager.env import ManagerBasedEnvCfg, ManagerEnv -from motrix_env_core.numba.manager.metrics import metric -from motrix_env_core.numba.manager.observations import ( +from motrix_env_core.manager.commands import CommandCfg, CommandTerm, ManagerCommandsCfg +from motrix_env_core.manager.context import ContextMap, ManagerContext, MdpContext +from motrix_env_core.manager.dispatch import dispatch +from motrix_env_core.manager.env import ManagerBasedEnvCfg, ManagerEnv +from motrix_env_core.manager.kernel import ManagerWarmupResult +from motrix_env_core.manager.kernel_data import SharedArray, kernel_data +from motrix_env_core.manager.metrics import metric +from motrix_env_core.manager.observations import ( ManagerObservationGroupCfg, ManagerObservationsCfg, ObservationTermCfg, ObsTerm, ) -from motrix_env_core.numba.manager.rewards import ManagerRewardsCfg, RewardTerm, RewardTermCfg -from motrix_env_core.numba.manager.sim_reset import ResetTerm -from motrix_env_core.numba.manager.terminations import ( +from motrix_env_core.manager.rewards import ManagerRewardsCfg, RewardTerm, RewardTermCfg +from motrix_env_core.manager.sim_reset import ResetTerm +from motrix_env_core.manager.terminations import ( ManagerTerminationsCfg, TerminationManager, TerminationTerm, TerminationTermCfg, ) -from motrix_env_core.numba.manager.terms import BaseTerm +from motrix_env_core.manager.terms import BaseTerm from motrix_env_core.sim import SimQueriesCfg __all__ = [ "ActionCfg", "ActionTerm", + "Array", "ActionState", + "mx", "CommandCfg", "CommandTerm", + "ContextMap", + "MdpContext", "ManagerActionsCfg", "ManagerBasedEnvCfg", "ManagerCommandsCfg", @@ -68,5 +71,4 @@ "TerminationTerm", "kernel_data", "metric", - "njit", ] diff --git a/motrix_env_core/src/motrix_env_core/numba/manager/actions.py b/motrix_env_core/src/motrix_env_core/manager/actions.py similarity index 93% rename from motrix_env_core/src/motrix_env_core/numba/manager/actions.py rename to motrix_env_core/src/motrix_env_core/manager/actions.py index 2513ae3b..f116e807 100644 --- a/motrix_env_core/src/motrix_env_core/numba/manager/actions.py +++ b/motrix_env_core/src/motrix_env_core/manager/actions.py @@ -11,11 +11,11 @@ import numpy as np from motrix_env_core.config import configclass -from motrix_env_core.numba.kernel_data import SharedArray, canonicalize_kernel_data, kernel_data -from motrix_env_core.numba.manager.dispatch import dispatch +from motrix_env_core.manager.dispatch import dispatch +from motrix_env_core.manager.kernel_data import SharedArray, canonicalize_kernel_data, kernel_data if TYPE_CHECKING: - from motrix_env_core.numba.manager.env import ManagerEnv + from motrix_env_core.manager.env import ManagerEnv from motrix_env_core.sim.model import ActuatorSpec @@ -88,8 +88,8 @@ def _reset(self, env_ids: np.ndarray) -> None: class ActionCfg(abc.ABC): """Configuration that creates one environment-local action term.""" + # Actuator route: ``None`` selects all actuators; ``()`` selects none. actuator_names: tuple[str, ...] | None = None - """Actuator route: ``None`` selects all actuators; ``()`` selects none.""" @abc.abstractmethod def __call__(self, env: ManagerEnv, actuators: tuple[ActuatorSpec, ...] | None) -> ActionTerm: diff --git a/motrix_env_core/src/motrix_env_core/numba/manager/commands.py b/motrix_env_core/src/motrix_env_core/manager/commands.py similarity index 92% rename from motrix_env_core/src/motrix_env_core/numba/manager/commands.py rename to motrix_env_core/src/motrix_env_core/manager/commands.py index 0a2a9eab..77b23b3d 100644 --- a/motrix_env_core/src/motrix_env_core/numba/manager/commands.py +++ b/motrix_env_core/src/motrix_env_core/manager/commands.py @@ -10,12 +10,12 @@ import numpy as np from motrix_env_core.config import configclass -from motrix_env_core.numba.kernel_data import kernel_data -from motrix_env_core.numba.manager.dispatch import dispatch +from motrix_env_core.manager.dispatch import dispatch +from motrix_env_core.manager.kernel_data import kernel_data if TYPE_CHECKING: - from motrix_env_core.numba.manager.context import ManagerContext - from motrix_env_core.numba.manager.env import ManagerEnv + from motrix_env_core.manager.context import ManagerContext + from motrix_env_core.manager.env import ManagerEnv @dataclass diff --git a/motrix_env_core/src/motrix_env_core/numba/manager/context.py b/motrix_env_core/src/motrix_env_core/manager/context.py similarity index 71% rename from motrix_env_core/src/motrix_env_core/numba/manager/context.py rename to motrix_env_core/src/motrix_env_core/manager/context.py index e8f13aa8..4b19ceb4 100644 --- a/motrix_env_core/src/motrix_env_core/numba/manager/context.py +++ b/motrix_env_core/src/motrix_env_core/manager/context.py @@ -1,20 +1,39 @@ # Copyright Motphys Technology Co., Ltd. 2025, 2026 # SPDX-License-Identifier: Apache-2.0 +"""Logical context contract and runtime/build contexts for manager MDP functions.""" + from __future__ import annotations -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, Protocol, TypeVar import numpy as np -from motrix_env_core.numba.kernel_data import Map, kernel_data -from motrix_env_core.numba.manager.rand import RandValue +from motrix_env_core.manager.kernel_data import Map, kernel_data +from motrix_env_core.manager.rand import RandValue if TYPE_CHECKING: from motrix_env_core.base import EnvCfg - from motrix_env_core.numba.manager.env import ManagerEnv + from motrix_env_core.manager.env import ManagerEnv from motrix_env_core.sim import SimModel +T = TypeVar("T") + + +class ContextMap(Protocol[T]): + def __getitem__(self, key: str) -> T: ... + + +class MdpContext(Protocol): + """Logical context surface shared by NumPy/Numba and Torch MDP lowering.""" + + env_id: Any + actions: ContextMap[Any] + commands: ContextMap[Any] + metrics: ContextMap[Any] + sim: ContextMap[Any] + dt: Any + @kernel_data class ManagerContext: @@ -35,7 +54,7 @@ class ManagerContext: class BuildContext: - """What a term may touch while building its runtime record. + """NumPy build context retained for the current runtime builder. Handed to ``__call__`` after the model and read program are compiled. Exposes a curated surface of the environment — the term's address, the @@ -77,4 +96,4 @@ def command_terms(self) -> Any: return self._env.command_terms -__all__ = ["BuildContext", "ManagerContext"] +__all__ = ["BuildContext", "ContextMap", "ManagerContext", "MdpContext"] diff --git a/motrix_env_core/src/motrix_env_core/numba/manager/dispatch.py b/motrix_env_core/src/motrix_env_core/manager/dispatch.py similarity index 74% rename from motrix_env_core/src/motrix_env_core/numba/manager/dispatch.py rename to motrix_env_core/src/motrix_env_core/manager/dispatch.py index c4407774..f6fc6bff 100644 --- a/motrix_env_core/src/motrix_env_core/numba/manager/dispatch.py +++ b/motrix_env_core/src/motrix_env_core/manager/dispatch.py @@ -9,11 +9,12 @@ def dispatch(function: _Dispatchable) -> _Dispatchable: - """Mark a manager fused-kernel entry method for compiler dispatch. + """Mark one shared MDP function as a compiler dispatch entry. The decorator is intentionally transparent: it preserves the original Python - function so the manager compiler can inspect its signature and annotations. - Numba compilation remains owned by the manager compiler. + function so NumPy/Numba and Torch compilers can inspect the same signature and + annotations. Backend compilation remains owned by the selected manager + compiler. """ setattr(function, _DISPATCH_MARKER, True) return function diff --git a/motrix_env_core/src/motrix_env_core/numba/manager/env.py b/motrix_env_core/src/motrix_env_core/manager/env.py similarity index 84% rename from motrix_env_core/src/motrix_env_core/numba/manager/env.py rename to motrix_env_core/src/motrix_env_core/manager/env.py index dc800bc3..249bc0a1 100644 --- a/motrix_env_core/src/motrix_env_core/numba/manager/env.py +++ b/motrix_env_core/src/motrix_env_core/manager/env.py @@ -3,12 +3,11 @@ from __future__ import annotations -import abc import logging from collections.abc import Iterable from dataclasses import dataclass, field, fields from time import perf_counter -from typing import Any, TypeVar +from typing import TYPE_CHECKING, Any, TypeVar import gymnasium as gym import numpy as np @@ -16,16 +15,16 @@ from motrix_env_core.array.env import ArrayEnv, ArrayEnvState, EnvCfgType, FinalObs, NpObs from motrix_env_core.base import EnvCfg, ObsSpace from motrix_env_core.config import configclass -from motrix_env_core.config.sim_reset import ManagerResetCfg, ResetTermCfg -from motrix_env_core.numba.kernel import ( +from motrix_env_core.config.sim_reset import ManagerResetCfg +from motrix_env_core.manager.actions import ActionCfg, ActionTerm, ManagerActionsCfg +from motrix_env_core.manager.commands import CommandCfg, CommandTerm, ManagerCommandsCfg, ResetContext +from motrix_env_core.manager.kernel import ( + KernelOutputs, ManagerWarmupResult, - NumbaKernelOutputs, clone_kernel_value, - get_num_threads, - get_threading_layer, validate_kernel_context, ) -from motrix_env_core.numba.kernel_data import ( +from motrix_env_core.manager.kernel_data import ( KernelDataLayout, KernelDataScope, KernelLeafLayout, @@ -33,26 +32,24 @@ flatten_kernel_data, iter_layout_leaves, ) -from motrix_env_core.numba.manager.actions import ActionCfg, ActionTerm, ManagerActionsCfg -from motrix_env_core.numba.manager.commands import CommandCfg, CommandTerm, ManagerCommandsCfg, ResetContext -from motrix_env_core.numba.manager.compiler.plan import ManagerLayout -from motrix_env_core.numba.manager.metrics import collect_metrics, materialize_metrics -from motrix_env_core.numba.manager.observations import ( + +if TYPE_CHECKING: + from motrix_env_core.manager.resolver import ManagerLayout, ResolvedManagerProgram +from motrix_env_core.manager.metrics import collect_metrics, materialize_metrics +from motrix_env_core.manager.observations import ( ManagerObservationGroupCfg, ManagerObservationsCfg, ObservationGroupEntry, ObservationTermCfg, create_observation_groups, ) -from motrix_env_core.numba.manager.rewards import ManagerRewardsCfg, RewardTermCfg, create_reward_terms -from motrix_env_core.numba.manager.sim_reset import SimResetRuntime -from motrix_env_core.numba.manager.terminations import ( +from motrix_env_core.manager.rewards import ManagerRewardsCfg, RewardTermCfg, create_reward_terms +from motrix_env_core.manager.sim_reset import SimResetRuntime +from motrix_env_core.manager.terminations import ( ManagerTerminationsCfg, TerminationManager, TerminationTermCfg, ) -from motrix_env_core.numba.program import NumbaTaskProgram -from motrix_env_core.numba.threading import preload_tbb from motrix_env_core.sim import ( ModelQuery, PhysicsReadProgram, @@ -140,22 +137,6 @@ def command_cfgs(self) -> dict[str, CommandCfg]: def action_cfgs(self) -> dict[str, ActionCfg]: return _manager_group_to_dict(self.actions, ManagerActionsCfg, ActionCfg, label="action") - def sim_reset_cfgs(self) -> dict[str, ResetTermCfg]: - if not isinstance(self.sim_reset, ManagerResetCfg): - raise TypeError( - f"Manager simulator reset config must be a ManagerResetCfg, got {type(self.sim_reset).__name__}." - ) - return self.sim_reset.to_dict() - - def sim_query_cfgs(self) -> dict[str, SimDataQuery]: - """Return the task-declared simulator data queries. - - Term query requirements flow through ``SimDataQuery`` term arguments - and are merged at environment construction, not through this method. - """ - self.queries.validate() - return dict(self.queries.data) - def model_query_cfgs(self) -> dict[str, ModelQuery]: """Return the task-declared model queries.""" self.queries.validate() @@ -272,18 +253,10 @@ def prepare( return cls(source_name, value_type, layout, values) -@dataclass(frozen=True) -class CompiledManagerProgram(abc.ABC): - task: NumbaTaskProgram - read_plan: ReadPlan - layout: ManagerLayout - sources: tuple[str, str, str] - - @dataclass(frozen=True) class _ObsBufferSlot: observation: NpObs - kernel_outputs: NumbaKernelOutputs + kernel_outputs: KernelOutputs @classmethod def bind(cls, observation: NpObs, reward: np.ndarray, terminated: np.ndarray) -> _ObsBufferSlot: @@ -292,7 +265,7 @@ def bind(cls, observation: NpObs, reward: np.ndarray, terminated: np.ndarray) -> value = np.empty((observation.policy.shape[0], 0), dtype=observation.policy.dtype) return cls( observation=observation, - kernel_outputs=NumbaKernelOutputs( + kernel_outputs=KernelOutputs( policy_obs=observation.policy, value_obs=value, reward=reward, @@ -305,7 +278,7 @@ def bind(cls, observation: NpObs, reward: np.ndarray, terminated: np.ndarray) -> class SwapObsBuffer: """Alternate two preallocated observation slots without copying. - Each slot binds an exposed ``NpObs`` to ``NumbaKernelOutputs`` that write + Each slot binds an exposed ``NpObs`` to ``KernelOutputs`` that write into the same policy/value arrays. Swapping both together before the next kernel call keeps ``obs_t`` unchanged while a transition consumer retains it and the kernel produces ``obs_t+1`` in the other slot. @@ -332,7 +305,7 @@ def create( def current(self) -> _ObsBufferSlot: return self._slots[self._index] - def swap(self) -> tuple[NpObs, NumbaKernelOutputs]: + def swap(self) -> tuple[NpObs, KernelOutputs]: self._index = 1 - self._index slot = self.current return slot.observation, slot.kernel_outputs @@ -341,16 +314,20 @@ def swap(self) -> tuple[NpObs, NumbaKernelOutputs]: class ManagerEnv(ArrayEnv[EnvCfgType]): """NumPy environment whose task program is generated from manager config.""" - def __init__(self, cfg: EnvCfgType, num_envs: int = 1, backend: str | None = None, seed: int | None = None): - """Construct a manager environment using the selected simulator backend.""" + def __init__( + self, + cfg: EnvCfgType, + num_envs: int = 1, + backend: str | None = None, + seed: int | None = None, + ): + """Construct a manager environment using the selected simulator backend. + + The numerical task program runs on the generic Torch executor + (:class:`motrix_env_core.torch_manager.runtime.TorchManagerRuntime`). + """ if not isinstance(cfg, ManagerBasedEnvCfg): raise TypeError(f"{type(cfg).__name__} must inherit ManagerBasedEnvCfg.") - # Prefer the TBB threading layer for the parallel step kernels when the - # tbb extra is installed: the default OpenMP layer pays a thread-wakeup - # storm per kernel call that dominates short kernels on many-core - # machines. Must happen before the first parallel kernel executes. - if preload_tbb(): - logger.info("Manager env %r: numba TBB threading layer available", type(self).__name__) super().__init__(cfg, num_envs) self._rand_seed = 1 if seed is None else seed from motrix_env_core.sim.registry import create_sim_backend, default_sim_backend_name @@ -366,17 +343,17 @@ def __init__(self, cfg: EnvCfgType, num_envs: int = 1, backend: str | None = Non self._sim_reset_runtime = SimResetRuntime.create(self, cfg.sim_reset, self.sim) from motrix_env_core.mdp.state import _create_rand_value - self._task_program: NumbaTaskProgram | None = None + self._task_program: Any | None = None self._kernel_buffers: tuple[Any, ...] | None = None self.metrics: dict[str, np.ndarray] = {} - self._kernel_outputs: NumbaKernelOutputs | None = None + self._kernel_outputs: KernelOutputs | None = None self._observation_buffer: SwapObsBuffer | None = None self._final_obs_slot: _ObsBufferSlot | None = None # Persistent kernel-bound sim-reset-request buffer # (``ctx.sim_reset_requested``), cleared before each physics step and # consumed by _reset_done_envs. self._sim_reset_requested = np.zeros((self.num_envs, 1), dtype=bool) - self._compiled_manager_program: CompiledManagerProgram | None = None + self._compiled_manager_program: ResolvedManagerProgram | None = None self._action_cfgs = cfg.action_cfgs() if not self._action_cfgs: raise ValueError("Manager environment config requires at least one action config.") @@ -494,10 +471,8 @@ def init_state(self) -> ArrayEnvState: logger.info("Manager env %r startup begin: num_envs=%d", env_name, self.num_envs) self._task_program = self._build_task_program() # Initialize the simulator arena before the base lifecycle resets - # rows through the reset kernel, and compile that kernel so the - # first reset does not pay its compilation as a hidden stall. + # rows through the reset kernel. self._refresh_sim_reads() - self._precompile_reset_kernel() state = super().init_state() self._kernel_buffers = self._make_kernel_buffers(state) state.metrics = self._make_metrics_view() @@ -516,7 +491,6 @@ def init_state(self) -> ArrayEnvState: self._refresh_sim_reads() inputs = self._kernel_inputs self._validate_kernel_context(inputs) - self._compile_manager_specializations(inputs) warmup_started = perf_counter() self._execute_observe_kernel(inputs) logger.info( @@ -571,11 +545,6 @@ def _capture_final_obs(self, env_ids: np.ndarray) -> FinalObs: ), ) - def _execute_task_kernel(self, inputs: tuple[Any, ...]) -> None: - """Evaluate and observe without advancing host-side commands.""" - self._execute_evaluate_kernel(inputs) - self._execute_observe_kernel(inputs) - def _resolve_action_actuators(self) -> dict[str, tuple[ActuatorSpec, ...] | None]: routes = {} owners = {} @@ -809,32 +778,13 @@ def manager_layout(self) -> ManagerLayout: raise RuntimeError("Manager layout is available after init_state() builds the task program.") return self._compiled_manager_program.layout - def _build_task_program(self) -> NumbaTaskProgram: - from motrix_env_core.numba.manager.compiler import NumbaKernelCompiler - - self._compiled_manager_program = NumbaKernelCompiler(self).build() - return self._compiled_manager_program.task + def _build_task_program(self) -> Any: + from motrix_env_core.torch_manager.binding import build_manager_binding + from motrix_env_core.torch_manager.runtime import TorchManagerRuntime - def _precompile_reset_kernel(self) -> None: - from motrix_env_core.numba.manager.compiler import NumbaKernelCompiler - - NumbaKernelCompiler(self).precompile_reset_kernel(self._kernel_inputs) - - def _compile_manager_specializations(self, inputs: tuple[Any, ...]) -> None: - from motrix_env_core.numba.manager.compiler import NumbaKernelCompiler - - NumbaKernelCompiler(self).compile_specializations(inputs) - - def compile(self) -> float: - """Compile all manager term and kernel specializations for this environment.""" - if self._state is None: - self.init_state() - self._refresh_sim_reads() - inputs = self._kernel_inputs - self._validate_kernel_context(inputs) - started = perf_counter() - self._compile_manager_specializations(inputs) - return perf_counter() - started + self._compiled_manager_program = build_manager_binding(self) + self._torch_runtime = TorchManagerRuntime(self, self._compiled_manager_program) + return self._torch_runtime.task_program() @property def observation_groups(self) -> dict[str, ObservationGroupEntry]: @@ -878,9 +828,6 @@ def _refresh_sim_reads(self, env_ids: np.ndarray | None = None) -> None: def _validate_kernel_context(self, inputs: tuple[Any, ...]) -> None: assert self._task_program is not None - kernel = self._task_program.evaluate_kernel - if not callable(kernel) or not hasattr(kernel, "nopython_signatures"): - raise TypeError(f"{type(self).__name__} manager compiler must produce a Numba dispatcher.") assert self._kernel_buffers is not None assert self._kernel_outputs is not None try: @@ -905,6 +852,12 @@ def _kernel_input_batch_axes(self) -> tuple[bool, ...]: ) def warmup(self) -> ManagerWarmupResult: + """Run one isolated evaluate/observe pass to exercise the executor graphs. + + Inputs and buffers are cloned so the warmup cannot perturb live state; + the Torch executor compiles lazily on first real use, so only the + first-execution time is reported. + """ if self._state is None: self.init_state() assert self._state is not None @@ -916,37 +869,22 @@ def warmup(self) -> ManagerWarmupResult: self._refresh_sim_reads() inputs = self._kernel_inputs self._validate_kernel_context(inputs) - warmup_args = tuple( - clone_kernel_value(value) - for value in (inputs, self._task_program.reward_weights, self._kernel_buffers, self._kernel_outputs) - ) - observe_args = (clone_kernel_value(inputs), clone_kernel_value(self._kernel_outputs)) - started = perf_counter() - self._compile_manager_specializations(inputs) - kernel_compile_seconds = perf_counter() - started + import torch started = perf_counter() - result = self._task_program.evaluate_kernel(*warmup_args) - observe_result = self._task_program.observe_kernel(*observe_args) - first_execution_seconds = perf_counter() - started - if result is not None or observe_result is not None: - raise TypeError(f"{type(self).__name__} manager kernels must return None.") - term_compile_seconds = 0.0 - - signatures = ( - tuple(getattr(self._task_program.evaluate_kernel, "nopython_signatures", ())) - + tuple(getattr(self._task_program.observe_kernel, "nopython_signatures", ())) - + tuple(getattr(self._task_program.reset_kernel, "nopython_signatures", ())) + self._task_program.evaluate_kernel( + clone_kernel_value(inputs), + self._task_program.reward_weights, + clone_kernel_value(self._kernel_buffers), + clone_kernel_value(self._kernel_outputs), ) - if not signatures: - raise RuntimeError(f"{type(self).__name__} manager kernels did not produce nopython signatures.") - warmup = ManagerWarmupResult( - compile_seconds=term_compile_seconds + kernel_compile_seconds, - signatures=tuple(str(signature) for signature in signatures), - threading_layer=get_threading_layer(), - num_threads=get_num_threads(), - term_compile_seconds=term_compile_seconds, - kernel_compile_seconds=kernel_compile_seconds, - first_execution_seconds=first_execution_seconds, + self._task_program.observe_kernel(clone_kernel_value(inputs), clone_kernel_value(self._kernel_outputs)) + return ManagerWarmupResult( + compile_seconds=0.0, + signatures=("torch eager",), + threading_layer="openmp", + num_threads=torch.get_num_threads(), + term_compile_seconds=0.0, + kernel_compile_seconds=0.0, + first_execution_seconds=perf_counter() - started, ) - return warmup diff --git a/motrix_env_core/src/motrix_env_core/numba/fingerprint.py b/motrix_env_core/src/motrix_env_core/manager/fingerprint.py similarity index 100% rename from motrix_env_core/src/motrix_env_core/numba/fingerprint.py rename to motrix_env_core/src/motrix_env_core/manager/fingerprint.py diff --git a/motrix_env_core/src/motrix_env_core/numba/kernel.py b/motrix_env_core/src/motrix_env_core/manager/kernel.py similarity index 86% rename from motrix_env_core/src/motrix_env_core/numba/kernel.py rename to motrix_env_core/src/motrix_env_core/manager/kernel.py index 4ba54acb..0897a183 100644 --- a/motrix_env_core/src/motrix_env_core/numba/kernel.py +++ b/motrix_env_core/src/motrix_env_core/manager/kernel.py @@ -2,24 +2,15 @@ # SPDX-License-Identifier: Apache-2.0 from dataclasses import dataclass -from typing import Any, NamedTuple, Protocol, TypeVar, cast +from typing import Any, NamedTuple, TypeVar, cast -import numba import numpy as np T = TypeVar("T") -class CPUDispatcher(Protocol): - """Runtime surface used from a compiled Numba dispatcher.""" - - nopython_signatures: list[Any] - - def __call__(self, *args: Any, **kwargs: Any) -> Any: ... - - -class NumbaKernelOutputs(NamedTuple): - """Preallocated arrays written by a Numba task kernel.""" +class KernelOutputs(NamedTuple): + """Preallocated arrays written by one manager task kernel.""" policy_obs: np.ndarray value_obs: np.ndarray @@ -44,7 +35,7 @@ def validate_kernel_context( inputs: Any, reward_weights: np.ndarray, buffers: Any, - outputs: NumbaKernelOutputs, + outputs: KernelOutputs, num_envs: int, input_batch_axes: tuple[bool, ...] | None = None, ) -> None: @@ -120,7 +111,7 @@ def _validate_value(value: Any, path: str, *, num_envs: int, batched: bool, writ ) -def _validate_outputs(outputs: NumbaKernelOutputs, num_envs: int) -> None: +def _validate_outputs(outputs: KernelOutputs, num_envs: int) -> None: _validate_value(outputs, "outputs", num_envs=num_envs, batched=True, writable=True) if outputs.reward.shape != (num_envs,): raise ValueError(f"outputs.reward must have shape ({num_envs},), got {outputs.reward.shape}.") @@ -128,15 +119,3 @@ def _validate_outputs(outputs: NumbaKernelOutputs, num_envs: int) -> None: raise ValueError(f"outputs.terminated must have shape ({num_envs},), got {outputs.terminated.shape}.") if outputs.terminated.dtype != np.bool_: raise TypeError(f"outputs.terminated must use bool dtype, got {outputs.terminated.dtype}.") - - -def get_threading_layer() -> str | None: - """Return the active Numba threading layer when a parallel target initialized one.""" - try: - return numba.threading_layer() - except ValueError: - return None - - -def get_num_threads() -> int: - return numba.get_num_threads() diff --git a/motrix_env_core/src/motrix_env_core/numba/kernel_data/__init__.py b/motrix_env_core/src/motrix_env_core/manager/kernel_data/__init__.py similarity index 81% rename from motrix_env_core/src/motrix_env_core/numba/kernel_data/__init__.py rename to motrix_env_core/src/motrix_env_core/manager/kernel_data/__init__.py index 8e051282..8bc4666b 100644 --- a/motrix_env_core/src/motrix_env_core/numba/kernel_data/__init__.py +++ b/motrix_env_core/src/motrix_env_core/manager/kernel_data/__init__.py @@ -1,8 +1,8 @@ # Copyright Motphys Technology Co., Ltd. 2025, 2026 # SPDX-License-Identifier: Apache-2.0 -from motrix_env_core.numba.kernel_data.canonical import canonicalize_kernel_data -from motrix_env_core.numba.kernel_data.lowering import ( +from motrix_env_core.manager.kernel_data.canonical import canonicalize_kernel_data +from motrix_env_core.manager.kernel_data.lowering import ( KernelDataLayout, KernelDataLowering, KernelDataScope, @@ -16,8 +16,8 @@ proxy_types, rebuild_lowered, ) -from motrix_env_core.numba.kernel_data.map import Map, map_keys, map_proxy -from motrix_env_core.numba.kernel_data.tree import ( +from motrix_env_core.manager.kernel_data.map import Map, map_keys, map_proxy +from motrix_env_core.manager.kernel_data.tree import ( BaseDef, FieldDef, KernelMapDef, diff --git a/motrix_env_core/src/motrix_env_core/numba/kernel_data/canonical.py b/motrix_env_core/src/motrix_env_core/manager/kernel_data/canonical.py similarity index 87% rename from motrix_env_core/src/motrix_env_core/numba/kernel_data/canonical.py rename to motrix_env_core/src/motrix_env_core/manager/kernel_data/canonical.py index 9ca83a8f..1918df79 100644 --- a/motrix_env_core/src/motrix_env_core/numba/kernel_data/canonical.py +++ b/motrix_env_core/src/motrix_env_core/manager/kernel_data/canonical.py @@ -5,8 +5,8 @@ import numpy as np -from motrix_env_core.numba.kernel_data.lowering import KernelDataLowering -from motrix_env_core.numba.kernel_data.tree import ( +from motrix_env_core.manager.kernel_data.lowering import KernelDataLowering +from motrix_env_core.manager.kernel_data.tree import ( flatten_kernel_data, is_kernel_data, unflatten_kernel_data, diff --git a/motrix_env_core/src/motrix_env_core/numba/kernel_data/lowering.py b/motrix_env_core/src/motrix_env_core/manager/kernel_data/lowering.py similarity index 97% rename from motrix_env_core/src/motrix_env_core/numba/kernel_data/lowering.py rename to motrix_env_core/src/motrix_env_core/manager/kernel_data/lowering.py index fb56f528..d8adc007 100644 --- a/motrix_env_core/src/motrix_env_core/numba/kernel_data/lowering.py +++ b/motrix_env_core/src/motrix_env_core/manager/kernel_data/lowering.py @@ -11,10 +11,10 @@ import numpy as np -from motrix_env_core.numba.fingerprint import function_fingerprint -from motrix_env_core.numba.kernel_data.map import map_proxy -from motrix_env_core.numba.kernel_data.methods import dispatch_methods, register_proxy_method -from motrix_env_core.numba.kernel_data.tree import KernelMapDef, LeafDef, TreeClassDef +from motrix_env_core.manager.fingerprint import function_fingerprint +from motrix_env_core.manager.kernel_data.map import map_proxy +from motrix_env_core.manager.kernel_data.methods import dispatch_methods +from motrix_env_core.manager.kernel_data.tree import KernelMapDef, LeafDef, TreeClassDef _LOWERING_SCHEMA_VERSION = 2 _LAYOUT_CACHE: dict[ @@ -392,8 +392,6 @@ def _lowered_proxy( lowered.__module__ = logical_type.__module__ lowered.__qualname__ = name setattr(module, name, lowered) - for method_name, method in dispatch_methods(logical_type).items(): - register_proxy_method(lowered, method_name, method) return lowered @staticmethod diff --git a/motrix_env_core/src/motrix_env_core/numba/kernel_data/map.py b/motrix_env_core/src/motrix_env_core/manager/kernel_data/map.py similarity index 81% rename from motrix_env_core/src/motrix_env_core/numba/kernel_data/map.py rename to motrix_env_core/src/motrix_env_core/manager/kernel_data/map.py index a28491ed..1b23402f 100644 --- a/motrix_env_core/src/motrix_env_core/numba/kernel_data/map.py +++ b/motrix_env_core/src/motrix_env_core/manager/kernel_data/map.py @@ -4,14 +4,10 @@ from __future__ import annotations import hashlib -import operator import sys from collections.abc import Iterator, Mapping from typing import Any, Generic, NamedTuple, TypeVar -from numba import types -from numba.extending import overload - _MAP_KEYS = "__motrix_kernel_data_map_keys__" _MAP_SCHEMA = "__motrix_kernel_data_map_schema__" @@ -89,25 +85,6 @@ def map_keys(value_type: type[Any]) -> tuple[str, ...] | None: return keys if isinstance(keys, tuple) else None -@overload(operator.getitem, prefer_literal=True) -def _overload_map_getitem(value: Any, key: Any) -> Any: - instance_class = getattr(value, "instance_class", None) - keys = map_keys(instance_class) if isinstance(instance_class, type) else None - if keys is None: - return None - if not isinstance(key, types.StringLiteral): - return None - literal = key.literal_value - if literal not in keys: - raise KeyError(f"Map has no entry {literal!r}; available entries: {list(keys)}.") - index = keys.index(literal) - - def impl(value: Any, key: Any) -> Any: - return value[index] - - return impl - - __all__ = [ "Map", "map_keys", diff --git a/motrix_env_core/src/motrix_env_core/manager/kernel_data/methods.py b/motrix_env_core/src/motrix_env_core/manager/kernel_data/methods.py new file mode 100644 index 00000000..37ea4fc9 --- /dev/null +++ b/motrix_env_core/src/motrix_env_core/manager/kernel_data/methods.py @@ -0,0 +1,27 @@ +# Copyright Motphys Technology Co., Ltd. 2025, 2026 +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import inspect +from collections.abc import Callable +from typing import Any + +from motrix_env_core.manager.dispatch import _DISPATCH_MARKER + + +def dispatch_methods(logical_type: type) -> dict[str, Callable[..., Any]]: + """Resolve marked instance methods with normal Python MRO shadowing.""" + methods: dict[str, Callable[..., Any]] = {} + seen: set[str] = set() + for base in logical_type.__mro__: + for name, value in vars(base).items(): + if name in seen: + continue + seen.add(name) + if inspect.isfunction(value) and getattr(value, _DISPATCH_MARKER, False): + methods[name] = value + return methods + + +__all__ = ["dispatch_methods"] diff --git a/motrix_env_core/src/motrix_env_core/numba/kernel_data/tree.py b/motrix_env_core/src/motrix_env_core/manager/kernel_data/tree.py similarity index 99% rename from motrix_env_core/src/motrix_env_core/numba/kernel_data/tree.py rename to motrix_env_core/src/motrix_env_core/manager/kernel_data/tree.py index 9cb2862e..00ce8778 100644 --- a/motrix_env_core/src/motrix_env_core/numba/kernel_data/tree.py +++ b/motrix_env_core/src/motrix_env_core/manager/kernel_data/tree.py @@ -9,7 +9,7 @@ from typing_extensions import dataclass_transform -from motrix_env_core.numba.kernel_data.map import Map +from motrix_env_core.manager.kernel_data.map import Map _KernelDataType = TypeVar("_KernelDataType", bound=type) _MARKER = "__motrix_kernel_data__" diff --git a/motrix_env_core/src/motrix_env_core/manager/math/__init__.py b/motrix_env_core/src/motrix_env_core/manager/math/__init__.py deleted file mode 100644 index 06d70aa0..00000000 --- a/motrix_env_core/src/motrix_env_core/manager/math/__init__.py +++ /dev/null @@ -1,8 +0,0 @@ -# Copyright Motphys Technology Co., Ltd. 2025, 2026 -# SPDX-License-Identifier: Apache-2.0 - -"""Math helpers supported in manager term computation.""" - -from motrix_env_core.manager.math import quaternion - -__all__ = ["quaternion"] diff --git a/motrix_env_core/src/motrix_env_core/manager/math/quaternion.py b/motrix_env_core/src/motrix_env_core/manager/math/quaternion.py deleted file mode 100644 index b1c204a7..00000000 --- a/motrix_env_core/src/motrix_env_core/manager/math/quaternion.py +++ /dev/null @@ -1,15 +0,0 @@ -# Copyright Motphys Technology Co., Ltd. 2025, 2026 -# SPDX-License-Identifier: Apache-2.0 - -"""Quaternion helpers supported in manager term computation.""" - -from motrix_env_core.numba.math.quaternion import ( - inverse, - mul, - rotate_inverse, - rotate_vector, - rotation_distance, - to_matrix_first_two_rows, -) - -__all__ = ["inverse", "mul", "rotate_inverse", "rotate_vector", "rotation_distance", "to_matrix_first_two_rows"] diff --git a/motrix_env_core/src/motrix_env_core/numba/manager/metrics.py b/motrix_env_core/src/motrix_env_core/manager/metrics.py similarity index 100% rename from motrix_env_core/src/motrix_env_core/numba/manager/metrics.py rename to motrix_env_core/src/motrix_env_core/manager/metrics.py diff --git a/motrix_env_core/src/motrix_env_core/numba/manager/observations.py b/motrix_env_core/src/motrix_env_core/manager/observations.py similarity index 95% rename from motrix_env_core/src/motrix_env_core/numba/manager/observations.py rename to motrix_env_core/src/motrix_env_core/manager/observations.py index b10c0a49..f3b7ce93 100644 --- a/motrix_env_core/src/motrix_env_core/numba/manager/observations.py +++ b/motrix_env_core/src/motrix_env_core/manager/observations.py @@ -12,11 +12,11 @@ import numpy as np from motrix_env_core.config import configclass -from motrix_env_core.numba.manager.context import BuildContext -from motrix_env_core.numba.manager.terms import BaseTerm, canonicalize_term_args +from motrix_env_core.manager.context import BuildContext +from motrix_env_core.manager.terms import BaseTerm, canonicalize_term_args if TYPE_CHECKING: - from motrix_env_core.numba.manager.env import ManagerBasedEnvCfg, ManagerEnv + from motrix_env_core.manager.env import ManagerBasedEnvCfg, ManagerEnv @configclass(kw_only=True) diff --git a/motrix_env_core/src/motrix_env_core/numba/manager/rand.py b/motrix_env_core/src/motrix_env_core/manager/rand.py similarity index 58% rename from motrix_env_core/src/motrix_env_core/numba/manager/rand.py rename to motrix_env_core/src/motrix_env_core/manager/rand.py index 875602b6..68a190e0 100644 --- a/motrix_env_core/src/motrix_env_core/numba/manager/rand.py +++ b/motrix_env_core/src/motrix_env_core/manager/rand.py @@ -1,14 +1,11 @@ # Copyright Motphys Technology Co., Ltd. 2025, 2026 # SPDX-License-Identifier: Apache-2.0 -"""SplitMix64 PRNG primitives for Numba manager environments.""" +"""SplitMix64 PRNG primitives for manager environments.""" -import numba import numpy as np -from numba import types -from numba.extending import overload_method -from motrix_env_core.numba.kernel_data import kernel_data +from motrix_env_core.manager.kernel_data import kernel_data _GOLDEN_RATIO_64 = np.uint64(0x9E3779B97F4A7C15) _ENV_SALT_64 = np.uint64(0xD2B74407B1CE6E93) @@ -27,38 +24,6 @@ def next_uniform(self) -> np.float32: return next_uniform(self.state) -@overload_method(types.BaseNamedTuple, "next_uniform") -def _overload_lowered_rand_value_next_uniform(value_type): - """Expose ``RandValue.next_uniform`` on its compiler-owned tuple proxy. - - Manager kernels reconstruct ``@kernel_data`` values as schema-specific - ``NamedTuple`` proxies. Numba does not inherit Python methods onto those - proxies, so the method is registered at the common named-tuple typing - boundary and selectively enabled for lowered ``RandValue`` records. - """ - proxy_type = getattr(value_type, "instance_class", None) - if not getattr(proxy_type, "__name__", "").startswith("_RandValueKernelData_"): - return None - - def impl(value_type): - return next_uniform(value_type[0]) - - return impl - - -@overload_method(types.BaseNamedTuple, "uniform_range") -def _overload_lowered_rand_value_uniform_range(value_type, low, high): - """Expose ``RandValue.uniform_range`` on its compiler-owned tuple proxy.""" - proxy_type = getattr(value_type, "instance_class", None) - if not getattr(proxy_type, "__name__", "").startswith("_RandValueKernelData_"): - return None - - def impl(value_type, low, high): - return uniform_range(value_type[0], low, high) - - return impl - - def initialize_rand_states(num_envs: int, rand_seed: int) -> np.ndarray: """Derive one deterministic, independently mutable SplitMix64 state per environment.""" @@ -76,13 +41,11 @@ def _mix_uint64_array(value: np.ndarray) -> np.ndarray: return value ^ (value >> np.uint64(31)) -@numba.njit(inline="always") def uniform_range(rng_state: np.ndarray, low: np.float32, high: np.float32) -> np.float32: """Uniform sample in ``[low, high)`` from the ``[-1, 1)`` PRNG primitive.""" return low + (high - low) * np.float32(0.5) * (next_uniform(rng_state) + np.float32(1.0)) -@numba.njit(inline="always") def next_uniform(rng_state: np.ndarray) -> np.float32: state = rng_state[0] + _GOLDEN_RATIO_64 rng_state[0] = state diff --git a/motrix_env_core/src/motrix_env_core/numba/manager/compiler/compiler.py b/motrix_env_core/src/motrix_env_core/manager/resolver.py similarity index 56% rename from motrix_env_core/src/motrix_env_core/numba/manager/compiler/compiler.py rename to motrix_env_core/src/motrix_env_core/manager/resolver.py index 264f1072..ece04dd4 100644 --- a/motrix_env_core/src/motrix_env_core/numba/manager/compiler/compiler.py +++ b/motrix_env_core/src/motrix_env_core/manager/resolver.py @@ -1,27 +1,41 @@ # Copyright Motphys Technology Co., Ltd. 2025, 2026 # SPDX-License-Identifier: Apache-2.0 -import hashlib -import importlib.util +"""Backend-neutral resolver for manager task programs. + +Lowers one constructed manager environment into its execution metadata: the +simulator read plan, the flat kernel-input layout, prepared term arguments, +and one invocation per manager term (command hooks, observations, rewards, +terminations, simulator resets). Numerical executors — currently the Torch +executor in :mod:`motrix_env_core.torch_manager` — consume this metadata and +schedule the invocations; no code generation or dispatcher compilation +happens here. +""" + +from __future__ import annotations + import inspect import logging -import pickle -import sys -from collections.abc import Callable, Mapping +from collections.abc import Mapping from dataclasses import dataclass -from pathlib import Path from time import perf_counter from types import MappingProxyType from typing import Any, get_args, get_origin, get_type_hints -import numba import numpy as np -from numba.extending import register_jitable +import torch -from motrix_env_core.numba.fingerprint import function_fingerprint -from motrix_env_core.numba.kernel_data import ( +from motrix_env_core.manager.context import ManagerContext +from motrix_env_core.manager.dispatch import is_dispatched +from motrix_env_core.manager.env import ( + KernelInputSource, + ManagerEnv, + ReadPlan, +) +from motrix_env_core.manager.kernel_data import ( KernelDataLayout, KernelDataLowering, + KernelDataScope, Map, SharedArray, flatten_kernel_data, @@ -29,59 +43,168 @@ iter_layout_leaves, kernel_data, lane_expression, - map_proxy, - proxy_symbol, - proxy_types, -) -from motrix_env_core.numba.manager.compiler.cache import ( - _KERNEL_CACHE, - _TERM_CACHE, - generated_source_path, - has_generated_kernel_cache, - invalidate_generated_cache, - invalidate_term_cache, - materialize_source, -) -from motrix_env_core.numba.manager.compiler.codegen import KernelSourceGenerator -from motrix_env_core.numba.manager.compiler.fingerprint import ( - plan_key, - type_name, -) -from motrix_env_core.numba.manager.compiler.plan import ( - ActionTermLayout, - InputSlotLayout, - ManagerLayout, - MetricFieldLayout, - ObservationGroupLayout, - ObservationTermLayout, - ScalarTermLayout, - SimInputLayout, -) -from motrix_env_core.numba.manager.compiler.program import ( - PreparedInvocation, - ResolvedManagerContext, - ResolvedSimReset, - _CompiledManagerProgram, -) -from motrix_env_core.numba.manager.context import ManagerContext -from motrix_env_core.numba.manager.dispatch import is_dispatched -from motrix_env_core.numba.manager.env import ( - CompiledManagerProgram, - KernelInputSource, - ManagerBasedEnvCfg, - ManagerEnv, - ReadPlan, ) -from motrix_env_core.numba.manager.observations import ObservationGroupEntry, ObsTerm -from motrix_env_core.numba.manager.rewards import RewardTerm -from motrix_env_core.numba.manager.terminations import TerminationTerm -from motrix_env_core.numba.program import NumbaTaskProgram +from motrix_env_core.manager.observations import ObservationGroupEntry, ObsTerm +from motrix_env_core.manager.rewards import RewardTerm +from motrix_env_core.manager.terminations import TerminationTerm from motrix_env_core.sim import PhysicsReadProgram, SimDataQuery -_SCHEMA_VERSION = 34 logger = logging.getLogger(__name__) -_KERNEL_KINDS = ("evaluate", "observe", "reset") + +# --------------------------------------------------------------------------- +# Resolved layouts +# --------------------------------------------------------------------------- + + +@dataclass(frozen=True) +class InputSlotLayout: + index: int + source: str + scope: KernelDataScope + + +@dataclass(frozen=True) +class SimInputLayout: + slot: int + key: str + shape: tuple[int, ...] + consumers: tuple[str, ...] + + +@dataclass(frozen=True) +class ActionTermLayout: + name: str + input_slice: slice + actuator_names: tuple[str, ...] | None + + +@dataclass(frozen=True) +class ObservationTermLayout: + name: str + output_slice: slice + + +@dataclass(frozen=True) +class ObservationGroupLayout: + name: str + terms: tuple[ObservationTermLayout, ...] + size: int + + +@dataclass(frozen=True) +class ScalarTermLayout: + name: str + index: int + + +@dataclass(frozen=True) +class MetricFieldLayout: + name: str + index: int + dtype: np.dtype + + +@dataclass(frozen=True) +class ManagerLayout: + inputs: tuple[InputSlotLayout, ...] + sim_inputs: tuple[SimInputLayout, ...] + actions: tuple[ActionTermLayout, ...] + observations: dict[str, ObservationGroupLayout] + rewards: tuple[ScalarTermLayout, ...] + terminations: tuple[ScalarTermLayout, ...] + metrics: tuple[MetricFieldLayout, ...] + + def dump(self) -> str: + lines = [ + "inputs:", + ] + lines.extend(f" [{slot.index}] scope={slot.scope.value} source={slot.source}" for slot in self.inputs) + lines.append("sim_inputs:") + lines.extend( + f" [{sim_input.slot}] {sim_input.key}: shape={sim_input.shape} consumers={sim_input.consumers}" + for sim_input in self.sim_inputs + ) + lines.append("actions:") + lines.extend( + f" {term.name}: input=[{term.input_slice.start}:{term.input_slice.stop}] actuators={term.actuator_names}" + for term in self.actions + ) + lines.append("observations:") + lines.extend( + f" {group.name}.{term.name}: [{term.output_slice.start}:{term.output_slice.stop}]" + for group in self.observations.values() + for term in group.terms + ) + lines.append("rewards:") + lines.extend(f" [{term.index}] {term.name}" for term in self.rewards) + lines.append("terminations:") + lines.extend(f" [{term.index}] {term.name}" for term in self.terminations) + lines.append("metrics:") + lines.extend(f" [{metric.index}] dtype={metric.dtype.name} {metric.name}" for metric in self.metrics) + return "\n".join(lines) + + +@dataclass(frozen=True) +class PreparedInvocation: + """Runtime metadata required to execute one manager term.""" + + context: str + kind: str + dispatcher: Any + receiver_key: str | None + args_expressions: tuple[str, ...] = () + args_values: tuple[Any, ...] = () + args_prepared_indices: tuple[int | None, ...] = () + output_size: int = 0 + + +@dataclass(frozen=True) +class ResolvedManagerContext: + prepared_index: int + expression: str + + +@dataclass(frozen=True) +class ResolvedSimReset: + invocation: PreparedInvocation + sim_writes_type: str + output_offset: int + output_count: int + + +@dataclass(frozen=True) +class ResolvedManagerProgram: + """Resolved execution metadata for one manager environment.""" + + read_plan: ReadPlan + layout: ManagerLayout + invocations: tuple[PreparedInvocation, ...] + prepared_terms: tuple[KernelInputSource, ...] + input_offsets: tuple[int, ...] + context: ResolvedManagerContext + + +__all__ = [ + "ActionTermLayout", + "InputSlotLayout", + "ManagerLayout", + "ManagerProgramResolver", + "MetricFieldLayout", + "ObservationGroupLayout", + "ObservationTermLayout", + "PreparedInvocation", + "ResolvedManagerContext", + "ResolvedManagerProgram", + "ResolvedSimReset", + "ScalarTermLayout", + "SimInputLayout", +] + + +# --------------------------------------------------------------------------- +# Term-argument buffer wrappers +# --------------------------------------------------------------------------- @kernel_data @@ -89,10 +212,8 @@ class TermArrayBuffer: """Shared array slot for one ndarray term argument. Arrays passed directly as manager term arguments are lowered with shared - scope (every kernel lane sees the full array): they hold env-invariant - reference data such as default poses, weights, or lookup tables. Unlike - :class:`TermScalarBuffer` the full value participates in the plan - fingerprint, mirroring kernel-data array leaves. + scope (every execution lane sees the full array): they hold env-invariant + reference data such as default poses, weights, or lookup tables. """ value: SharedArray @@ -102,11 +223,9 @@ class TermArrayBuffer: class TermScalarBuffer: """Runtime buffer slot for one numeric term argument. - Numeric tuning arguments (thresholds, scales, gains, ...) are lowered into - a fixed-dtype shared input slot: only the type/layout fingerprint enters - the plan key, and the value itself is read at runtime. This keeps - tuning-variant manager configs on one compiled plan instead of triggering - a full recompile per value. + Numeric tuning arguments (thresholds, scales, gains, ...) are lowered + into a fixed-dtype shared input slot so executors read the value at + runtime instead of baking it into a compiled artifact. """ value: np.float32 @@ -128,30 +247,26 @@ def _manager_sim_inputs(sim_data: PhysicsReadProgram) -> tuple[Mapping[str, Mana return MappingProxyType(inputs), fingerprint -class NumbaKernelCompiler: - """Lower one fixed manager config to a fused Numba task program.""" +# --------------------------------------------------------------------------- +# Resolver +# --------------------------------------------------------------------------- + + +class ManagerProgramResolver: + """Resolve one fixed manager config into executor-neutral metadata.""" def __init__(self, env: ManagerEnv): self._env = env + self._sim_inputs: Mapping[str, ManagerSimInput] self._sim_inputs, sim_input_fingerprint = _manager_sim_inputs(env.sim_data) self._sim_slots_by_key: dict[str, int] = {} self._prepared_terms: list[KernelInputSource] = [] self._input_offsets: list[int] = [] self._flat_input_count = 0 - # Plan parts are partitioned per fused kernel so every kernel gets an - # independent plan key and disk cache: editing reward terms must not - # recompile the observation or reset kernels. Parts shared by all - # kernels (schema, context layout, sim inputs) live in _shared_parts. - self._shared_parts: list[Any] = [_SCHEMA_VERSION, numba.__version__] - self._evaluate_parts: list[Any] = [] - self._observe_parts: list[Any] = [] - self._reset_parts: list[Any] = [] - self._term_functions: dict[str, Any] = {} - self._prepared_types: dict[str, type[tuple]] = {} self._kernel_data_lowering = KernelDataLowering() - self._shared_parts.append(("sim_inputs", sim_input_fingerprint)) + self._layout_identity: list[Any] = [sim_input_fingerprint] - def build(self) -> CompiledManagerProgram: + def build(self) -> ResolvedManagerProgram: build_started = perf_counter() self._resolve_runtime_terms() observation_groups = self._env.observation_groups @@ -162,7 +277,7 @@ def build(self) -> CompiledManagerProgram: command_advances = self._resolve_command_hooks("advance") command_resets = self._resolve_command_hooks("reset_env") observation_terms, observation_layout = self._resolve_observations(observation_groups) - reward_invocations, reward_layout, reward_weights = self._resolve_rewards(self._env.cfg, reward_terms) + reward_invocations, reward_layout, _ = self._resolve_rewards(self._env.cfg, reward_terms) termination_invocations, termination_layout = self._resolve_terminations(termination_terms) reset_invocations = self._resolve_resets() per_env_metric_layout = tuple( @@ -172,47 +287,6 @@ def build(self) -> CompiledManagerProgram: ) self._validate_observation_spaces(observation_groups) - # Reward weights stay unscaled in the runtime buffer; the dt scaling - # happens inside the kernel via ctx.dt, so ctrl_dt does not need to - # participate in the plan key. - kernel_reward_weights = np.asarray(reward_weights, dtype=np.float32) - # The kernel kind participates in the key: two kernels with identical - # term sets (e.g. both empty) still have different generated sources. - part_lists = { - "evaluate": self._evaluate_parts, - "observe": self._observe_parts, - "reset": self._reset_parts, - } - plan_keys = { - kind: plan_key([("kernel_kind", kind)] + self._shared_parts + part_lists[kind]) for kind in _KERNEL_KINDS - } - logger.info( - "Manager startup %s: build plan finished in %.3fs (evaluate=%.12s observe=%.12s reset=%.12s)", - self._env_name(), - perf_counter() - build_started, - plan_keys["evaluate"], - plan_keys["observe"], - plan_keys["reset"], - ) - sources = KernelSourceGenerator(self._flat_input_count).generate( - observation_terms, - observation_layout, - reward_invocations, - termination_invocations, - context, - command_updates, - command_advances, - command_resets, - reset_invocations, - ) - kernels: dict[str, Any] = {} - generated_filenames: dict[str, str] = {} - for kind, source in zip(_KERNEL_KINDS, sources): - kernels[kind], generated_filenames[kind] = self._load_kernel(kind, source, plan_keys[kind]) - evaluate_kernel = kernels["evaluate"] - observe_kernel = kernels["observe"] - reset_kernel = kernels["reset"] - layout = ManagerLayout( inputs=tuple( InputSlotLayout( @@ -238,12 +312,6 @@ def build(self) -> CompiledManagerProgram: rewards=reward_layout, terminations=termination_layout, metrics=per_env_metric_layout, - plan_keys=(plan_keys["evaluate"], plan_keys["observe"], plan_keys["reset"]), - generated_filenames=( - generated_filenames["evaluate"], - generated_filenames["observe"], - generated_filenames["reset"], - ), ) invocations = ( command_updates @@ -252,27 +320,30 @@ def build(self) -> CompiledManagerProgram: + tuple(term for terms in observation_terms.values() for term in terms) + reward_invocations + termination_invocations + + tuple(reset.invocation for reset in reset_invocations) ) - return _CompiledManagerProgram( - task=NumbaTaskProgram( - evaluate_kernel=evaluate_kernel, - observe_kernel=observe_kernel, - reset_kernel=reset_kernel, - reward_weights=kernel_reward_weights, - ), + logger.info( + "Manager startup %s: resolve finished in %.3fs (%d invocations, %d inputs)", + type(self._env).__name__, + perf_counter() - build_started, + len(invocations), + self._flat_input_count, + ) + return ResolvedManagerProgram( read_plan=ReadPlan( self._env.sim_data, tuple(self._prepared_terms), tuple(value for prepared in self._prepared_terms for value in prepared.values), ), layout=layout, - sources=sources, invocations=invocations, prepared_terms=tuple(self._prepared_terms), input_offsets=tuple(self._input_offsets), context=context, ) + # -- resets ------------------------------------------------------------- + def _resolve_resets(self) -> tuple[ResolvedSimReset, ...]: resolved = [] output_offset = 0 @@ -302,7 +373,7 @@ def _resolve_reset( if len(parameters) != 2 + len(term.args): raise TypeError( f"Manager simulator reset term {name!r} dispatch must take a ManagerContext parameter, a " - f"Map[np.ndarray] sim-writes parameter, and {len(term.args)} positional args; got " + f"Map[torch.Tensor] sim-writes parameter, and {len(term.args)} positional args; got " f"{tuple(parameter.name for parameter in parameters)}." ) annotations = get_type_hints(function, include_extras=True) @@ -312,23 +383,13 @@ def _resolve_reset( f"ManagerContext (conventionally named ctx)." ) sim_writes_annotation = annotations.get(parameters[1].name) - if get_origin(sim_writes_annotation) is not Map or get_args(sim_writes_annotation) != (np.ndarray,): + if get_origin(sim_writes_annotation) is not Map or get_args(sim_writes_annotation) != (torch.Tensor,): raise TypeError( f"Manager simulator reset term {name!r} dispatch must annotate its second parameter as " - f"Map[np.ndarray] (conventionally named sim_writes)." + f"Map[torch.Tensor] (conventionally named sim_writes)." ) - sim_writes_proxy = map_proxy( - f"sim_reset_{name}_writes", - fields, - module_name=__name__, - schema_fingerprint=hashlib.sha256(repr(descriptors).encode()).hexdigest(), - ) - self._prepared_types[proxy_symbol(sim_writes_proxy)] = sim_writes_proxy - dispatcher = self._compile_term(function) - symbol = f"term_reset_{term_index}" - self._term_functions[symbol] = self._generated_term(function, dispatcher) + dispatcher = self._term_entry(function) expressions = [] - plan_expressions = [] prepared_indices = [] for index, value in enumerate(term.args): if is_kernel_data(value): @@ -343,24 +404,10 @@ def _resolve_reset( ) prepared_indices.append(prepared_index) expressions.append(expression) - plan_expressions.append(expression) else: - expression, _, plan_part, prepared_index = self._resolve_plain_arg( - f"sim_reset.{name}.args[{index}]", value - ) + expression, _, prepared_index = self._resolve_plain_arg(f"sim_reset.{name}.args[{index}]", value) prepared_indices.append(prepared_index) expressions.append(expression) - plan_expressions.append(plan_part) - self._reset_parts.append( - ( - "sim_reset", - name, - self._function_fingerprint(function), - plan_expressions, - repr(descriptors), - fields, - ) - ) invocation = PreparedInvocation( "sim_reset." + name, "reset", @@ -372,11 +419,13 @@ def _resolve_reset( ) return ResolvedSimReset( invocation, - proxy_symbol(sim_writes_proxy), + f"sim_reset_{name}_writes", output_offset, len(fields), ) + # -- layouts ------------------------------------------------------------ + def _sim_input_layout(self) -> tuple[SimInputLayout, ...]: slots_by_key = self._sim_slots_by_key return tuple( @@ -392,11 +441,11 @@ def _sim_input_layout(self) -> tuple[SimInputLayout, ...]: def _resolve_runtime_terms(self) -> None: for name, action_term in self._env.action_terms.items(): _, tree_def = flatten_kernel_data(action_term.state) - self._shared_parts.append( + self._layout_identity.append( ( "action", name, - self._type_name(type(action_term)), + type(action_term).__qualname__, tree_def.fingerprint, self._env.action_slices[name], None @@ -406,11 +455,11 @@ def _resolve_runtime_terms(self) -> None: ) for name, command_term in self._env.command_terms.items(): _, tree_def = flatten_kernel_data(command_term) - self._shared_parts.append( + self._layout_identity.append( ( "command", name, - self._type_name(type(command_term)), + type(command_term).__qualname__, tree_def.fingerprint, ) ) @@ -421,6 +470,7 @@ def _resolve_manager_context( reward_terms: dict[str, RewardTerm], termination_terms: dict[str, TerminationTerm], ) -> ResolvedManagerContext: + del observation_groups, reward_terms, termination_terms actions = {name: term.state for name, term in self._env.action_terms.items()} sim = {key: binding.value for key, binding in self._sim_inputs.items()} commands = dict(self._env.command_terms) @@ -440,16 +490,16 @@ def _resolve_manager_context( context="ManagerContext", ) prepared_index, expression = self._register_prepared("manager_context", value, ManagerContext, layout) - self._shared_parts.append(("manager_context", layout.fingerprint)) - self._shared_parts.append( + self._layout_identity.append(("manager_context", layout.fingerprint)) + self._layout_identity.append( ( "metrics", tuple((name, value.shape, value.dtype.str) for name, value in self._env.metrics.items()), ) ) # Map each declared sim key to its global input slot inside the - # flattened context, so term args naming a sim key can resolve to the - # key's lane view instead of crossing the kernel as a string. + # flattened context, so term args naming a sim key resolve to the + # key's batched view instead of crossing the executor as a string. context_source = self._prepared_terms[prepared_index] context_offset = self._input_offsets[prepared_index] self._sim_slots_by_key = { @@ -463,6 +513,8 @@ def _resolve_command_hooks(self, function_name: str) -> tuple[PreparedInvocation hooks = [] for index, (name, command_term) in enumerate(self._env.command_terms.items()): function = inspect.getattr_static(type(command_term), function_name) + if isinstance(function, staticmethod): + function = function.__func__ hooks.append( self._resolve_receiver_term( "command", @@ -477,6 +529,8 @@ def _resolve_command_hooks(self, function_name: str) -> tuple[PreparedInvocation ) return tuple(hooks) + # -- observations -------------------------------------------------------- + def _resolve_observations( self, groups: dict[str, ObservationGroupEntry], @@ -495,7 +549,7 @@ def _resolve_observations( offset = output_slice.stop term_index += 1 if offset != group.size: - raise RuntimeError(f"Observation group {group_name!r} resolved width changed during compilation.") + raise RuntimeError(f"Observation group {group_name!r} resolved width changed during resolution.") resolved_groups[group_name] = tuple(terms) layout_groups[group_name] = ObservationGroupLayout(group_name, tuple(layouts), group.size) return resolved_groups, layout_groups @@ -513,20 +567,19 @@ def _resolve_observation_dispatch( if len(parameters) != 2 + len(term.args): raise TypeError( f"Manager observation {group}.{name} dispatch must take a ManagerContext parameter, an " - f"np.ndarray output parameter, and {len(term.args)} positional args; got " + f"torch.Tensor output parameter, and {len(term.args)} positional args; got " f"{tuple(parameter.name for parameter in parameters)}." ) annotations = get_type_hints(function, include_extras=True) if ( annotations.get(parameters[0].name) is not ManagerContext - or annotations.get(parameters[1].name) is not np.ndarray + or annotations.get(parameters[1].name) is not torch.Tensor ): raise TypeError( f"Manager observation {group}.{name} dispatch must annotate its first two parameters as " - f"ManagerContext and np.ndarray (conventionally named ctx and out)." + f"ManagerContext and torch.Tensor (conventionally named ctx and out)." ) expressions = [] - plan_expressions = [] prepared_indices = [] warmup_values = [] for index, value in enumerate(term.args): @@ -540,23 +593,16 @@ def _resolve_observation_dispatch( ) prepared_indices.append(prepared_index) expressions.append(expression) - plan_expressions.append(expression) warmup_values.append(value) else: - expression, warmup, plan_part, prepared_index = self._resolve_plain_arg( + expression, warmup, prepared_index = self._resolve_plain_arg( f"observation.{group}.{name}.args[{index}]", value ) prepared_indices.append(prepared_index) expressions.append(expression) - plan_expressions.append(plan_part) warmup_values.append(warmup) expressions = tuple(expressions) - dispatcher = self._compile_term(function) - symbol = f"term_observation_{term_index}" - self._term_functions[symbol] = self._generated_term(function, dispatcher) - self._observe_parts.append( - (group, name, "observation_dispatch", self._function_fingerprint(function), output_size, plan_expressions) - ) + dispatcher = self._term_entry(function) return PreparedInvocation( f"{group}.{name}", "observation", @@ -578,9 +624,11 @@ def _resolve_observation( ) -> PreparedInvocation: return self._resolve_observation_dispatch(group, name, term, term_index, output_size) + # -- rewards and terminations -------------------------------------------- + def _resolve_rewards( self, - cfg: ManagerBasedEnvCfg, + cfg: Any, reward_terms: dict[str, RewardTerm], ) -> tuple[tuple[PreparedInvocation, ...], tuple[ScalarTermLayout, ...], tuple[float, ...]]: term_cfgs = cfg.reward_cfgs() @@ -593,7 +641,7 @@ def _resolve_rewards( try: weights.append(float(term_cfgs[name].weight)) except (TypeError, ValueError) as error: - raise TypeError(f"Numba reward term reward.{name} weight must be float-compatible.") from error + raise TypeError(f"Manager reward term reward.{name} weight must be float-compatible.") from error return tuple(terms), tuple(layout), tuple(weights) def _resolve_reward(self, name: str, term: RewardTerm, term_index: int) -> PreparedInvocation: @@ -636,7 +684,6 @@ def _resolve_dispatch_term( f"(conventionally named ctx)." ) expressions = [] - plan_expressions = [] prepared_indices = [] warmup_values = [] for index, value in enumerate(term.args): @@ -650,20 +697,13 @@ def _resolve_dispatch_term( ) prepared_indices.append(prepared_index) expressions.append(expression) - plan_expressions.append(expression) warmup_values.append(value) else: - expression, warmup, plan_part, prepared_index = self._resolve_plain_arg( - f"{group}.{name}.args[{index}]", value - ) + expression, warmup, prepared_index = self._resolve_plain_arg(f"{group}.{name}.args[{index}]", value) prepared_indices.append(prepared_index) expressions.append(expression) - plan_expressions.append(plan_part) warmup_values.append(warmup) - dispatcher = self._compile_term(function) - symbol = f"term_{group}_{term_index}" - self._term_functions[symbol] = self._generated_term(function, dispatcher) - self._evaluate_parts.append((group, name, "dispatch", self._function_fingerprint(function), plan_expressions)) + dispatcher = self._term_entry(function) return PreparedInvocation( f"{group}.{name}", group, @@ -682,6 +722,8 @@ def _resolve_termination( ) -> PreparedInvocation: return self._resolve_dispatch_term("termination", name, term, term_index) + # -- command hooks ------------------------------------------------------- + def _resolve_receiver_term( self, group: str, @@ -691,7 +733,7 @@ def _resolve_receiver_term( term_index: int, output_size: int, *, - function: Callable[..., Any] | None = None, + function: Any | None = None, receiver_key: str, ) -> PreparedInvocation: function = getattr(term, function_name) if function is None else function @@ -716,27 +758,15 @@ def _resolve_receiver_term( raise TypeError( f"Manager term {group}.{name} {function_name} has unsupported dependency parameters " f"{[parameter.name for parameter in argument_parameters]}; read manager and simulator data " - f"from the ManagerContext parameter." + "from the ManagerContext parameter." ) kind = f"command_{function_name}" if group == "command" else group context = f"{group}.{name} ({function.__module__}.{function.__qualname__})" - symbol = f"term_{kind}_{term_index}" - dispatcher = self._compile_term(function) - self._term_functions[symbol] = self._generated_term(function, dispatcher) - parts = self._reset_parts if function_name == "reset_env" else self._evaluate_parts - parts.append( - ( - group, - name, - kind, - function_name, - self._function_fingerprint(function), - output_size, - has_receiver, - ) - ) + dispatcher = self._term_entry(function) return PreparedInvocation(context, kind, dispatcher, receiver_key, output_size=output_size) + # -- input preparation ---------------------------------------------------- + def _register_prepared( self, source_name: str, @@ -749,16 +779,14 @@ def _register_prepared( self._prepared_terms.append(KernelInputSource.prepare(source_name, value_type, layout, value)) self._input_offsets.append(input_offset) self._flat_input_count += len(iter_layout_leaves(layout)) - for proxy in proxy_types(layout): - self._prepared_types[proxy_symbol(proxy)] = proxy return prepared_index, lane_expression(layout, input_offset) def _scalar_buffer_arg(self, source_name: str, value: float) -> tuple[str, np.float32]: """Lower one numeric term argument into a runtime scalar buffer slot. - Returns the kernel-lane argument expression and the normalized warmup - value. Only the fixed float32 layout enters the plan; the value itself - is read from the input slot at runtime. + Returns the lane argument expression and the normalized warmup value. + Only the fixed float32 layout enters the resolved metadata; the value + itself is read from the input slot at runtime. """ scalar = np.float32(value) wrapped = TermScalarBuffer(scalar) @@ -775,10 +803,8 @@ def _scalar_buffer_arg(self, source_name: str, value: float) -> tuple[str, np.fl def _array_buffer_arg(self, source_name: str, value: np.ndarray) -> tuple[str, np.ndarray]: """Lower one ndarray term argument into a per-environment input slot. - Returns the kernel-lane argument expression (the lane row) and the - warmup value. The array content participates in the plan fingerprint - via its shape and dtype, and the value is read from the input slot at - runtime. + Returns the lane argument expression (the lane row) and the warmup + value. """ wrapped = TermArrayBuffer(value=value) _, tree_def = flatten_kernel_data(wrapped) @@ -790,34 +816,32 @@ def _array_buffer_arg(self, source_name: str, value: np.ndarray) -> tuple[str, n leaf = iter_layout_leaves(layout)[0] return f"input_{self._input_offsets[prepared_index] + leaf.slot_index}", value - def _resolve_plain_arg(self, source_name: str, value: Any) -> tuple[str | None, Any, str, int | None]: + def _resolve_plain_arg(self, source_name: str, value: Any) -> tuple[str | None, Any, int | None]: """Resolve one non-kernel-data term argument. - Float-like scalars become runtime scalar buffers; every - other scalar stays a compile-time constant. Returns the kernel-lane - argument expression, the warmup value, the plan fingerprint part, and - the prepared index (``None`` for both plain scalars and buffers). + Float-like scalars become runtime scalar buffers; every other scalar + stays a plain constant. Returns the lane argument expression, the + warmup value, and the prepared index (``None`` for both plain scalars + and buffers). """ if isinstance(value, (float, np.floating)): expression, warmup = self._scalar_buffer_arg(source_name, value) - return expression, warmup, "scalar_buffer(np.float32)", None + return expression, warmup, None if isinstance(value, np.ndarray): expression, warmup = self._array_buffer_arg(source_name, value) - return expression, warmup, f"array_buffer({value.dtype.str}{value.shape})", None + return expression, warmup, None + if isinstance(value, torch.Tensor): + return f"torch.Tensor({value.dtype}{tuple(value.shape)})", value, None if isinstance(value, str): - # String arguments are inlined as source-level literals. Dispatch - # bodies keep them literal via ``numba.literally`` so Map lookups - # like ``ctx.actions[name]`` resolve at compile time. - return repr(value), value, f"str({value!r})", None + return repr(value), value, None if isinstance(value, SimDataQuery): # A query passed as a term argument resolves to its registered - # read-program lane view; the dispatch receives the array in the - # argument's position. + # read-program batched view; the dispatch receives the array in + # the argument's position. key = self._env._term_query_keys[value] slot = self._sim_slots_by_key[key] - lane_view = self._sim_inputs[key].value[0] - return f"input_{slot}[env_id]", lane_view, f"sim_input({key!r})", None - return repr(value), value, repr(value), None + return f"input_{slot}[env_id]", self._sim_inputs[key].value[0], None + return repr(value), value, None def _validate_observation_spaces(self, groups: dict[str, ObservationGroupEntry]) -> None: expected_policy = self._env.policy_observation_space.shape @@ -840,184 +864,8 @@ def _validate_observation_spaces(self, groups: dict[str, ObservationGroupEntry]) elif self._env.has_value_observation: raise ValueError("Environment declares a value observation space but manager config has no 'value' group.") - # Cache and fingerprint helpers live in sibling modules; they stay bound as - # staticmethods so callers (and tests) can intercept them on the class. - _generated_source_path = staticmethod(generated_source_path) - _materialize_source = staticmethod(materialize_source) - _has_generated_kernel_cache = staticmethod(has_generated_kernel_cache) - _invalidate_generated_cache = staticmethod(invalidate_generated_cache) - _type_name = staticmethod(type_name) - _function_fingerprint = staticmethod(function_fingerprint) - - def compile_specializations(self, inputs: tuple[Any, ...]) -> None: - """Compile all term and fused-kernel specializations for the environment.""" - started = perf_counter() - try: - self._compile_specializations_once(inputs) - except ( - AttributeError, - EOFError, - ImportError, - IndexError, - KeyError, - OSError, - pickle.UnpicklingError, - ValueError, - ) as error: - plan_keys = self._env.manager_layout.plan_keys - logger.warning("Manager specialization cache fallback: plan_keys=%s error=%s", plan_keys, error) - for plan_key in plan_keys: - self._invalidate_generated_cache(plan_key) - _KERNEL_CACHE.pop(plan_key, None) - invalidate_term_cache() - self._env._task_program = self.build().task - self._compile_specializations_once(inputs) - logger.info( - "Manager startup %s: term specializations finished in %.3fs", - self._env_name(), - perf_counter() - started, - ) - - def precompile_reset_kernel(self, inputs: tuple[Any, ...]) -> None: - """Compile the reset kernel before the first reset lifecycle call. - - The base ``init_state`` lifecycle resets rows through the reset kernel - before term specializations run; without this precompile the reset - kernel compiles lazily inside that first reset. env_ids always come - from ``np.flatnonzero`` (int64). Cache errors are deferred to the - specialization phase, whose fallback invalidates and rebuilds. - """ - task = self._env._task_program - if task is None: - return - reset_args = (inputs, np.arange(2, dtype=np.int64), self._env._sim_reset_runtime.buffers) - try: - self._compile_specialization("reset", task.reset_kernel, reset_args) - except ( - AttributeError, - EOFError, - ImportError, - IndexError, - KeyError, - OSError, - pickle.UnpicklingError, - ValueError, - ): - logger.warning("Manager reset kernel precompile deferred to specializations", exc_info=True) - - def _env_name(self) -> str: - return type(self._env).__name__ - - def _compile_specializations_once(self, inputs: tuple[Any, ...]) -> None: - task = self._env._task_program - assert task is not None - # Terms are not compiled standalone: their dispatch bodies are inlined - # into the fused kernels below, which type-checks the whole plan in one - # compilation instead of paying a separate dispatcher compile per term. - # compile() only consumes types, not values. Copying input arrays here - # changes readonly/strided views into writable contiguous arrays and - # compiles a specialization that real reset/step calls cannot reuse. - warmup_args = (inputs, task.reward_weights, self._env._kernel_buffers, self._env._kernel_outputs) - self._compile_specialization("evaluate", task.evaluate_kernel, warmup_args) - observe_args = (inputs, self._env._kernel_outputs) - self._compile_specialization("observe", task.observe_kernel, observe_args) - # Precompile the reset kernel here instead of paying its compilation on - # the first reset; env_ids always come from np.flatnonzero (int64). - reset_args = (inputs, np.arange(2, dtype=np.int64), self._env._sim_reset_runtime.buffers) - self._compile_specialization("reset", task.reset_kernel, reset_args) - - def _compile_specialization(self, kind: str, kernel: Any, args: tuple[Any, ...]) -> None: - signature = tuple(numba.typeof(arg) for arg in args) - started = perf_counter() - kernel.compile(signature) - logger.info( - "Manager startup %s: %s specialization finished in %.3fs", - self._env_name(), - kind, - perf_counter() - started, - ) - logger.debug("Manager %s specialization signature: %s", kind, signature) - - def _load_kernel(self, kind: str, source: str, plan_key: str) -> tuple[Any, str]: - """Load one fused kernel from the in-process or disk cache, else compile it.""" - cached = _KERNEL_CACHE.get(plan_key) - if cached is not None: - logger.info( - "Manager startup %s: %s kernel reused (cache=in-process, plan_key=%.12s)", - self._env_name(), - kind, - plan_key, - ) - return cached, NumbaKernelCompiler._generated_source_path(plan_key) - disk_cache = self._has_generated_kernel_cache(plan_key) - if disk_cache: - logger.info( - "Manager startup %s: %s kernel cache found, loading (cache=disk, plan_key=%.12s)", - self._env_name(), - kind, - plan_key, - ) - else: - logger.info( - "Manager startup %s: no cache for this %s kernel, first compile may take a while " - "(cache=miss, plan_key=%.12s)", - self._env_name(), - kind, - plan_key, - ) - compile_started = perf_counter() - filename = self._materialize_source(source, plan_key) - try: - kernel = self._compile_kernel(kind, source, filename) - except (ImportError, OSError, RuntimeError, TypeError, ValueError) as error: - logger.warning("Manager kernel cache fallback: plan_key=%s error=%s", plan_key, error) - self._invalidate_generated_cache(plan_key) - filename = self._materialize_source(source, plan_key) - kernel = self._compile_kernel(kind, source, filename) - _KERNEL_CACHE[plan_key] = kernel - logger.info( - "Manager startup %s: %s kernel dispatcher prepared in %.3fs (cache=%s)", - self._env_name(), - kind, - perf_counter() - compile_started, - "disk" if disk_cache else "miss", - ) - return kernel, filename - - def _compile_kernel(self, kind: str, source: str, filename: str) -> Any: - module_name = f"motrixlab_generated_{kind}_{Path(filename).stem}" - spec = importlib.util.spec_from_file_location(module_name, filename) - if spec is None or spec.loader is None: - raise ImportError(f"Unable to load generated Manager module from {filename!r}.") - module = importlib.util.module_from_spec(spec) - module.__dict__.update({"numba": numba, "np": np, **self._term_functions, **self._prepared_types}) - sys.modules[module_name] = module - spec.loader.exec_module(module) - options = {"cache": True, "nogil": True, "parallel": True} - return numba.njit(**options)(getattr(module, f"generated_{kind}_kernel")) - - @staticmethod - def _generated_term(function: Callable[..., Any], dispatcher: Any) -> Any: - module_name = function.__module__ or "" - if "." in module_name and not module_name.startswith("test") and "" not in function.__qualname__: - return register_jitable(function) - return dispatcher - - def _compile_term(self, function: Callable[..., Any]) -> Any: + def _term_entry(self, function: Any) -> Any: + """Validate one dispatch entry and keep it as the Python callable.""" if not is_dispatched(function): raise TypeError(f"Manager kernel entry {function.__qualname__!r} must be decorated with @dispatch.") - dispatcher = _TERM_CACHE.get(function) - if dispatcher is None: - # Functions declared inside a caller (notably test fixtures) do not - # have a stable importable locator, so their cache entries cannot be - # restored safely in another interpreter. - module_name = function.__module__ or "" - cache = ( - "." in module_name and not module_name.startswith("test") and "" not in function.__qualname__ - ) - dispatcher = numba.njit(cache=cache, nogil=True, inline="always")(function) - _TERM_CACHE[function] = dispatcher - return dispatcher - - -__all__ = ["NumbaKernelCompiler"] + return function diff --git a/motrix_env_core/src/motrix_env_core/numba/manager/rewards.py b/motrix_env_core/src/motrix_env_core/manager/rewards.py similarity index 82% rename from motrix_env_core/src/motrix_env_core/numba/manager/rewards.py rename to motrix_env_core/src/motrix_env_core/manager/rewards.py index f746832b..84c35f08 100644 --- a/motrix_env_core/src/motrix_env_core/numba/manager/rewards.py +++ b/motrix_env_core/src/motrix_env_core/manager/rewards.py @@ -9,11 +9,11 @@ from typing import TYPE_CHECKING, Any from motrix_env_core.config import configclass -from motrix_env_core.numba.manager.context import BuildContext -from motrix_env_core.numba.manager.terms import BaseTerm, canonicalize_term_args +from motrix_env_core.manager.context import BuildContext +from motrix_env_core.manager.terms import BaseTerm, canonicalize_term_args if TYPE_CHECKING: - from motrix_env_core.numba.manager.env import ManagerEnv + from motrix_env_core.manager.env import ManagerEnv @dataclass(frozen=True, slots=True, init=False) @@ -29,8 +29,8 @@ class RewardTermCfg(abc.ABC): weight: float @abc.abstractmethod - def __call__(self, ctx) -> RewardTerm: - """Assemble the runtime term (dispatch plus static arguments).""" + def __call__(self, ctx: BuildContext) -> RewardTerm: + """Assemble the current NumPy/Numba runtime term.""" @configclass @@ -63,4 +63,9 @@ def create_reward_terms(cfg: dict[str, RewardTermCfg], env: ManagerEnv) -> dict[ return terms -__all__ = ["ManagerRewardsCfg", "RewardTerm", "RewardTermCfg", "create_reward_terms"] +__all__ = [ + "ManagerRewardsCfg", + "RewardTerm", + "RewardTermCfg", + "create_reward_terms", +] diff --git a/motrix_env_core/src/motrix_env_core/numba/manager/sim_reset.py b/motrix_env_core/src/motrix_env_core/manager/sim_reset.py similarity index 92% rename from motrix_env_core/src/motrix_env_core/numba/manager/sim_reset.py rename to motrix_env_core/src/motrix_env_core/manager/sim_reset.py index f90f9da1..442619ba 100644 --- a/motrix_env_core/src/motrix_env_core/numba/manager/sim_reset.py +++ b/motrix_env_core/src/motrix_env_core/manager/sim_reset.py @@ -12,16 +12,20 @@ import numpy as np from motrix_env_core.config.sim_reset import ManagerResetCfg -from motrix_env_core.numba.manager.context import BuildContext -from motrix_env_core.numba.manager.terms import BaseTerm, canonicalize_term_args +from motrix_env_core.manager.context import BuildContext +from motrix_env_core.manager.terms import BaseTerm, canonicalize_term_args from motrix_env_core.perf import active_perf_scope from motrix_env_core.sim.write import SimWrite, WriteProgram if TYPE_CHECKING: - from motrix_env_core.numba.manager.env import ManagerEnv - from motrix_env_core.numba.program import NumbaTaskProgram + from typing import Protocol + + from motrix_env_core.manager.env import ManagerEnv from motrix_env_core.sim.backend import SimBackend + class _TaskProgram(Protocol): + reset_kernel: Any + @dataclass(frozen=True, slots=True, init=False) class ResetTerm(BaseTerm): @@ -95,7 +99,7 @@ def create( def apply( self, - task_program: NumbaTaskProgram, + task_program: _TaskProgram, inputs: tuple[Any, ...], env_ids: np.ndarray, ) -> None: diff --git a/motrix_env_core/src/motrix_env_core/numba/manager/terminations.py b/motrix_env_core/src/motrix_env_core/manager/terminations.py similarity index 92% rename from motrix_env_core/src/motrix_env_core/numba/manager/terminations.py rename to motrix_env_core/src/motrix_env_core/manager/terminations.py index 56d0d750..3052835e 100644 --- a/motrix_env_core/src/motrix_env_core/numba/manager/terminations.py +++ b/motrix_env_core/src/motrix_env_core/manager/terminations.py @@ -11,11 +11,11 @@ import numpy as np from motrix_env_core.config import configclass -from motrix_env_core.numba.manager.context import BuildContext -from motrix_env_core.numba.manager.terms import BaseTerm, canonicalize_term_args +from motrix_env_core.manager.context import BuildContext +from motrix_env_core.manager.terms import BaseTerm, canonicalize_term_args if TYPE_CHECKING: - from motrix_env_core.numba.manager.env import ManagerEnv + from motrix_env_core.manager.env import ManagerEnv @dataclass(frozen=True, slots=True, init=False) @@ -46,8 +46,8 @@ class TerminationTermCfg(abc.ABC): """ @abc.abstractmethod - def __call__(self, ctx) -> TerminationTerm: - """Assemble the runtime term (dispatch plus static arguments).""" + def __call__(self, ctx: BuildContext) -> TerminationTerm: + """Assemble the current NumPy/Numba runtime term.""" @configclass diff --git a/motrix_env_core/src/motrix_env_core/numba/manager/terms.py b/motrix_env_core/src/motrix_env_core/manager/terms.py similarity index 90% rename from motrix_env_core/src/motrix_env_core/numba/manager/terms.py rename to motrix_env_core/src/motrix_env_core/manager/terms.py index 5b26e568..a4fe864a 100644 --- a/motrix_env_core/src/motrix_env_core/numba/manager/terms.py +++ b/motrix_env_core/src/motrix_env_core/manager/terms.py @@ -9,11 +9,16 @@ from typing import Any import numpy as np +import torch -from motrix_env_core.numba.kernel_data import canonicalize_kernel_data, is_kernel_data +from motrix_env_core.manager.kernel_data import canonicalize_kernel_data, is_kernel_data from motrix_env_core.sim.read import SimDataQuery +def _is_torch_tensor(value: object) -> bool: + return isinstance(value, torch.Tensor) + + def canonicalize_term_args(args: tuple[Any, ...], *, context: str) -> tuple[Any, ...]: """Validate and canonicalize positional Numba-compatible term arguments. @@ -31,6 +36,8 @@ def canonicalize_term_args(args: tuple[Any, ...], *, context: str) -> tuple[Any, values.append(value) elif isinstance(value, np.ndarray): values.append(value) + elif _is_torch_tensor(value): + values.append(value) elif isinstance(value, SimDataQuery): values.append(value) else: diff --git a/motrix_env_core/src/motrix_env_core/math/__init__.py b/motrix_env_core/src/motrix_env_core/math/__init__.py index 9d495c8f..aad97831 100644 --- a/motrix_env_core/src/motrix_env_core/math/__init__.py +++ b/motrix_env_core/src/motrix_env_core/math/__init__.py @@ -1,6 +1,6 @@ # Copyright Motphys Technology Co., Ltd. 2025, 2026 # SPDX-License-Identifier: Apache-2.0 -from motrix_env_core.math import quaternion, utils +from motrix_env_core.math import namespace, quaternion, quaternion_array, utils -__all__ = ["quaternion", "utils"] +__all__ = ["namespace", "quaternion", "quaternion_array", "utils"] diff --git a/motrix_env_core/src/motrix_env_core/math/namespace.py b/motrix_env_core/src/motrix_env_core/math/namespace.py new file mode 100644 index 00000000..ea417ca1 --- /dev/null +++ b/motrix_env_core/src/motrix_env_core/math/namespace.py @@ -0,0 +1,13 @@ +# Copyright Motphys Technology Co., Ltd. 2025, 2026 +# SPDX-License-Identifier: Apache-2.0 + +"""Type-specialized namespace selection for compiled batch computations.""" + +import array_api_compat.numpy as numpy_namespace +import array_api_compat.torch as torch_namespace +import torch + + +def array_namespace(value, *others): + # Dynamo specializes this type check; no runtime backend switch in kernels. + return torch_namespace if isinstance(value, torch.Tensor) else numpy_namespace diff --git a/motrix_env_core/src/motrix_env_core/math/quaternion_array.py b/motrix_env_core/src/motrix_env_core/math/quaternion_array.py new file mode 100644 index 00000000..9f7a222a --- /dev/null +++ b/motrix_env_core/src/motrix_env_core/math/quaternion_array.py @@ -0,0 +1,126 @@ +# Copyright Motphys Technology Co., Ltd. 2025, 2026 +# SPDX-License-Identifier: Apache-2.0 + +"""Shared NumPy/Torch batch quaternion operations, using xyzw convention.""" + +from motrix_env_core.math.namespace import array_namespace + + +def inverse(q): + xp = array_namespace(q) + norm = xp.maximum(xp.sum(q * q, axis=-1, keepdims=True), xp.full_like(q[..., :1], 1e-12)) + return xp.concat((-q[..., :3], q[..., 3:]), axis=-1) / norm + + +def mul(lhs, rhs): + xp = array_namespace(lhs, rhs) + x1, y1, z1, w1 = (lhs[..., i] for i in range(4)) + x2, y2, z2, w2 = (rhs[..., i] for i in range(4)) + return xp.stack( + ( + w1 * x2 + x1 * w2 + y1 * z2 - z1 * y2, + w1 * y2 - x1 * z2 + y1 * w2 + z1 * x2, + w1 * z2 + x1 * y2 - y1 * x2 + z1 * w2, + w1 * w2 - x1 * x2 - y1 * y2 - z1 * z2, + ), + axis=-1, + ) + + +def rotate_inverse(q, vector): + xp = array_namespace(q, vector) + x, y, z, w = (q[..., i] for i in range(4)) + vx, vy, vz = (vector[..., i] for i in range(3)) + cx = y * vz - z * vy - w * vx + cy = z * vx - x * vz - w * vy + cz = x * vy - y * vx - w * vz + return xp.stack((vx + 2 * (y * cz - z * cy), vy + 2 * (z * cx - x * cz), vz + 2 * (x * cy - y * cx)), axis=-1) + + +def matrix6d(q): + xp = array_namespace(q) + x, y, z, w = (q[..., i] for i in range(4)) + s = 2 / xp.sum(q * q, axis=-1) + return xp.stack( + ( + 1 - s * (y * y + z * z), + s * (x * y - z * w), + s * (x * z + y * w), + s * (x * y + z * w), + 1 - s * (x * x + z * z), + s * (y * z - x * w), + ), + axis=-1, + ) + + +def relative_orientation(reference, target): + return matrix6d(mul(inverse(reference), target)) + + +def align_targets(robot_position, robot_quat, reference_position, reference_quat, body_position, body_quat): + """Yaw-align corpus bodies, anchoring robot XY and corpus reference Z.""" + xp = array_namespace(robot_position) + delta = mul(robot_quat, inverse(reference_quat)) + x, y, z, w = (delta[..., i] for i in range(4)) + yaw = xp.atan2(2 * (w * z + x * y), 1 - 2 * (y * y + z * z)) + zero = xp.zeros_like(yaw) + yaw_quat = xp.stack((zero, zero, xp.sin(yaw / 2), xp.cos(yaw / 2)), axis=-1) + anchor = xp.concat((robot_position[..., :2], reference_position[..., 2:]), axis=-1) + # Forward rotation uses the same unit-quaternion algebra as scalar helper. + conjugate = xp.concat((-yaw_quat[..., :3], yaw_quat[..., 3:]), axis=-1) + aligned_position = rotate_inverse(conjugate[:, None, :], body_position - reference_position[:, None, :]) + return aligned_position + anchor[:, None, :], mul(yaw_quat[:, None, :], body_quat) + + +def rotation_distance(lhs, rhs): + """Geodesic angle between unit quaternions, without normalization.""" + xp = array_namespace(lhs, rhs) + x1, y1, z1, w1 = (lhs[..., i] for i in range(4)) + x2, y2, z2, w2 = (rhs[..., i] for i in range(4)) + dx = -w1 * x2 + x1 * w2 - y1 * z2 + z1 * y2 + dy = -w1 * y2 + x1 * z2 + y1 * w2 - z1 * x2 + dz = -w1 * z2 - x1 * y2 + y1 * x2 + z1 * w2 + imaginary_norm = xp.sqrt(dx * dx + dy * dy + dz * dz) + imaginary_norm = xp.where(imaginary_norm < 0.0, 0.0, xp.where(imaginary_norm > 1.0, 1.0, imaginary_norm)) + return 2.0 * xp.asin(imaginary_norm) + + +def from_euler(euler): + """Quaternion (xyzw) from ``(..., 3)`` roll/pitch/yaw angles; trig keeps the input dtype.""" + xp = array_namespace(euler) + roll, pitch, yaw = (euler[..., i] for i in range(3)) + cy = xp.cos(yaw * 0.5) + sy = xp.sin(yaw * 0.5) + cp = xp.cos(pitch * 0.5) + sp = xp.sin(pitch * 0.5) + cr = xp.cos(roll * 0.5) + sr = xp.sin(roll * 0.5) + return xp.stack( + ( + sr * cp * cy - cr * sp * sy, + cr * sp * cy + sr * cp * sy, + cr * cp * sy - sr * sp * cy, + cr * cp * cy + sr * sp * cy, + ), + dim=-1, + ) + + +def rotate_inverse_components(quaternion, vector): + """Rotate vectors by the inverse of ``[x, y, z, w]`` quaternions, batched. + + ``quaternion`` is ``(..., 4)`` and ``vector`` is ``(..., 3)``; the rotated + components are returned as three tensors of the leading shape. + """ + xp = array_namespace(quaternion, vector) + x, y, z, w = (quaternion[..., i] for i in range(4)) + vx, vy, vz = (vector[..., i] for i in range(3)) + cx = y * vz - z * vy - w * vx + cy = z * vx - x * vz - w * vy + cz = x * vy - y * vx - w * vz + return ( + vx + 2.0 * (y * cz - z * cy), + vy + 2.0 * (z * cx - x * cz), + vz + 2.0 * (x * cy - y * cx), + ) diff --git a/motrix_env_core/src/motrix_env_core/mdp/noise.py b/motrix_env_core/src/motrix_env_core/mdp/noise.py deleted file mode 100644 index e1408c72..00000000 --- a/motrix_env_core/src/motrix_env_core/mdp/noise.py +++ /dev/null @@ -1,36 +0,0 @@ -# Copyright Motphys Technology Co., Ltd. 2025, 2026 -# SPDX-License-Identifier: Apache-2.0 - -"""Stateful per-environment observation noise for Numba manager terms.""" - -import numba -import numpy as np - -from motrix_env_core.numba.manager.rand import next_uniform - - -@numba.njit(inline="always") -def add_uniform_noise(out: np.ndarray, amplitude: np.float32, rng_state: np.ndarray) -> None: - """Consume one random value per output component and add scaled uniform noise in place.""" - if amplitude != 0.0: - for component in range(out.shape[0]): - out[component] += next_uniform(rng_state) * amplitude - - -@numba.njit(cache=True, parallel=True) -def add_uniform_noise_batch( - out: np.ndarray, - amplitude: float, - rng_states: np.ndarray, - env_ids: np.ndarray, -) -> None: - """Advance selected environment states while adding noise to a reset batch.""" - if amplitude != 0.0: - for row in numba.prange(out.shape[0]): - add_uniform_noise(out[row], amplitude, rng_states[env_ids[row]]) - - -__all__ = [ - "add_uniform_noise", - "add_uniform_noise_batch", -] diff --git a/motrix_env_core/src/motrix_env_core/mdp/observations.py b/motrix_env_core/src/motrix_env_core/mdp/observations.py index 1a6888d6..c7b41f5a 100644 --- a/motrix_env_core/src/motrix_env_core/mdp/observations.py +++ b/motrix_env_core/src/motrix_env_core/mdp/observations.py @@ -8,15 +8,13 @@ from typing import TYPE_CHECKING import numpy as np -from numba import literally +import torch from motrix_env_core.config import configclass from motrix_env_core.config.scene import RobotCfg -from motrix_env_core.mdp.noise import add_uniform_noise -from motrix_env_core.numba.manager.context import BuildContext, ManagerContext -from motrix_env_core.numba.manager.dispatch import dispatch -from motrix_env_core.numba.manager.observations import ObservationTermCfg, ObsTerm -from motrix_env_core.numba.math.quaternion import rotate_inverse +from motrix_env_core.manager.context import BuildContext, ManagerContext +from motrix_env_core.manager.dispatch import dispatch +from motrix_env_core.manager.observations import ObservationTermCfg, ObsTerm from motrix_env_core.sim import ( BodyJointPositionQuery, BodyJointVelocityQuery, @@ -45,20 +43,36 @@ def _body(ctx: BuildContext, name: str) -> BodyModel: return ctx.model.bodies[name] +def add_uniform_noise(ctx: ManagerContext, out: torch.Tensor, amplitude: float) -> None: + """Consume one random value per output component and add scaled uniform noise in place.""" + if amplitude != 0.0: + out += ctx.rand.next_uniform(out.shape[1]) * amplitude + + +def _rotate_inverse(quaternion: torch.Tensor, vector: torch.Tensor, out: torch.Tensor) -> None: + """Rotate ``(N, 3)`` vectors by the inverse of ``(N, 4)`` ``[x, y, z, w]`` quaternions.""" + x, y, z, w = quaternion.unbind(dim=1) + vx, vy, vz = vector.unbind(dim=1) + cx = y * vz - z * vy - w * vx + cy = z * vx - x * vz - w * vy + cz = x * vy - y * vx - w * vz + out[:, 0] = vx + 2.0 * (y * cz - z * cy) + out[:, 1] = vy + 2.0 * (z * cx - x * cz) + out[:, 2] = vz + 2.0 * (x * cy - y * cx) + + @dispatch def body_joint_vel_obs( - ctx: ManagerContext, out: np.ndarray, dof_vel: np.ndarray, scale: np.float32, noise_amplitude: np.float32 + ctx: ManagerContext, out: torch.Tensor, dof_vel: torch.Tensor, scale: float, noise_amplitude: float ) -> None: - out[:] = dof_vel - out *= scale - add_uniform_noise(out, noise_amplitude, ctx.rand.state) + out.copy_(dof_vel) + out.mul_(scale) + add_uniform_noise(ctx, out, noise_amplitude) @dispatch -def actions_obs(ctx: ManagerContext, out: np.ndarray, action_name: str) -> None: - action_name = literally(action_name) - action = ctx.actions[action_name] - out[:] = action.current() +def actions_obs(ctx: ManagerContext, out: torch.Tensor, action_name: str) -> None: + out.copy_(ctx.actions[action_name].current()) @configclass(kw_only=True) @@ -73,10 +87,8 @@ def __call__(self, ctx: BuildContext) -> ObsTerm: @dispatch -def command_obs(ctx: ManagerContext, out: np.ndarray, command_name: str) -> None: - command_name = literally(command_name) - command = ctx.commands[command_name].command - out[:] = command +def command_obs(ctx: ManagerContext, out: torch.Tensor, command_name: str) -> None: + out.copy_(ctx.commands[command_name].command) @configclass(kw_only=True) @@ -94,9 +106,10 @@ def __call__(self, ctx: BuildContext) -> ObsTerm: class BodyJointVelObsCfg(ObservationTermCfg): """Joint-velocity observation of one scene body in its own frame. - The term passes its query as an argument; the compiler registers the - query and the dispatch receives the query's lane view in that position. - Equal queries across terms fold into one physical read at the backend. + The term passes its query as an argument; the resolver registers the + query and the dispatch receives the query's batched view in that + position. Equal queries across terms fold into one physical read at the + backend. """ body: str = "robot" @@ -117,15 +130,15 @@ def __call__(self, ctx: BuildContext) -> ObsTerm: @dispatch def body_linear_velocity_obs( ctx: ManagerContext, - out: np.ndarray, - base_quat: np.ndarray, - linear_velocity: np.ndarray, - scale: np.float32, - noise_amplitude: np.float32, + out: torch.Tensor, + base_quat: torch.Tensor, + linear_velocity: torch.Tensor, + scale: float, + noise_amplitude: float, ) -> None: - rotate_inverse(base_quat, linear_velocity, out) - out *= scale - add_uniform_noise(out, noise_amplitude, ctx.rand.state) + _rotate_inverse(base_quat, linear_velocity, out) + out.mul_(scale) + add_uniform_noise(ctx, out, noise_amplitude) @configclass(kw_only=True) @@ -151,15 +164,15 @@ def __call__(self, ctx: BuildContext) -> ObsTerm: @dispatch def body_angular_velocity_obs( ctx: ManagerContext, - out: np.ndarray, - base_quat: np.ndarray, - angular_velocity: np.ndarray, - scale: np.float32, - noise_amplitude: np.float32, + out: torch.Tensor, + base_quat: torch.Tensor, + angular_velocity: torch.Tensor, + scale: float, + noise_amplitude: float, ) -> None: - rotate_inverse(base_quat, angular_velocity, out) - out *= scale - add_uniform_noise(out, noise_amplitude, ctx.rand.state) + _rotate_inverse(base_quat, angular_velocity, out) + out.mul_(scale) + add_uniform_noise(ctx, out, noise_amplitude) @configclass(kw_only=True) @@ -184,10 +197,12 @@ def __call__(self, ctx: BuildContext) -> ObsTerm: @dispatch def body_projected_gravity_obs( - ctx: ManagerContext, out: np.ndarray, base_quat: np.ndarray, noise_amplitude: np.float32 + ctx: ManagerContext, out: torch.Tensor, base_quat: torch.Tensor, noise_amplitude: float ) -> None: - rotate_inverse(base_quat, (0.0, 0.0, -1.0), out) - add_uniform_noise(out, noise_amplitude, ctx.rand.state) + gravity = out.new_zeros(out.shape) + gravity[:, 2] = -1.0 + _rotate_inverse(base_quat, gravity, out) + add_uniform_noise(ctx, out, noise_amplitude) @configclass(kw_only=True) @@ -210,16 +225,15 @@ def __call__(self, ctx: BuildContext) -> ObsTerm: @dispatch def body_joint_pos_rel_obs( ctx: ManagerContext, - out: np.ndarray, - joint_pos: np.ndarray, - defaults: np.ndarray, - scale: np.float32, - noise_amplitude: np.float32, + out: torch.Tensor, + joint_pos: torch.Tensor, + defaults: torch.Tensor, + scale: float, + noise_amplitude: float, ) -> None: - out[:] = joint_pos - out -= defaults - out *= scale - add_uniform_noise(out, noise_amplitude, ctx.rand.state) + out.copy_(joint_pos - torch.as_tensor(np.asarray(defaults), dtype=out.dtype)) + out.mul_(scale) + add_uniform_noise(ctx, out, noise_amplitude) @configclass(kw_only=True) diff --git a/motrix_env_core/src/motrix_env_core/mdp/rewards.py b/motrix_env_core/src/motrix_env_core/mdp/rewards.py index 5504a4fa..b633a062 100644 --- a/motrix_env_core/src/motrix_env_core/mdp/rewards.py +++ b/motrix_env_core/src/motrix_env_core/mdp/rewards.py @@ -3,16 +3,14 @@ """Reusable reward terms for manager-based environments.""" -import math - import numpy as np -from numba import literally +import torch from motrix_env_core.config import configclass from motrix_env_core.manager import RewardTerm, RewardTermCfg -from motrix_env_core.numba.manager.context import BuildContext, ManagerContext -from motrix_env_core.numba.manager.dispatch import dispatch -from motrix_env_core.numba.math.quaternion import rotate_inverse_components +from motrix_env_core.manager.context import BuildContext, ManagerContext +from motrix_env_core.manager.dispatch import dispatch +from motrix_env_core.math.quaternion_array import rotate_inverse_components from motrix_env_core.sim import ( LinkAngularVelocityQuery, LinkLinearVelocityQuery, @@ -21,8 +19,8 @@ @dispatch -def alive_reward(_ctx: ManagerContext) -> float: - return 1.0 +def alive_reward(ctx: ManagerContext) -> torch.Tensor: + return torch.ones_like(ctx.env_id, dtype=torch.float32) @configclass(kw_only=True) @@ -30,16 +28,14 @@ class AliveRewardCfg(RewardTermCfg): """Constant alive bonus; zero weight disables it.""" def __call__(self, ctx) -> RewardTerm: - del ctx return RewardTerm(alive_reward) @dispatch -def action_rate_reward(ctx: ManagerContext, action_name: str) -> float: - action_name = literally(action_name) +def action_rate_reward(ctx: ManagerContext, action_name: str) -> torch.Tensor: action = ctx.actions[action_name] delta = action.current() - action.previous() - return float(np.dot(delta, delta)) + return (delta * delta).sum(dim=1) @configclass(kw_only=True) @@ -49,23 +45,21 @@ class ActionRateRewardCfg(RewardTermCfg): action_name: str = "joint_position" def __call__(self, ctx) -> RewardTerm: - del ctx return RewardTerm(action_rate_reward, self.action_name) @dispatch def tracking_lin_vel_xy_reward( ctx: ManagerContext, - sigma: np.float32, - base_quat: np.ndarray, - base_lin_vel: np.ndarray, + sigma: float, + base_quat: torch.Tensor, + base_lin_vel: torch.Tensor, command_name: str, -) -> float: - command_name = literally(command_name) +) -> torch.Tensor: command = ctx.commands[command_name].command vx, vy, _ = rotate_inverse_components(base_quat, base_lin_vel) - error = (command[0] - vx) * (command[0] - vx) + (command[1] - vy) * (command[1] - vy) - return math.exp(-error / sigma) + error = (command[:, 0] - vx) * (command[:, 0] - vx) + (command[:, 1] - vy) * (command[:, 1] - vy) + return torch.exp(-error / sigma) @configclass(kw_only=True) @@ -90,16 +84,15 @@ def __call__(self, ctx: BuildContext) -> RewardTerm: @dispatch def tracking_ang_vel_z_reward( ctx: ManagerContext, - sigma: np.float32, - base_quat: np.ndarray, - base_ang_vel: np.ndarray, + sigma: float, + base_quat: torch.Tensor, + base_ang_vel: torch.Tensor, command_name: str, -) -> float: - command_name = literally(command_name) +) -> torch.Tensor: command = ctx.commands[command_name].command _, _, wz = rotate_inverse_components(base_quat, base_ang_vel) - error = (command[2] - wz) * (command[2] - wz) - return math.exp(-error / sigma) + error = (command[:, 2] - wz) * (command[:, 2] - wz) + return torch.exp(-error / sigma) @configclass(kw_only=True) diff --git a/motrix_env_core/src/motrix_env_core/mdp/state.py b/motrix_env_core/src/motrix_env_core/mdp/state.py index abed0078..3cbf892a 100644 --- a/motrix_env_core/src/motrix_env_core/mdp/state.py +++ b/motrix_env_core/src/motrix_env_core/mdp/state.py @@ -7,7 +7,7 @@ from typing import TYPE_CHECKING -from motrix_env_core.numba.manager.rand import RandValue, initialize_rand_states +from motrix_env_core.manager.rand import RandValue, initialize_rand_states if TYPE_CHECKING: from motrix_env_core.manager import ManagerEnv diff --git a/motrix_env_core/src/motrix_env_core/mdp/terminations.py b/motrix_env_core/src/motrix_env_core/mdp/terminations.py index e03caaf4..365ddd70 100644 --- a/motrix_env_core/src/motrix_env_core/mdp/terminations.py +++ b/motrix_env_core/src/motrix_env_core/mdp/terminations.py @@ -3,35 +3,30 @@ """Reusable termination terms for manager-based environments.""" -import math - import numpy as np +import torch from motrix_env_core.config import configclass from motrix_env_core.manager import ManagerContext, TerminationTerm, TerminationTermCfg -from motrix_env_core.numba.manager.context import BuildContext -from motrix_env_core.numba.manager.dispatch import dispatch +from motrix_env_core.manager.context import BuildContext +from motrix_env_core.manager.dispatch import dispatch from motrix_env_core.sim import BodyJointVelocityQuery, GeomPairCollidingQuery @dispatch -def colliding_termination(ctx: ManagerContext, colliding: np.ndarray) -> bool: - return bool(colliding.any()) +def colliding_termination(ctx: ManagerContext, colliding: torch.Tensor) -> torch.Tensor: + return colliding.any(dim=1) @dispatch -def bad_dof_velocity_termination(ctx: ManagerContext, dof_vel: np.ndarray, threshold: np.float32) -> bool: - error = 0.0 - finite = True - for joint_id in range(dof_vel.shape[0]): - velocity = dof_vel[joint_id] - if math.isfinite(velocity): - error = max(error, abs(velocity)) - else: - finite = False - error = math.inf - ctx.metrics["dof_vel_abs_max"][0] = error - return (not finite) or error > threshold +def bad_dof_velocity_termination(ctx: ManagerContext, dof_vel: torch.Tensor, threshold: float) -> torch.Tensor: + # Non-finite components count as +inf, so the max over the masked + # absolute velocities terminates lanes with non-finite joint speeds + # (inf > threshold is always true). + masked = torch.where(torch.isfinite(dof_vel), dof_vel.abs(), torch.inf) + error = masked.max(dim=1).values + ctx.metrics["dof_vel_abs_max"][:, 0] = error + return error > threshold @configclass(kw_only=True) @@ -39,7 +34,7 @@ class CollidingTerminationCfg(TerminationTermCfg): """Terminate when any of ``termination_geoms`` contacts ``ground_geom``. The collision query is passed as an argument; the dispatch receives the - query's lane view (one bool per declared geom pair). + query's batched view (one bool per declared geom pair). """ termination_geoms: tuple[str, ...] = () diff --git a/motrix_env_core/src/motrix_env_core/mdp/terrain.py b/motrix_env_core/src/motrix_env_core/mdp/terrain.py index 1b77987c..c395137c 100644 --- a/motrix_env_core/src/motrix_env_core/mdp/terrain.py +++ b/motrix_env_core/src/motrix_env_core/mdp/terrain.py @@ -1,28 +1,24 @@ # Copyright Motphys Technology Co., Ltd. 2025, 2026 # SPDX-License-Identifier: Apache-2.0 -"""In-kernel terrain-height lookup over a static height-field grid.""" +"""Static terrain-height grid shared by manager reset and reward terms.""" import numpy as np -from numba import njit +import torch from motrix_env_core.manager import SharedArray, kernel_data @kernel_data class HeightFieldGrid: - """Static terrain grid for in-kernel bilinear height lookup. + """Static terrain grid for batched bilinear height lookup. ``enabled=False`` degenerates to the flat-ground ``constant`` height. - - 采样高度计算公式(见 :func:`heightfield_lookup`):: - - fx = (x - origin[0]) / spacing[0] - fy = (y - origin[1]) / spacing[1] # clamp 到 [0, n - 1 - eps] - col, row = int(fx), int(fy); tx, ty = fx - col, fy - row - top = heights[row, col] * (1 - tx) + heights[row, col + 1] * tx - bottom = heights[row + 1, col] * (1 - tx) + heights[row + 1, col + 1] * tx - h(x, y) = z0[0] + top * (1 - ty) + bottom * ty + Torch-side lookups (for example + ``motrix_envs.locomotion.humanoid.walk_manager_mdp.terrain.heightfield_lookup_batch``) + mirror the engine's hfield sample convention in float64: + bilinear interpolation of the four surrounding samples after mapping + world ``xy`` to fractional grid coordinates. Attributes: heights: ``SharedArray`` of shape ``(nrow, ncol)`` holding the @@ -35,8 +31,8 @@ class HeightFieldGrid: z0: 长度为 1 的 ``SharedArray``,叠加在插值结果之上的基准高度偏移。 constant: 平地退化模式下的常数高度;仅在 ``enabled=False`` 时生效, 对应平地 ground geom 的世界 z 坐标。 - enabled: 是否启用网格插值;``False`` 时 ``heightfield_lookup`` - 直接返回 ``constant``,忽略所有网格字段。 + enabled: 是否启用网格插值;``False`` 时高度查找直接返回 + ``constant``,忽略所有网格字段。 """ heights: SharedArray @@ -47,31 +43,35 @@ class HeightFieldGrid: enabled: bool -@njit(inline="always") -def heightfield_lookup(grid: HeightFieldGrid, x: float, y: float) -> float: - """Bilinear grid lookup matching the engine's hfield sample convention.""" - if not grid.enabled: - return grid.constant - nrow = grid.heights.shape[0] - ncol = grid.heights.shape[1] - fx = (x - grid.origin[0]) / grid.spacing[0] - fy = (y - grid.origin[1]) / grid.spacing[1] - fx = min(max(fx, 0.0), ncol - 1.0 - 1e-6) - fy = min(max(fy, 0.0), nrow - 1.0 - 1e-6) - col = int(fx) - row = int(fy) +__all__ = [ + "heightfield_lookup", + "HeightFieldGrid", +] + + +def heightfield_lookup(grid, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: + """Batched torch bilinear height lookup over a :class:`HeightFieldGrid`. + + ``x``/``y`` are ``(N,)`` world coordinates. The interpolation runs in + float64, matching the engine's hfield sample convention. + """ + if not bool(grid.enabled): + return torch.full_like(x.double(), float(grid.constant)) + heights = torch.as_tensor(np.asarray(grid.heights)).to(torch.float64) + origin = torch.as_tensor(np.asarray(grid.origin)).to(torch.float64) + spacing = torch.as_tensor(np.asarray(grid.spacing)).to(torch.float64) + z0 = torch.as_tensor(np.asarray(grid.z0)).to(torch.float64) + nrow, ncol = heights.shape + fx = ((x.to(torch.float64) - origin[0]) / spacing[0]).clamp(0.0, ncol - 1.0 - 1e-6) + fy = ((y.to(torch.float64) - origin[1]) / spacing[1]).clamp(0.0, nrow - 1.0 - 1e-6) + col = fx.to(torch.long) + row = fy.to(torch.long) tx = fx - col ty = fy - row - h00 = grid.heights[row, col] - h01 = grid.heights[row, col + 1] - h10 = grid.heights[row + 1, col] - h11 = grid.heights[row + 1, col + 1] + h00 = heights[row, col] + h01 = heights[row, col + 1] + h10 = heights[row + 1, col] + h11 = heights[row + 1, col + 1] top = h00 * (1.0 - tx) + h01 * tx bottom = h10 * (1.0 - tx) + h11 * tx - return grid.z0[0] + top * (1.0 - ty) + bottom * ty - - -__all__ = [ - "HeightFieldGrid", - "heightfield_lookup", -] + return z0[0] + top * (1.0 - ty) + bottom * ty diff --git a/motrix_env_core/src/motrix_env_core/numba/__init__.py b/motrix_env_core/src/motrix_env_core/numba/__init__.py deleted file mode 100644 index a4bdf551..00000000 --- a/motrix_env_core/src/motrix_env_core/numba/__init__.py +++ /dev/null @@ -1,18 +0,0 @@ -# Copyright Motphys Technology Co., Ltd. 2025, 2026 -# SPDX-License-Identifier: Apache-2.0 - -from motrix_env_core.numba.kernel_data import ( - SharedArray, - TreeClassDef, - flatten_kernel_data, - kernel_data, - unflatten_kernel_data, -) - -__all__ = [ - "SharedArray", - "TreeClassDef", - "flatten_kernel_data", - "kernel_data", - "unflatten_kernel_data", -] diff --git a/motrix_env_core/src/motrix_env_core/numba/kernel_data/methods.py b/motrix_env_core/src/motrix_env_core/numba/kernel_data/methods.py deleted file mode 100644 index 6dcc4b7b..00000000 --- a/motrix_env_core/src/motrix_env_core/numba/kernel_data/methods.py +++ /dev/null @@ -1,51 +0,0 @@ -# Copyright Motphys Technology Co., Ltd. 2025, 2026 -# SPDX-License-Identifier: Apache-2.0 - -from __future__ import annotations - -import inspect -from collections.abc import Callable -from typing import Any - -from numba import types -from numba.extending import overload_method - -from motrix_env_core.numba.manager.dispatch import _DISPATCH_MARKER - -_METHOD_OVERLOADS: dict[str, dict[type, Callable[..., Any]]] = {} - - -def dispatch_methods(logical_type: type) -> dict[str, Callable[..., Any]]: - """Resolve marked instance methods with normal Python MRO shadowing.""" - methods: dict[str, Callable[..., Any]] = {} - seen: set[str] = set() - for base in logical_type.__mro__: - for name, value in vars(base).items(): - if name in seen: - continue - seen.add(name) - if inspect.isfunction(value) and getattr(value, _DISPATCH_MARKER, False): - methods[name] = value - return methods - - -def register_proxy_method(proxy: type, name: str, method: Callable[..., Any]) -> None: - """Register one marked KernelData method for a lowered proxy type.""" - methods = _METHOD_OVERLOADS.get(name) - if methods is not None: - methods[proxy] = method - return - methods = {proxy: method} - _METHOD_OVERLOADS[name] = methods - - def getter(self: Any, *args: Any, **kwargs: Any) -> Any: - return methods.get(self.instance_class) - - # Register one getter per name: competing attribute templates for the same - # tuple method can hide later proxies. Homogeneous records (including - # single-field records) use NamedUniTuple. - for tuple_type in (types.NamedTuple, types.NamedUniTuple): - overload_method(tuple_type, name, strict=False, inline="always")(getter) - - -__all__ = ["dispatch_methods", "register_proxy_method"] diff --git a/motrix_env_core/src/motrix_env_core/numba/manager/__init__.py b/motrix_env_core/src/motrix_env_core/numba/manager/__init__.py deleted file mode 100644 index 942236b9..00000000 --- a/motrix_env_core/src/motrix_env_core/numba/manager/__init__.py +++ /dev/null @@ -1,4 +0,0 @@ -# Copyright Motphys Technology Co., Ltd. 2025, 2026 -# SPDX-License-Identifier: Apache-2.0 - -"""Internal Numba manager implementation.""" diff --git a/motrix_env_core/src/motrix_env_core/numba/manager/compiler/__init__.py b/motrix_env_core/src/motrix_env_core/numba/manager/compiler/__init__.py deleted file mode 100644 index 73934bc8..00000000 --- a/motrix_env_core/src/motrix_env_core/numba/manager/compiler/__init__.py +++ /dev/null @@ -1,19 +0,0 @@ -# Copyright Motphys Technology Co., Ltd. 2025, 2026 -# SPDX-License-Identifier: Apache-2.0 - -from typing import TYPE_CHECKING, Any - -from motrix_env_core.numba.manager.compiler.plan import ManagerLayout - -if TYPE_CHECKING: - from motrix_env_core.numba.manager.compiler.compiler import NumbaKernelCompiler - -__all__ = ["ManagerLayout", "NumbaKernelCompiler"] - - -def __getattr__(name: str) -> Any: - if name == "NumbaKernelCompiler": - from motrix_env_core.numba.manager.compiler.compiler import NumbaKernelCompiler - - return NumbaKernelCompiler - raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/motrix_env_core/src/motrix_env_core/numba/manager/compiler/cache.py b/motrix_env_core/src/motrix_env_core/numba/manager/compiler/cache.py deleted file mode 100644 index 30dc6a26..00000000 --- a/motrix_env_core/src/motrix_env_core/numba/manager/compiler/cache.py +++ /dev/null @@ -1,106 +0,0 @@ -# Copyright Motphys Technology Co., Ltd. 2025, 2026 -# SPDX-License-Identifier: Apache-2.0 - -"""In-process and on-disk caches for manager fused kernels and terms.""" - -import logging -import os -import tempfile -from collections.abc import Callable -from pathlib import Path -from typing import Any - -logger = logging.getLogger(__name__) - -_KERNEL_CACHE: dict[str, Any] = {} -_TERM_CACHE: dict[Callable[..., Any], Any] = {} - - -def invalidate_term_cache() -> None: - for dispatcher in _TERM_CACHE.values(): - try: - dispatcher._cache.flush() - except (AttributeError, OSError): - logger.debug("Unable to remove invalid manager term cache", exc_info=True) - _TERM_CACHE.clear() - - -def generated_cache_dir() -> str: - cache_dir = os.environ.get("NUMBA_CACHE_DIR") - if not cache_dir: - cache_dir = os.path.join(os.path.expanduser("~"), ".cache", "motrixlab", "numba-manager") - return cache_dir - - -def generated_source_path(plan_key: str) -> str: - return str(Path(generated_cache_dir()) / "generated" / f"{plan_key}.py") - - -def materialize_source(source: str, plan_key: str) -> str: - cache_dir = generated_cache_dir() - path = Path(cache_dir) / "generated" / f"{plan_key}.py" - path.parent.mkdir(parents=True, exist_ok=True) - if not path.exists() or path.read_text(encoding="utf-8") != source: - temporary_name: str | None = None - try: - with tempfile.NamedTemporaryFile( - mode="w", encoding="utf-8", dir=path.parent, prefix=f".{path.name}.", suffix=".tmp", delete=False - ) as temporary: - temporary.write(source) - temporary_name = temporary.name - os.replace(temporary_name, path) - finally: - if temporary_name is not None: - Path(temporary_name).unlink(missing_ok=True) - return str(path) - - -def generated_cache_paths(plan_key: str) -> tuple[Path, ...]: - """Return compiled-kernel cache files for one plan key. - - Numba stores the cached specializations of the generated module next to - its cache locator: under ``generated_/`` when ``NUMBA_CACHE_DIR`` - is set, and under ``generated/__pycache__/`` (beside the generated - source) when it is not. Both layouts are covered so probe and - invalidation agree with where the files actually live. - """ - cache_dir = Path(generated_cache_dir()) - paths: list[Path] = [] - for pattern in (f"generated_*/*{plan_key}*", f"generated/__pycache__/*{plan_key}*"): - try: - paths.extend(cache_dir.glob(pattern)) - except OSError: - logger.debug("Unable to scan generated kernel cache: %s", cache_dir, exc_info=True) - return tuple(paths) - - -def has_generated_kernel_cache(plan_key: str) -> bool: - """Return whether compiled-kernel cache files exist for one plan key.""" - return any(path.suffix in {".nbi", ".nbc"} for path in generated_cache_paths(plan_key)) - - -def invalidate_generated_cache(plan_key: str) -> None: - cache_dir = generated_cache_dir() - source_path = Path(cache_dir) / "generated" / f"{plan_key}.py" - try: - source_path.unlink(missing_ok=True) - except OSError: - logger.debug("Unable to remove invalid generated source: %s", source_path, exc_info=True) - for cache_path in generated_cache_paths(plan_key): - try: - cache_path.unlink(missing_ok=True) - except OSError: - logger.debug("Unable to remove invalid generated cache: %s", cache_path, exc_info=True) - - -__all__ = [ - "_KERNEL_CACHE", - "_TERM_CACHE", - "generated_cache_dir", - "generated_cache_paths", - "generated_source_path", - "has_generated_kernel_cache", - "invalidate_generated_cache", - "invalidate_term_cache", - "materialize_source", -] diff --git a/motrix_env_core/src/motrix_env_core/numba/manager/compiler/codegen.py b/motrix_env_core/src/motrix_env_core/numba/manager/compiler/codegen.py deleted file mode 100644 index aba338d6..00000000 --- a/motrix_env_core/src/motrix_env_core/numba/manager/compiler/codegen.py +++ /dev/null @@ -1,202 +0,0 @@ -# Copyright Motphys Technology Co., Ltd. 2025, 2026 -# SPDX-License-Identifier: Apache-2.0 - -from dataclasses import dataclass - -from motrix_env_core.numba.manager.compiler.plan import ObservationGroupLayout -from motrix_env_core.numba.manager.compiler.program import PreparedInvocation, ResolvedManagerContext, ResolvedSimReset - - -@dataclass(frozen=True) -class KernelSourceGenerator: - """Generate schema-specialized manager kernels for one compilation plan. - - Each kernel is emitted as its own module source so every kernel gets an - independent plan key and on-disk cache: editing reward terms must not - recompile the observation or reset kernels. - """ - - flat_input_count: int - - def generate( - self, - observations: dict[str, tuple[PreparedInvocation, ...]], - observation_layout: dict[str, ObservationGroupLayout], - rewards: tuple[PreparedInvocation, ...], - terminations: tuple[PreparedInvocation, ...], - context: ResolvedManagerContext, - command_updates: tuple[PreparedInvocation, ...], - command_advances: tuple[PreparedInvocation, ...], - command_resets: tuple[PreparedInvocation, ...], - reset: tuple[ResolvedSimReset, ...], - ) -> tuple[str, str, str]: - evaluate_lane_body, observe_lane_body = self._generate_lane_bodies( - observations, - observation_layout, - rewards, - terminations, - command_updates, - command_advances, - ) - evaluate_source = self._evaluate_module(evaluate_lane_body, context) - observe_source = self._observe_module(observe_lane_body, context) - reset_source = self._reset_module(context, command_resets, reset) - return evaluate_source, observe_source, reset_source - - def _evaluate_module(self, evaluate_lane_body: list[str], context: ResolvedManagerContext) -> str: - lines = [ - "def generated_evaluate_kernel(inputs, reward_weights, buffers, outputs):", - ] - lines.extend(f" input_{index} = inputs[{index}]" for index in range(self.flat_input_count)) - lines.extend( - [ - " def _evaluate_lane(", - " ctx, reward_weights, reward_terms, weighted_reward_terms, termination_masks,", - " reward, terminated, env_id,", - " ):", - ] - ) - lines.extend(f" {line}" for line in evaluate_lane_body) - lines.extend( - [ - " reward_terms = buffers[0]", - " weighted_reward_terms = buffers[1]", - " termination_masks = buffers[2]", - " reward = outputs.reward", - " terminated = outputs.terminated", - " for env_id in numba.prange(reward.shape[0]):", - ] - ) - lines.append(f" ctx = {context.expression}") - lines.extend( - [ - " _evaluate_lane(", - " ctx,", - " reward_weights,", - " reward_terms,", - " weighted_reward_terms,", - " termination_masks,", - " reward,", - " terminated,", - " env_id,", - " )", - ] - ) - return "\n".join(lines) + "\n" - - def _observe_module(self, observe_lane_body: list[str], context: ResolvedManagerContext) -> str: - lines = [ - "def generated_observe_kernel(inputs, outputs):", - ] - lines.extend(f" input_{index} = inputs[{index}]" for index in range(self.flat_input_count)) - lines.extend( - [ - " policy_obs = outputs.policy_obs", - " value_obs = outputs.value_obs", - " for env_id in numba.prange(policy_obs.shape[0]):", - f" ctx = {context.expression}", - ] - ) - lines.extend(f" {line}" for line in observe_lane_body) - return "\n".join(lines) + "\n" - - def _reset_module( - self, - context: ResolvedManagerContext, - command_resets: tuple[PreparedInvocation, ...], - reset: tuple[ResolvedSimReset, ...], - ) -> str: - lines = [ - "def generated_reset_kernel(inputs, env_ids, reset_buffers):", - ] - lines.extend(f" input_{index} = inputs[{index}]" for index in range(self.flat_input_count)) - lines.extend([" for row in numba.prange(env_ids.shape[0]):", " env_id = env_ids[row]"]) - lines.append(f" ctx = {context.expression}") - for index, command_reset in enumerate(command_resets): - lines.append(f" {self._term_call(command_reset, index, 'ctx')}") - for index, reset_term in enumerate(reset): - output_args = tuple( - f"reset_buffers[{reset_term.output_offset + field_index}][row]" - for field_index in range(reset_term.output_count) - ) - lines.append(f" sim_writes_{index} = {reset_term.sim_writes_type}({', '.join(output_args)})") - call = self._term_call( - reset_term.invocation, - index, - "ctx", - f"sim_writes_{index}", - ) - lines.append(f" {call}") - return "\n".join(lines) + "\n" - - @classmethod - def _generate_lane_bodies( - cls, - observations: dict[str, tuple[PreparedInvocation, ...]], - observation_layout: dict[str, ObservationGroupLayout], - rewards: tuple[PreparedInvocation, ...], - terminations: tuple[PreparedInvocation, ...], - command_updates: tuple[PreparedInvocation, ...], - command_advances: tuple[PreparedInvocation, ...], - ) -> tuple[list[str], list[str]]: - """Emit straight-line local helpers that Numba inlines during the kernel frontend pass. - - Keeping these helpers local avoids an extra dispatcher compilation. Their straight-line control flow also lets - the enclosing ``prange`` remain parallel, unlike heterogeneous ``literal_unroll`` loops. - """ - evaluate_lines = [] - for index, update in enumerate(command_updates): - evaluate_lines.append(cls._term_call(update, index, "ctx")) - evaluate_lines.append("total_reward = 0.0") - for index, term in enumerate(rewards): - evaluate_lines.extend( - [ - f"reward_value_{index} = {cls._term_call(term, index, 'ctx')}", - f"reward_terms[env_id, {index}] = reward_value_{index}", - f"weighted_reward_{index} = reward_value_{index} * reward_weights[{index}] * ctx.dt", - f"weighted_reward_terms[env_id, {index}] = weighted_reward_{index}", - f"total_reward += weighted_reward_{index}", - ] - ) - evaluate_lines.append("reward[env_id] = total_reward") - evaluate_lines.append("is_terminated = False") - for index, term in enumerate(terminations): - evaluate_lines.extend( - [ - f"termination_value_{index} = {cls._term_call(term, index, 'ctx')}", - f"termination_masks[env_id, {index}] = termination_value_{index}", - f"is_terminated = is_terminated or termination_value_{index}", - ] - ) - evaluate_lines.append("terminated[env_id] = is_terminated") - for index, advance in enumerate(command_advances): - evaluate_lines.append(cls._term_call(advance, index, "ctx")) - - observe_lines = [] - observation_index = 0 - for group_name, group_terms in observations.items(): - if group_name not in {"policy", "value"}: - raise ValueError(f"Unsupported observation group {group_name!r}.") - output_name = "policy_obs" if group_name == "policy" else "value_obs" - for term, term_layout in zip(group_terms, observation_layout[group_name].terms, strict=True): - output_slice = ( - f"{output_name}[env_id, {term_layout.output_slice.start}:{term_layout.output_slice.stop}]" - ) - observe_lines.append(cls._term_call(term, observation_index, "ctx", output_slice)) - observation_index += 1 - if not observe_lines: - observe_lines.append("pass") - return evaluate_lines, observe_lines - - @staticmethod - def _term_call(invocation: PreparedInvocation, term_index: int, *arguments: str) -> str: - """Emit one positional call for a dispatch term or legacy receiver term.""" - symbol = f"term_{invocation.kind}_{term_index}" - if invocation.receiver_key is not None: - store_name = "commands" - arguments = (f"ctx.{store_name}[{invocation.receiver_key!r}]", *arguments) - arguments = (*arguments, *invocation.args_expressions) - return f"{symbol}({', '.join(arguments)})" - - -__all__ = ["KernelSourceGenerator"] diff --git a/motrix_env_core/src/motrix_env_core/numba/manager/compiler/fingerprint.py b/motrix_env_core/src/motrix_env_core/numba/manager/compiler/fingerprint.py deleted file mode 100644 index cbf5222b..00000000 --- a/motrix_env_core/src/motrix_env_core/numba/manager/compiler/fingerprint.py +++ /dev/null @@ -1,18 +0,0 @@ -# Copyright Motphys Technology Co., Ltd. 2025, 2026 -# SPDX-License-Identifier: Apache-2.0 - -"""Plan-key helpers for the manager kernel compiler.""" - -import hashlib -from typing import Any - - -def plan_key(parts: list[Any]) -> str: - return hashlib.sha256(repr(parts).encode()).hexdigest() - - -def type_name(value_type: type[Any]) -> str: - return f"{value_type.__module__}.{value_type.__qualname__}" - - -__all__ = ["plan_key", "type_name"] diff --git a/motrix_env_core/src/motrix_env_core/numba/manager/compiler/plan.py b/motrix_env_core/src/motrix_env_core/numba/manager/compiler/plan.py deleted file mode 100644 index d6a1978a..00000000 --- a/motrix_env_core/src/motrix_env_core/numba/manager/compiler/plan.py +++ /dev/null @@ -1,112 +0,0 @@ -# Copyright Motphys Technology Co., Ltd. 2025, 2026 -# SPDX-License-Identifier: Apache-2.0 - -from dataclasses import dataclass - -import numpy as np - -from motrix_env_core.numba.kernel_data import KernelDataScope - - -@dataclass(frozen=True) -class InputSlotLayout: - index: int - source: str - scope: KernelDataScope - - -@dataclass(frozen=True) -class SimInputLayout: - slot: int - key: str - shape: tuple[int, ...] - consumers: tuple[str, ...] - - -@dataclass(frozen=True) -class ActionTermLayout: - name: str - input_slice: slice - actuator_names: tuple[str, ...] | None - - -@dataclass(frozen=True) -class ObservationTermLayout: - name: str - output_slice: slice - - -@dataclass(frozen=True) -class ObservationGroupLayout: - name: str - terms: tuple[ObservationTermLayout, ...] - size: int - - -@dataclass(frozen=True) -class ScalarTermLayout: - name: str - index: int - - -@dataclass(frozen=True) -class MetricFieldLayout: - name: str - index: int - dtype: np.dtype - - -@dataclass(frozen=True) -class ManagerLayout: - inputs: tuple[InputSlotLayout, ...] - sim_inputs: tuple[SimInputLayout, ...] - actions: tuple[ActionTermLayout, ...] - observations: dict[str, ObservationGroupLayout] - rewards: tuple[ScalarTermLayout, ...] - terminations: tuple[ScalarTermLayout, ...] - metrics: tuple[MetricFieldLayout, ...] - plan_keys: tuple[str, str, str] - generated_filenames: tuple[str, str, str] - - def dump(self) -> str: - lines = [ - f"plan_keys (evaluate, observe, reset): {', '.join(key[:12] for key in self.plan_keys)}", - f"generated_filenames: {', '.join(self.generated_filenames)}", - "inputs:", - ] - lines.extend(f" [{slot.index}] scope={slot.scope.value} source={slot.source}" for slot in self.inputs) - lines.append("sim_inputs:") - lines.extend( - f" [{sim_input.slot}] {sim_input.key}: shape={sim_input.shape} consumers={sim_input.consumers}" - for sim_input in self.sim_inputs - ) - lines.append("actions:") - lines.extend( - f" {term.name}: input=[{term.input_slice.start}:{term.input_slice.stop}] actuators={term.actuator_names}" - for term in self.actions - ) - lines.append("observations:") - lines.extend( - f" {group.name}.{term.name}: [{term.output_slice.start}:{term.output_slice.stop}]" - for group in self.observations.values() - for term in group.terms - ) - lines.append("rewards:") - lines.extend(f" [{term.index}] {term.name}" for term in self.rewards) - lines.append("terminations:") - lines.extend(f" [{term.index}] {term.name}" for term in self.terminations) - lines.append("metrics:") - lines.extend(f" [{metric.index}] dtype={metric.dtype.name} {metric.name}" for metric in self.metrics) - return "\n".join(lines) - - -__all__ = [ - "ActionTermLayout", - "InputSlotLayout", - "ManagerLayout", - "MetricFieldLayout", - "ObservationGroupLayout", - "ObservationTermLayout", - "ScalarTermLayout", - "SimInputLayout", -] diff --git a/motrix_env_core/src/motrix_env_core/numba/manager/compiler/program.py b/motrix_env_core/src/motrix_env_core/numba/manager/compiler/program.py deleted file mode 100644 index b6ebe726..00000000 --- a/motrix_env_core/src/motrix_env_core/numba/manager/compiler/program.py +++ /dev/null @@ -1,51 +0,0 @@ -# Copyright Motphys Technology Co., Ltd. 2025, 2026 -# SPDX-License-Identifier: Apache-2.0 - -from dataclasses import dataclass -from typing import Any - -from motrix_env_core.numba.manager.env import CompiledManagerProgram, KernelInputSource - - -@dataclass(frozen=True) -class PreparedInvocation: - """Runtime metadata required to generate one compiled manager term.""" - - context: str - kind: str - dispatcher: Any - receiver_key: str | None - args_expressions: tuple[str, ...] = () - args_values: tuple[Any, ...] = () - args_prepared_indices: tuple[int | None, ...] = () - output_size: int = 0 - - -@dataclass(frozen=True) -class ResolvedManagerContext: - prepared_index: int - expression: str - - -@dataclass(frozen=True) -class ResolvedSimReset: - invocation: PreparedInvocation - sim_writes_type: str - output_offset: int - output_count: int - - -@dataclass(frozen=True) -class _CompiledManagerProgram(CompiledManagerProgram): - invocations: tuple[PreparedInvocation, ...] - prepared_terms: tuple[KernelInputSource, ...] - input_offsets: tuple[int, ...] - context: ResolvedManagerContext - - -__all__ = [ - "PreparedInvocation", - "ResolvedManagerContext", - "ResolvedSimReset", - "_CompiledManagerProgram", -] diff --git a/motrix_env_core/src/motrix_env_core/numba/math/__init__.py b/motrix_env_core/src/motrix_env_core/numba/math/__init__.py deleted file mode 100644 index ac1f5a42..00000000 --- a/motrix_env_core/src/motrix_env_core/numba/math/__init__.py +++ /dev/null @@ -1,4 +0,0 @@ -# Copyright Motphys Technology Co., Ltd. 2025, 2026 -# SPDX-License-Identifier: Apache-2.0 - -"""Numba-compatible math helpers.""" diff --git a/motrix_env_core/src/motrix_env_core/numba/math/quaternion.py b/motrix_env_core/src/motrix_env_core/numba/math/quaternion.py deleted file mode 100644 index ddbfbb72..00000000 --- a/motrix_env_core/src/motrix_env_core/numba/math/quaternion.py +++ /dev/null @@ -1,175 +0,0 @@ -# Copyright Motphys Technology Co., Ltd. 2025, 2026 -# SPDX-License-Identifier: Apache-2.0 - -"""Quaternion operations for scalar Numba kernels.""" - -import math - -import numba -import numpy as np - - -@numba.njit(inline="always") -def from_euler( - roll: float, - pitch: float, - yaw: float, - out: np.ndarray | None = None, -) -> np.ndarray: - """Convert roll/pitch/yaw angles to an ``[x, y, z, w]`` quaternion. - - The scalar implementation is intended for use inside fused Numba lanes. - An optional caller-owned output avoids an allocation in reset and other - hot paths. - """ - if out is None: - out = np.empty(4, dtype=np.float32) - cy = math.cos(yaw * 0.5) - sy = math.sin(yaw * 0.5) - cp = math.cos(pitch * 0.5) - sp = math.sin(pitch * 0.5) - cr = math.cos(roll * 0.5) - sr = math.sin(roll * 0.5) - out[0] = sr * cp * cy - cr * sp * sy - out[1] = cr * sp * cy + sr * cp * sy - out[2] = cr * cp * sy - sr * sp * cy - out[3] = cr * cp * cy + sr * sp * sy - return out - - -@numba.njit(inline="always") -def inverse( - quaternion: np.ndarray, - out: np.ndarray | None = None, -) -> np.ndarray: - """Invert one ``[x, y, z, w]`` quaternion and return the output array.""" - if out is None: - out = np.empty(4, dtype=quaternion.dtype) - x, y, z, w = quaternion - norm_sq = max(x * x + y * y + z * z + w * w, 1e-12) - out[0] = -x / norm_sq - out[1] = -y / norm_sq - out[2] = -z / norm_sq - out[3] = w / norm_sq - return out - - -@numba.njit(inline="always") -def mul( - lhs: np.ndarray, - rhs: np.ndarray, - out: np.ndarray | None = None, -) -> np.ndarray: - """Multiply two ``[x, y, z, w]`` quaternions and return the output array.""" - if out is None: - out = np.empty(4, dtype=lhs.dtype) - x1, y1, z1, w1 = lhs - x2, y2, z2, w2 = rhs - out[0] = w1 * x2 + x1 * w2 + y1 * z2 - z1 * y2 - out[1] = w1 * y2 - x1 * z2 + y1 * w2 + z1 * x2 - out[2] = w1 * z2 + x1 * y2 - y1 * x2 + z1 * w2 - out[3] = w1 * w2 - x1 * x2 - y1 * y2 - z1 * z2 - return out - - -@numba.njit(inline="always") -def rotate_vector( - quaternion: np.ndarray, - vector: np.ndarray | tuple[float, float, float], - out: np.ndarray | None = None, -) -> np.ndarray: - """Rotate one 3D vector by the quaternion and return the output array.""" - if out is None: - out = np.empty(3, dtype=quaternion.dtype) - x, y, z, w = quaternion - vx, vy, vz = vector - tx = 2.0 * (y * vz - z * vy) - ty = 2.0 * (z * vx - x * vz) - tz = 2.0 * (x * vy - y * vx) - out[0] = vx + w * tx + y * tz - z * ty - out[1] = vy + w * ty + z * tx - x * tz - out[2] = vz + w * tz + x * ty - y * tx - return out - - -@numba.njit(inline="always") -def rotate_inverse( - quaternion: np.ndarray, - vector: np.ndarray | tuple[float, float, float], - out: np.ndarray | None = None, -) -> np.ndarray: - """Rotate one 3D vector by the inverse quaternion and return the output array.""" - if out is None: - out = np.empty(3, dtype=quaternion.dtype) - x, y, z, w = quaternion - vx, vy, vz = vector - cx = y * vz - z * vy - w * vx - cy = z * vx - x * vz - w * vy - cz = x * vy - y * vx - w * vz - out[0] = vx + 2.0 * (y * cz - z * cy) - out[1] = vy + 2.0 * (z * cx - x * cz) - out[2] = vz + 2.0 * (x * cy - y * cx) - return out - - -@numba.njit(inline="always") -def rotation_distance(lhs: np.ndarray, rhs: np.ndarray) -> float: - """Compute the rotation distance between two quaternions in radians.""" - x1, y1, z1, w1 = lhs - x2, y2, z2, w2 = rhs - dx = -w1 * x2 + x1 * w2 - y1 * z2 + z1 * y2 - dy = -w1 * y2 + x1 * z2 + y1 * w2 - z1 * x2 - dz = -w1 * z2 - x1 * y2 + y1 * x2 + z1 * w2 - imaginary_norm = min(max(math.sqrt(dx * dx + dy * dy + dz * dz), 0.0), 1.0) - return 2.0 * math.asin(imaginary_norm) - - -@numba.njit(inline="always") -def rotate_inverse_components(quaternion: np.ndarray, vector): - """Rotate one 3D vector by the inverse ``[x, y, z, w]`` quaternion. - - Allocation-free variant of :func:`rotate_inverse` that returns the three - components as a tuple, so callers inside fused kernels can destructure - them without a scratch array. - """ - x, y, z, w = quaternion - vx, vy, vz = vector - cx = y * vz - z * vy - w * vx - cy = z * vx - x * vz - w * vy - cz = x * vy - y * vx - w * vz - return ( - vx + 2.0 * (y * cz - z * cy), - vy + 2.0 * (z * cx - x * cz), - vz + 2.0 * (x * cy - y * cx), - ) - - -@numba.njit(inline="always") -def to_matrix_first_two_rows( - quaternion: np.ndarray, - out: np.ndarray | None = None, -) -> np.ndarray: - """Write the first two rows of the quaternion rotation matrix in row-major order.""" - if out is None: - out = np.empty(6, dtype=quaternion.dtype) - x, y, z, w = quaternion - two_s = 2.0 / (x * x + y * y + z * z + w * w) - out[0] = 1.0 - two_s * (y * y + z * z) - out[1] = two_s * (x * y - z * w) - out[2] = two_s * (x * z + y * w) - out[3] = two_s * (x * y + z * w) - out[4] = 1.0 - two_s * (x * x + z * z) - out[5] = two_s * (y * z - x * w) - return out - - -__all__ = [ - "from_euler", - "inverse", - "mul", - "rotate_inverse", - "rotate_inverse_components", - "rotate_vector", - "rotation_distance", - "to_matrix_first_two_rows", -] diff --git a/motrix_env_core/src/motrix_env_core/numba/program.py b/motrix_env_core/src/motrix_env_core/numba/program.py deleted file mode 100644 index 92975080..00000000 --- a/motrix_env_core/src/motrix_env_core/numba/program.py +++ /dev/null @@ -1,18 +0,0 @@ -# Copyright Motphys Technology Co., Ltd. 2025, 2026 -# SPDX-License-Identifier: Apache-2.0 - -from dataclasses import dataclass - -import numpy as np - -from motrix_env_core.numba.kernel import CPUDispatcher - - -@dataclass(frozen=True) -class NumbaTaskProgram: - """Fused reset/evaluate/observe kernels paired with environment-local reward weights.""" - - evaluate_kernel: CPUDispatcher - observe_kernel: CPUDispatcher - reset_kernel: CPUDispatcher - reward_weights: np.ndarray diff --git a/motrix_env_core/src/motrix_env_core/numba/threading.py b/motrix_env_core/src/motrix_env_core/numba/threading.py deleted file mode 100644 index 611292d0..00000000 --- a/motrix_env_core/src/motrix_env_core/numba/threading.py +++ /dev/null @@ -1,59 +0,0 @@ -# Copyright Motphys Technology Co., Ltd. 2025, 2026 -# SPDX-License-Identifier: Apache-2.0 - -"""Runtime threading-layer helpers for numba parallel kernels. - -On many-core machines the default OpenMP layer pays a per-region thread-wakeup -storm (``OMP_WAIT_POLICY=PASSIVE`` parks every worker between calls, and each -``prange`` region must wake them all before any work starts). For short kernels -called every environment step this overhead dominates the kernel itself. The -TBB layer keeps a resident work-stealing pool and does not have this problem. - -The PyPI ``tbb`` wheel installs ``libtbb.so.12`` into ``/lib`` instead of -``site-packages``, so numba's plain ``CDLL("libtbb.so.12")`` lookup cannot find -it without ``LD_LIBRARY_PATH``. Preloading the library by full path with -``RTLD_GLOBAL`` registers its SONAME in the process, and numba's later lookup by -the same SONAME reuses the loaded handle. -""" - -from __future__ import annotations - -import ctypes -import glob -import sys -from pathlib import Path - -_TBB_SONAME = "libtbb.so.12" -_tbb_preload_result: bool | None = None - - -def preload_tbb() -> bool: - """Preload ``libtbb.so.12`` from the active environment, if present. - - Returns ``True`` when the library was loaded (or was already loaded), which - lets numba select the TBB threading layer for parallel kernels. Returns - ``False`` when the ``tbb`` extra is not installed — callers keep the numba - default layer in that case. The probe runs once and is cached: every - ``ManagerEnv`` construction calls this, and the glob + CDLL work should - not repeat per instance. - """ - global _tbb_preload_result - if _tbb_preload_result is None: - _tbb_preload_result = _probe_tbb() - return _tbb_preload_result - - -def _probe_tbb() -> bool: - if not sys.platform.startswith("linux"): - return False - candidates = sorted(glob.glob(str(Path(sys.prefix) / "lib" / f"{_TBB_SONAME}*")), reverse=True) - for path in candidates: - try: - ctypes.CDLL(path, mode=ctypes.RTLD_GLOBAL) - except OSError: - continue - return True - return False - - -__all__ = ["preload_tbb"] diff --git a/motrix_env_core/src/motrix_env_core/registry.py b/motrix_env_core/src/motrix_env_core/registry.py index aec3f7fa..c16e6b0a 100644 --- a/motrix_env_core/src/motrix_env_core/registry.py +++ b/motrix_env_core/src/motrix_env_core/registry.py @@ -71,7 +71,7 @@ def _construct_env( runtime and rejects an explicit ``sim``. """ from motrix_env_core.direct.env import DirectEnv - from motrix_env_core.numba.manager.env import ManagerEnv + from motrix_env_core.manager.env import ManagerEnv if issubclass(env_cls, ManagerEnv): return env_cls( diff --git a/motrix_env_core/src/motrix_env_core/torch_manager/__init__.py b/motrix_env_core/src/motrix_env_core/torch_manager/__init__.py new file mode 100644 index 00000000..c963f48c --- /dev/null +++ b/motrix_env_core/src/motrix_env_core/torch_manager/__init__.py @@ -0,0 +1,8 @@ +# Copyright Motphys Technology Co., Ltd. 2025, 2026 +# SPDX-License-Identifier: Apache-2.0 + +"""Torch executor for resolved Manager MDP programs.""" + +from motrix_env_core.torch_manager.runtime import TorchManagerRuntime, TorchTaskProgram + +__all__ = ["TorchManagerRuntime", "TorchTaskProgram"] diff --git a/motrix_env_core/src/motrix_env_core/torch_manager/actions.py b/motrix_env_core/src/motrix_env_core/torch_manager/actions.py new file mode 100644 index 00000000..c22e0636 --- /dev/null +++ b/motrix_env_core/src/motrix_env_core/torch_manager/actions.py @@ -0,0 +1,63 @@ +# Copyright Motphys Technology Co., Ltd. 2025, 2026 +# SPDX-License-Identifier: Apache-2.0 + +"""Pure tensor joint-position processing and selected-lane action-history reset.""" + +from __future__ import annotations + +import torch + + +def process_joint_position_actions( + action_queue: torch.Tensor, + action_ptr: torch.Tensor, + actions: torch.Tensor, + delay_steps: torch.Tensor, + action_scales: torch.Tensor, + default_angles: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Advance raw history and map delayed policy actions to position targets. + + Args: + action_queue: Raw policy-action ring buffer ``(N, W, A)``, ``W >= 2``. + action_ptr: Shared ``(1,)`` int64 cursor pointing at the latest action. + actions: New raw policy actions ``(N, A)``. + delay_steps: Per-lane int64 delay ``(N,)`` in ``[0, W - 1]``. Use + zeros when delay is disabled. + action_scales: Per-actuator target scales ``(A,)``. + default_angles: Per-actuator default positions ``(A,)``. + + Returns: + ``(targets, new_queue, new_ptr, current, previous)``. Targets, current, + and previous have shape ``(N, A)``; history and cursor retain their + input shapes. Current and previous are *raw*, not delayed or scaled. + + Matches ``ActionTerm.process`` followed by + ``JointPositionActionTerm._process``: advance the shared cursor, write the + new actions, then read the delayed slot. Inputs are not mutated. Cursor + arithmetic and indexing stay on tensors for full-graph compilation. + """ + num_envs, width, num_actions = action_queue.shape + new_ptr = (action_ptr + 1) % width + current_indices = new_ptr.reshape(1, 1, 1).expand(num_envs, 1, num_actions) + new_queue = action_queue.scatter(1, current_indices, actions.unsqueeze(1)) + delayed_indices = ((new_ptr - delay_steps) % width).reshape(num_envs, 1, 1).expand(-1, 1, num_actions) + delayed = new_queue.gather(1, delayed_indices).squeeze(1) + current = new_queue.gather(1, current_indices).squeeze(1) + previous_indices = ((new_ptr - 1) % width).reshape(1, 1, 1).expand(num_envs, 1, num_actions) + previous = new_queue.gather(1, previous_indices).squeeze(1) + targets = delayed * action_scales + default_angles + return targets, new_queue, new_ptr, current, previous + + +def reset_action_history(action_queue: torch.Tensor, reset_mask: torch.Tensor) -> torch.Tensor: + """Return history with selected lanes zeroed, without changing the cursor. + + ``reset_mask`` is a fixed-shape ``(N,)`` boolean tensor, not dynamic lane + indices. Unselected lanes retain all ``(W, A)`` entries. This pure helper + does not resample delays; callers own reset-time randomization separately. + """ + return torch.where(reset_mask[:, None, None], torch.zeros_like(action_queue), action_queue) + + +__all__ = ["process_joint_position_actions", "reset_action_history"] diff --git a/motrix_env_core/src/motrix_env_core/torch_manager/binding.py b/motrix_env_core/src/motrix_env_core/torch_manager/binding.py new file mode 100644 index 00000000..828ea31d --- /dev/null +++ b/motrix_env_core/src/motrix_env_core/torch_manager/binding.py @@ -0,0 +1,20 @@ +# Copyright Motphys Technology Co., Ltd. 2025, 2026 +# SPDX-License-Identifier: Apache-2.0 + +"""Manager metadata binding for the Torch executor.""" + +from motrix_env_core.manager.env import ManagerEnv +from motrix_env_core.manager.resolver import ManagerProgramResolver, ResolvedManagerProgram + + +def build_manager_binding(env: ManagerEnv) -> ResolvedManagerProgram: + """Resolve an already constructed manager environment before ``init_state``. + + The result supplies the read plan, layouts, prepared inputs, and Python + term entries to the Torch runtime; this helper neither initializes the + environment nor compiles any numerical kernels. + """ + return ManagerProgramResolver(env).build() + + +__all__ = ["build_manager_binding"] diff --git a/motrix_env_core/src/motrix_env_core/torch_manager/rand.py b/motrix_env_core/src/motrix_env_core/torch_manager/rand.py new file mode 100644 index 00000000..3a78a9c8 --- /dev/null +++ b/motrix_env_core/src/motrix_env_core/torch_manager/rand.py @@ -0,0 +1,49 @@ +# Copyright Motphys Technology Co., Ltd. 2025, 2026 +# SPDX-License-Identifier: Apache-2.0 + +"""Torch SplitMix64 primitives for manager environments.""" + +import torch + +_GOLDEN_RATIO_64 = -0x61C8864680B583EB # 0x9E3779B97F4A7C15 as int64 +_MIX_MULTIPLIER_1 = -4658895280553007687 # 0xBF58476D1CE4E5B9 as int64 +_MIX_MULTIPLIER_2 = -7723592293110705685 # 0x94D049BB133111EB as int64 + +# Masks make an arithmetic signed right shift equivalent to a uint64 logical shift. +_MASK_SHIFT_30 = (1 << 34) - 1 +_MASK_SHIFT_27 = (1 << 37) - 1 +_MASK_SHIFT_31 = (1 << 33) - 1 +_MASK_SHIFT_40 = (1 << 24) - 1 + + +def _logical_rshift(value: torch.Tensor, amount: int, mask: int) -> torch.Tensor: + return (value >> amount) & mask + + +def _mix(value: torch.Tensor) -> torch.Tensor: + value = (value ^ _logical_rshift(value, 30, _MASK_SHIFT_30)) * _MIX_MULTIPLIER_1 + value = (value ^ _logical_rshift(value, 27, _MASK_SHIFT_27)) * _MIX_MULTIPLIER_2 + return value ^ _logical_rshift(value, 31, _MASK_SHIFT_31) + + +def draw_uniform(state: torch.Tensor, count: int) -> tuple[torch.Tensor, torch.Tensor]: + """Draw ``count`` SplitMix64 values per row, returning samples and updated state. + + ``state`` is an ``(N, 1)`` signed int64 tensor whose bits are interpreted as + uint64. Arithmetic intentionally wraps at 64 bits, and samples are float32 + values in ``[-1, 1)``. The count is fixed for a call and may be zero. + """ + if count == 0: + return torch.empty((*state.shape[:-1], 0), device=state.device, dtype=torch.float32), state.clone() + # Pre-wrap fixed offsets in Python: Inductor may otherwise use signed C++ + # loop-index multiplication, whose overflow is undefined rather than wrapping. + offsets = tuple(((index * _GOLDEN_RATIO_64 + (1 << 63)) % (1 << 64)) - (1 << 63) for index in range(1, count + 1)) + increments = torch.tensor(offsets, device=state.device, dtype=torch.int64) + draw_states = state + increments + mixed = _mix(draw_states) + top = _logical_rshift(mixed, 40, _MASK_SHIFT_40) + samples = top.to(torch.float32) * (1.0 / (1 << 23)) - 1.0 + return samples, draw_states[..., -1:] + + +__all__ = ["draw_uniform"] diff --git a/motrix_env_core/src/motrix_env_core/torch_manager/runtime.py b/motrix_env_core/src/motrix_env_core/torch_manager/runtime.py new file mode 100644 index 00000000..1c80a877 --- /dev/null +++ b/motrix_env_core/src/motrix_env_core/torch_manager/runtime.py @@ -0,0 +1,868 @@ +# Copyright Motphys Technology Co., Ltd. 2025, 2026 +# SPDX-License: Apache-2.0 + +"""Generic Torch executor for resolved Manager programs. + +The runtime consumes the metadata produced by the manager resolver (read +plan, layout, and prepared invocations) and executes the same term schedule +as the previous fused Numba kernels — command ``update`` hooks, rewards, +terminations, command ``advance`` hooks for evaluate; per-group observation +terms for observe; command ``reset_env`` hooks plus simulator reset terms +for reset — with one batched tensor call per term instead of a per-lane +loop. All writes go back into the persistent numpy buffers the host +lifecycle already owns, so ``ManagerEnv`` keeps its current API and state +surfaces unchanged. + +By default each kernel schedule is additionally fused with ``torch.compile`` +into a compiled closure over the captured persistent tensors (one fused +graph per kernel instead of one eager dispatch per tensor op; the reset +graph computes over the full batch with a lane mask so every shape stays +static). The eager schedule +remains available as a debugging fallback +(``TorchManagerRuntime(env, program, compile=False)``). +""" + +from __future__ import annotations + +import os +import re +import warnings +from dataclasses import dataclass +from typing import Any + +import numpy as np +import torch + +from motrix_env_core.manager.kernel_data import KernelDataScope +from motrix_env_core.torch_manager.rand import draw_uniform + +_INPUT_EXPRESSION = re.compile(r"^input_(\d+)(\[env_id\])?$") + +# Simulator-backed read-only arrays are wrapped as tensors but never written +# by the executor; suppress PyTorch's non-writable-array warning once instead +# of copying every read-only leaf on every view build. +warnings.filterwarnings( + "ignore", + message="The given NumPy array is not writable", + category=UserWarning, +) + + +def _compile_options() -> dict[str, Any]: + """Inductor options for compiled term schedules. + + Small fused graphs oversubscribe when every generated loop grabs all + OpenMP threads; capping intra-graph threads near 8 is faster for + manager-term-sized ops. + """ + threads = int(os.environ.get("MOTRIX_TORCH_COMPILE_THREADS", "0")) or min(8, torch.get_num_threads()) + # Inductor serializes any loop whose per-thread work would fall below + # ``cpp.min_chunk_size`` (default 4096 elements — with capped threads that + # means whole-batch loops under ~16k elements run single-threaded inside + # ``#pragma omp single``). Manager term loops are exactly that size, so + # lower the threshold to let them parallelize. + chunk = int(os.environ.get("MOTRIX_TORCH_MIN_CHUNK", "256")) + return {"cpp.threads": threads, "cpp.min_chunk_size": chunk} + + +@dataclass(frozen=True) +class TorchTaskProgram: + """Drop-in replacement task program backed by Torch kernels.""" + + evaluate_kernel: Any + observe_kernel: Any + reset_kernel: Any + reward_weights: np.ndarray + + +class _TensorCache: + """Tensor views over one persistent flat-input tuple. + + Simulator-backed arrays are rewritten in place by the read program, so + views stay valid for the lifetime of the input tuple. The cache rebuilds + only when the tuple identity changes. + """ + + def __init__(self) -> None: + self._inputs: tuple[Any, ...] | None = None + self._tensors: dict[int, torch.Tensor] = {} + + def tensors(self, inputs: tuple[Any, ...]) -> dict[int, torch.Tensor]: + if self._inputs is not inputs: + self._inputs = inputs + self._tensors = {} + return self._tensors + + def tensor(self, inputs: tuple[Any, ...], slot: int) -> torch.Tensor: + tensors = self.tensors(inputs) + tensor = tensors.get(slot) + if tensor is None: + value = inputs[slot] + if isinstance(value, np.ndarray): + # 0-dim scalar leaves are configuration constants (body + # indices, sizes, noise amplitudes, mode flags); as tensors + # they force tensor-dependent control flow that Dynamo cannot + # trace (``if amplitude != 0`` / ``if sequential_clips``) and + # a scalar extraction Inductor cannot compile. Python scalars + # carry no dtype-precision concern and specialize branches at + # trace time. + if value.shape == (): + if np.issubdtype(value.dtype, np.integer): + tensors[slot] = int(value) + return tensors[slot] + if np.issubdtype(value.dtype, np.floating): + tensors[slot] = float(value) + return tensors[slot] + if np.issubdtype(value.dtype, np.bool_): + tensors[slot] = bool(value) + return tensors[slot] + if value.dtype == np.uint64: + value = value.view(np.int64) + tensor = torch.from_numpy(value) + elif isinstance(value, np.integer): + tensors[slot] = int(value) + return tensors[slot] + elif isinstance(value, (np.floating, float)): + tensors[slot] = float(value) + return tensors[slot] + elif isinstance(value, (np.bool_, bool)): + tensors[slot] = bool(value) + return tensors[slot] + else: + tensor = torch.as_tensor(value) + tensors[slot] = tensor + return tensor + + +def _leaf_getter(cache: _TensorCache, inputs: tuple[Any, ...]): + """Leaf resolver over one input tuple. + + All tensor views are materialized eagerly so the getter body is a plain + list index — no per-slot branches for Dynamo to specialize when a graph + break forces it to compile as its own frame. + """ + + tensors = [cache.tensor(inputs, slot) for slot in range(len(inputs))] + + def get(slot: int) -> torch.Tensor: + return tensors[slot] + + return get + + +class _Namespace: + """Attribute access over one subtree of the flattened manager context. + + Leaf paths resolve to tensors through ``get``; intermediate paths build + child namespaces; names absent from the layout fall back to the owning + host object (plain Python attributes such as scalar parameters). + Resolved children are memoized per instance: the layout is immutable, and + tensor views stay valid for the lifetime of the enclosing input tuple. + """ + + def __init__(self, paths: dict[tuple[str, ...], int], prefix: tuple[str, ...], host: Any, get) -> None: + self._paths = paths + self._prefix = prefix + self._host = host + self._get = get + self._resolved: dict[str, Any] = {} + + def __getattr__(self, name: str) -> Any: + if name.startswith("_"): + raise AttributeError(name) + resolved = self._resolved.get(name, _MISSING) + if resolved is not _MISSING: + return resolved + resolved = self._resolve(name) + self._resolved[name] = resolved + return resolved + + def _resolve(self, name: str) -> Any: + path = (*self._prefix, name) + if path in self._paths: + return self._get(self._paths[path]) + children = { + rest: slot for p, slot in self._paths.items() if p[: len(path)] == path and (rest := p[len(path) :]) + } + if children: + # Single-leaf wrappers (for example scalar/array term buffers) + # resolve directly to their only leaf; structured records with + # nested kernel_data build a child namespace so chained paths + # like ``motion.clip.joint_pos`` stay reachable. + if len(children) == 1: + first = next(iter(children)) + if len(first) == 1: + return self._get(children[first]) + return _Namespace(self._paths, path, self._host, self._get) + if self._host is None: + raise AttributeError(name) + value = getattr(self._host, name) + if isinstance(value, np.ndarray): + # Same 0-dim integer-constant rule as _TensorCache.tensor. + if value.shape == () and np.issubdtype(value.dtype, np.integer): + return int(value) + if value.dtype == np.uint64: + value = value.view(np.int64) + return torch.from_numpy(value) + return value + + +_MISSING = object() + + +class _BatchRand: + """Batched SplitMix64 draws over the shared per-env state buffer.""" + + def __init__(self, state: torch.Tensor) -> None: + self.state = state + + def next_uniform(self, count: int = 1) -> torch.Tensor: + """Draw ``count`` values per lane in ``[-1, 1)``; shape ``(N, count)``.""" + samples, state = draw_uniform(self.state, count) + self.state.copy_(state) + return samples + + def uniform_range(self, count: int, low: float, high: float) -> torch.Tensor: + """Draw ``count`` values per lane in ``[low, high)``; shape ``(N, count)``.""" + return low + (high - low) * np.float32(0.5) * (self.next_uniform(count) + np.float32(1.0)) + + +class _BatchActions: + """One action term's ring-buffer state with batched current/previous.""" + + def __init__(self, namespace: _Namespace) -> None: + self._namespace = namespace + + @property + def action_queue(self) -> torch.Tensor: + return self._namespace.action_queue + + @property + def action_ptr(self) -> torch.Tensor: + return self._namespace.action_ptr + + def current(self) -> torch.Tensor: + # ``action_ptr`` is the shared (1,) cursor; gather every lane's row + # from the same slot. + queue = self.action_queue + index = self.action_ptr.view(1, 1, 1).expand(queue.shape[0], 1, queue.shape[-1]) + return queue.gather(1, index).squeeze(1) + + def previous(self) -> torch.Tensor: + queue = self.action_queue + prev = (self.action_ptr.view(1, 1) - 1) % queue.shape[1] + index = prev.expand(queue.shape[0], 1, queue.shape[-1]) + return queue.gather(1, index).squeeze(1) + + def __getattr__(self, name: str) -> Any: + return getattr(self._namespace, name) + + +class _SimMap: + def __init__(self, get, slots: dict[str, int]) -> None: + self._get = get + self._slots = slots + + def __getitem__(self, key: str) -> torch.Tensor: + return self._get(self._slots[key]) + + +class _CommandsMap: + def __init__(self, paths: dict[tuple[str, ...], int], command_terms: dict[str, Any], get) -> None: + self._paths = paths + self._command_terms = command_terms + self._get = get + self._cache: dict[str, _Namespace] = {} + + def __getitem__(self, name: str) -> _Namespace: + namespace = self._cache.get(name) + if namespace is None: + prefix = ("commands", name) + if not any(path[:2] == prefix for path in self._paths): + raise KeyError(name) + namespace = _Namespace(self._paths, prefix, self._command_terms[name], self._get) + self._cache[name] = namespace + return namespace + + +class _ActionsMap: + def __init__(self, paths: dict[tuple[str, ...], int], action_terms: dict[str, Any], get) -> None: + self._paths = paths + self._action_terms = action_terms + self._get = get + self._cache: dict[str, _BatchActions] = {} + + def __getitem__(self, name: str) -> _BatchActions: + actions = self._cache.get(name) + if actions is None: + prefix = ("actions", name) + if not any(path[:2] == prefix for path in self._paths): + raise KeyError(name) + namespace = _Namespace(self._paths, prefix, self._action_terms[name].state, self._get) + actions = _BatchActions(namespace) + self._cache[name] = actions + return actions + + +class _MetricsMap: + def __init__(self, paths: dict[tuple[str, ...], int], get) -> None: + self._paths = paths + self._get = get + + def __getitem__(self, name: str) -> torch.Tensor: + return self._get(self._paths[("metrics", name)]) + + +class _BatchContext: + """Batched stand-in for the per-lane ``ManagerContext``. + + Per-env fields are full ``(N, ...)`` tensors; terms written for this + context operate on whole batches. + """ + + def __init__(self, inputs: tuple[Any, ...], cache: _TensorCache, paths: dict[tuple[str, ...], int], env: Any): + get = _leaf_getter(cache, inputs) + self._paths = paths + self.env_id = get(paths[("env_id",)]) + self.dt = float(inputs[paths[("dt",)]]) + self.sim_reset_requested = get(paths[("sim_reset_requested",)]) + self.rand = _BatchRand(get(paths[("rand", "state")])) + self.sim = _SimMap(get, {path[1]: slot for path, slot in paths.items() if len(path) == 2 and path[0] == "sim"}) + self.commands = _CommandsMap(paths, env.command_terms, get) + self.actions = _ActionsMap(paths, env.action_terms, get) + self.metrics = _MetricsMap(paths, get) + + +class _WritesMap: + def __init__(self, writes: dict[str, torch.Tensor]) -> None: + self._writes = writes + + def __getitem__(self, name: str) -> torch.Tensor: + return self._writes[name] + + +def _parse_input_expression(expression: str) -> tuple[int, bool] | None: + """Return ``(slot, is_lane)`` for ``input_N`` / ``input_N[env_id]``.""" + match = _INPUT_EXPRESSION.match(expression) + if match is None: + return None + return int(match.group(1)), match.group(2) is not None + + +class TorchManagerRuntime: + """Execute one resolved manager program with batched Torch kernels. + + ``compile`` fuses each kernel's whole term schedule into one Inductor + graph over the captured persistent tensors; ``compile=False`` keeps the + per-term eager schedule for debugging. + """ + + def __init__(self, env: Any, program: Any, compile: bool = True) -> None: + self.env = env + self.program = program + self.layout = program.layout + self.compile_enabled = compile + self._cache = _TensorCache() + + context_source = program.prepared_terms[program.context.prepared_index] + context_base = program.input_offsets[program.context.prepared_index] + self._context_paths: dict[tuple[str, ...], int] = { + tuple(field.path): context_base + field.slot_index for field in context_source.fields + } + self._per_env_slots = frozenset( + context_base + field.slot_index for field in context_source.fields if field.scope is KernelDataScope.PER_ENV + ) + + # Command hooks are static functions on the CommandTerm class taking + # the command's kernel-data namespace as their receiver. + self.invocations = tuple(program.invocations) + # The resolver groups invocations by term kind; the evaluate schedule + # must instead follow the fused-kernel statement order: command + # update hooks -> rewards -> terminations -> command advance hooks + # (rewards and terminations read pre-advance command state). + self.evaluate_invocations = tuple( + sorted( + (i for i in self.invocations if i.kind in _EVALUATE_KINDS), + key=lambda invocation: _EVALUATE_ORDER[invocation.kind], + ) + ) + self.observe_invocations = tuple(i for i in self.invocations if i.kind == "observation") + self._reward_indices = { + i.context: index for index, i in enumerate(i for i in self.evaluate_invocations if i.kind == "reward") + } + self._termination_indices = { + i.context: index for index, i in enumerate(i for i in self.evaluate_invocations if i.kind == "termination") + } + self._observation_by_name = {i.context: i for i in self.observe_invocations} + reward_cfgs = env.cfg.reward_cfgs() + self.reward_weights = np.asarray( + [float(reward_cfgs[term.name].weight) for term in self.layout.rewards], dtype=np.float32 + ) + # Reset-term compact-write layout mirrors SimResetRuntime.buffers: + # one entry per declared write, ordered by term then output name. + self._reset_terms: tuple[tuple[Any, int, tuple[str, ...]], ...] = () + offsets: dict[str, int] = {} + offset = 0 + for name, writes in env.sim_reset_writes.items(): + fields = tuple(writes) + offsets[name] = offset + offset += len(fields) + self._reset_terms = tuple( + ( + invocation, + offsets[name], + tuple(env.sim_reset_writes[name]), + ) + for invocation in program.invocations + if invocation.kind == "reset" + for name in (invocation.context[len("sim_reset.") :],) + ) + self._reset_env_invocations = tuple(i for i in self.invocations if i.kind == "command_reset_env") + + # Per-inputs-identity caches: the host reuses one persistent input + # tuple (and at most two kernel-output buffer sets), so the batched + # context and per-invocation argument tuples are built once. + self._batch_ctx: _BatchContext | None = None + self._batch_ctx_inputs: tuple[Any, ...] | None = None + self._invocation_args_cache: dict[int, dict[int, tuple[Any, ...]]] = {} + self._compiled_kernels: dict[Any, Any] = {} + + # -- kernel API --------------------------------------------------------- + + def task_program(self) -> TorchTaskProgram: + return TorchTaskProgram( + evaluate_kernel=self.evaluate, + observe_kernel=self.observe, + reset_kernel=self.reset, + reward_weights=self.reward_weights, + ) + + def evaluate(self, inputs: tuple[Any, ...], reward_weights: np.ndarray, buffers, outputs) -> None: + if self.compile_enabled: + kernel = self._compiled_kernels.get(("evaluate", id(inputs), id(buffers))) + if kernel is None: + kernel = self._build_compiled_evaluate(inputs, reward_weights, buffers, outputs) + self._compiled_kernels[("evaluate", id(inputs), id(buffers))] = kernel + kernel(torch.from_numpy(outputs.reward), torch.from_numpy(outputs.terminated)) + return + self._evaluate_eager(inputs, reward_weights, buffers, outputs) + + def _evaluate_eager(self, inputs: tuple[Any, ...], reward_weights: np.ndarray, buffers, outputs) -> None: + ctx = self._context_for(inputs) + weights = torch.from_numpy(np.asarray(reward_weights, dtype=np.float32)) + raw_terms = torch.from_numpy(buffers[0]) + weighted_terms = torch.from_numpy(buffers[1]) + termination_masks = torch.from_numpy(buffers[2]) + dt = ctx.dt + total = None + terminated = None + for invocation, args in self._evaluate_plan(inputs): + kind = invocation.kind + if kind in ("command_update", "command_advance"): + receiver = ctx.commands[invocation.receiver_key] + invocation.dispatcher(receiver, ctx) + elif kind == "reward": + value = invocation.dispatcher(ctx, *args).to(torch.float32) + index = self._reward_indices[invocation.context] + weighted = value * weights[index] * dt + raw_terms[:, index] = value + weighted_terms[:, index] = weighted + total = weighted if total is None else total + weighted + elif kind == "termination": + mask = invocation.dispatcher(ctx, *args).to(torch.bool) + index = self._termination_indices[invocation.context] + termination_masks[:, index] = mask + terminated = mask if terminated is None else terminated | mask + else: # pragma: no cover - schedule is fixed at resolve time + raise RuntimeError(f"Unsupported evaluate invocation kind {kind!r}.") + reward_view = torch.from_numpy(outputs.reward) + reward_view.copy_(total if total is not None else torch.zeros_like(reward_view)) + terminated_view = torch.from_numpy(outputs.terminated) + terminated_view.copy_(terminated if terminated is not None else torch.zeros_like(terminated_view)) + + def observe(self, inputs: tuple[Any, ...], outputs) -> None: + if self.compile_enabled: + kernel = self._compiled_kernels.get(("observe", id(inputs))) + if kernel is None: + kernel = self._build_compiled_observe(inputs, outputs) + self._compiled_kernels[("observe", id(inputs))] = kernel + value_obs = outputs.value_obs + kernel( + torch.from_numpy(outputs.policy_obs), + torch.from_numpy(value_obs) if value_obs is not None and value_obs.size else None, + ) + return + self._observe_eager(inputs, outputs) + + def _observe_eager(self, inputs: tuple[Any, ...], outputs) -> None: + ctx = self._context_for(inputs) + group_tensors = {"policy": torch.from_numpy(outputs.policy_obs)} + if outputs.value_obs is not None and outputs.value_obs.size: + group_tensors["value"] = torch.from_numpy(outputs.value_obs) + for group_name, group in self.layout.observations.items(): + group_tensor = group_tensors[group_name] + for term in group.terms: + invocation = self._observation_by_name[f"{group_name}.{term.name}"] + out = group_tensor[:, term.output_slice] + invocation.dispatcher(ctx, out, *self._invocation_args_for(invocation, inputs)) + + def reset(self, inputs: tuple[Any, ...], env_ids: np.ndarray, reset_buffers) -> None: + if not self._reset_env_invocations and not self._reset_terms: + return + ids = torch.from_numpy(np.ascontiguousarray(env_ids, dtype=np.int64)) + if self.compile_enabled: + kernel = self._compiled_kernels.get(("reset", id(inputs), id(reset_buffers))) + if kernel is None: + kernel = self._build_compiled_reset(inputs, reset_buffers) + self._compiled_kernels[("reset", id(inputs), id(reset_buffers))] = kernel + kernel(ids) + return + self._reset_eager(inputs, ids, reset_buffers) + + def _reset_eager(self, inputs: tuple[Any, ...], ids: torch.Tensor, reset_buffers) -> None: + ctx = self._context_for(inputs) + + gathered: dict[int, torch.Tensor] = {} + + def gather(slot: int, tensor: torch.Tensor) -> torch.Tensor: + cached = gathered.get(slot) + if cached is None: + cached = tensor.index_select(0, ids) + gathered[slot] = cached + return cached + + get = _leaf_getter(self._cache, inputs) + per_env = self._per_env_slots + + def lane_get(slot: int) -> torch.Tensor: + full = get(slot) + return gather(slot, full) if slot in per_env else full + + lane_rand = _LaneRand(self._context_paths[("rand", "state")], ctx, gather) + lane_ctx = _LaneContext( + ctx, + _CommandsMap(self._context_paths, self.env.command_terms, lane_get), + lane_rand, + _SimMap( + lane_get, + {path[1]: slot for path, slot in self._context_paths.items() if len(path) == 2 and path[0] == "sim"}, + ), + _MetricsMap(self._context_paths, lane_get), + ids, + actions=_ActionsMap(self._context_paths, self.env.action_terms, lane_get), + ) + for invocation in self._reset_env_invocations: + receiver = lane_ctx.commands[invocation.receiver_key] + invocation.dispatcher(receiver, lane_ctx) + for invocation, offset, fields in self._reset_terms: + sim_writes = { + name: torch.from_numpy(reset_buffers[offset + index][: ids.shape[0]]) + for index, name in enumerate(fields) + } + invocation.dispatcher(lane_ctx, _WritesMap(sim_writes), *self._invocation_args_for(invocation, inputs)) + for slot, tensor in gathered.items(): + full = self._cache.tensor(inputs, slot) + if tensor.shape == full.shape: + full.copy_(tensor) + else: + full.index_copy_(0, ids, tensor) + + # -- compiled schedules ------------------------------------------------- + + def _build_compiled_evaluate(self, inputs, reward_weights, buffers, outputs): + ctx = self._context_for(inputs) + weights = torch.from_numpy(np.asarray(reward_weights, dtype=np.float32)) + raw_terms = torch.from_numpy(buffers[0]) + weighted_terms = torch.from_numpy(buffers[1]) + termination_masks = torch.from_numpy(buffers[2]) + dt = ctx.dt + reward_indices = self._reward_indices + termination_indices = self._termination_indices + plan = tuple((invocation.kind, invocation, args) for invocation, args in self._evaluate_plan(inputs)) + + def kernel(reward_view: torch.Tensor, terminated_view: torch.Tensor) -> None: + total = None + terminated = None + for kind, invocation, args in plan: + if kind in ("command_update", "command_advance"): + receiver = ctx.commands[invocation.receiver_key] + invocation.dispatcher(receiver, ctx) + elif kind == "reward": + value = invocation.dispatcher(ctx, *args).to(torch.float32) + index = reward_indices[invocation.context] + weighted = value * weights[index] * dt + raw_terms[:, index] = value + weighted_terms[:, index] = weighted + total = weighted if total is None else total + weighted + else: + mask = invocation.dispatcher(ctx, *args).to(torch.bool) + index = termination_indices[invocation.context] + termination_masks[:, index] = mask + terminated = mask if terminated is None else terminated | mask + reward_view.copy_(total if total is not None else torch.zeros_like(reward_view)) + terminated_view.copy_(terminated if terminated is not None else torch.zeros_like(terminated_view)) + + return torch.compile(kernel, dynamic=False, options=_compile_options()) + + def _build_compiled_observe(self, inputs, outputs): + ctx = self._context_for(inputs) + observation_by_name = self._observation_by_name + plan = tuple( + (group_name, term, observation_by_name[f"{group_name}.{term.name}"]) + for group_name, group in self.layout.observations.items() + for term in group.terms + ) + invocation_args = {id(inv): self._invocation_args_for(inv, inputs) for _, _, inv in plan} + # Each term writes a private contiguous buffer; a single fused + # ``cat + copy`` assembles the group output. Writing through strided + # slice views of the group buffer forces per-term unfusable store + # regions in the generated code (aliased memory), which multiplies + # the OpenMP fork/join cost of the fused kernel. + num_envs = outputs.policy_obs.shape[0] + term_bufs = { + group_name: tuple( + torch.empty(num_envs, term.output_slice.stop - term.output_slice.start, dtype=torch.float32) + for term in group.terms + ) + for group_name, group in self.layout.observations.items() + } + per_term_plan = tuple( + ( + group_name, + tuple( + (inv, invocation_args[id(inv)], buf) + for (_group_name, _term, inv), buf in zip( + tuple(entry for entry in plan if entry[0] == group_name), + bufs, + strict=True, + ) + ), + ) + for group_name, bufs in term_bufs.items() + ) + + # One compiled graph per observation group: smaller generated kernels + # keep their code hot through the interleaved physics/evaluate passes, + # which measurably beats a single fused graph on many-core hosts. + + def make_group_kernel(entries): + def group_kernel(target: torch.Tensor) -> None: + for invocation, args, buf in entries: + invocation.dispatcher(ctx, buf, *args) + target.copy_(torch.cat([buf for _inv, _args, buf in entries], dim=1)) + + return torch.compile(group_kernel, dynamic=False, options=_compile_options()) + + group_kernels = {group_name: make_group_kernel(entries) for group_name, entries in per_term_plan} + policy_kernel = group_kernels.get("policy") + value_kernel = group_kernels.get("value") + + def kernel(policy_obs: torch.Tensor, value_obs: torch.Tensor | None) -> None: + if policy_kernel is not None: + policy_kernel(policy_obs) + if value_obs is not None and value_kernel is not None: + value_kernel(value_obs) + + return kernel + + def _build_compiled_reset(self, inputs, reset_buffers): + """Compile the reset schedule as a static full-batch masked graph. + + The reset math runs for ALL lanes with the selected lanes chosen by a + boolean mask; per-env state written by terms is clone-modified and + merged back with ``where(mask, clone, original)``, so non-selected + lanes (including their RNG state) keep their values while selected + lanes advance exactly as the gathered schedule computed them (each + lane's SplitMix64 sequence is independent). Compaction into the + reset buffers happens outside the graph via one ``index_select`` per + write field. Static shapes let Inductor parallelize every loop. + """ + ctx = self._context_for(inputs) + rand_slot = self._context_paths[("rand", "state")] + per_env = self._per_env_slots + get = _leaf_getter(self._cache, inputs) + sim_slots = {path[1]: slot for path, slot in self._context_paths.items() if len(path) == 2 and path[0] == "sim"} + sim_slot_ids = frozenset(sim_slots.values()) + command_reset_env = self._reset_env_invocations + num_envs = ctx.rand.state.shape[0] + mask = torch.zeros(num_envs, dtype=torch.bool) + all_ids = torch.arange(num_envs) + write_plan = tuple( + ( + invocation, + tuple(torch.from_numpy(reset_buffers[offset + index]) for index in range(len(fields))), + tuple( + torch.empty_like(torch.from_numpy(reset_buffers[offset + index])) for index in range(len(fields)) + ), + fields, + self._invocation_args_for(invocation, inputs), + ) + for invocation, offset, fields in self._reset_terms + ) + + def kernel() -> None: + clones: dict[int, torch.Tensor] = {} + + def lane_get(slot: int) -> torch.Tensor: + cached = clones.get(slot) + if cached is None: + full = get(slot) + # Simulator reads are never written by reset terms (writes + # go through the sim_writes map), so they pass through + # uncloned; every other per-env leaf may be mutated. + cached = full.clone() if slot in per_env and slot not in sim_slot_ids else full + clones[slot] = cached + return cached + + lane_rand = _LaneRand(rand_slot, ctx, lambda slot, _tensor: lane_get(slot)) + lane_ctx = _LaneContext( + ctx, + _CommandsMap(self._context_paths, self.env.command_terms, lane_get), + lane_rand, + _SimMap(lane_get, sim_slots), + _MetricsMap(self._context_paths, lane_get), + all_ids, + actions=_ActionsMap(self._context_paths, self.env.action_terms, lane_get), + ) + for invocation in command_reset_env: + receiver = lane_ctx.commands[invocation.receiver_key] + invocation.dispatcher(receiver, lane_ctx) + for invocation, _buffers, staging, fields, args in write_plan: + sim_writes = {name: buffer for name, buffer in zip(fields, staging, strict=True)} + invocation.dispatcher(lane_ctx, _WritesMap(sim_writes), *args) + for slot, tensor in clones.items(): + if tensor is not get(slot): + full = get(slot) + lane_mask = mask.view((num_envs,) + (1,) * (tensor.dim() - 1)) + full.copy_(torch.where(lane_mask, tensor, full)) + + compiled = torch.compile(kernel, dynamic=False, options=_compile_options()) + mask_np = mask.numpy() + compact_plan = tuple( + (buffer.numpy(), stage.numpy()) + for _invocation, buffers, staging, _fields, _args in write_plan + for buffer, stage in zip(buffers, staging, strict=True) + ) + + def wrapper(ids: torch.Tensor) -> None: + ids_np = ids.numpy() + mask_np.fill(False) + mask_np[ids_np] = True + compiled() + count = ids_np.shape[0] + for buffer, stage in compact_plan: + buffer[:count] = stage[ids_np] + + return wrapper + + # -- helpers ------------------------------------------------------------ + + def _context_for(self, inputs: tuple[Any, ...]) -> _BatchContext: + if self._batch_ctx_inputs is not inputs: + self._batch_ctx = _BatchContext(inputs, self._cache, self._context_paths, self.env) + self._batch_ctx_inputs = inputs + self._invocation_args_cache = {} + return self._batch_ctx + + def _evaluate_plan(self, inputs: tuple[Any, ...]) -> tuple[tuple[Any, tuple[Any, ...]], ...]: + plan = self._plan_cache(inputs).get("evaluate") + if plan is None: + plan = tuple( + (invocation, self._invocation_args_for(invocation, inputs)) for invocation in self.evaluate_invocations + ) + self._plan_cache(inputs)["evaluate"] = plan + return plan + + def _plan_cache(self, inputs: tuple[Any, ...]) -> dict[str, Any]: + cache = self._invocation_args_cache.get(id(inputs)) + if cache is None: + cache = {} + self._invocation_args_cache[id(inputs)] = cache + return cache + + def _invocation_args_for(self, invocation: Any, inputs: tuple[Any, ...]) -> tuple[Any, ...]: + cache = self._plan_cache(inputs) + key = id(invocation) + args = cache.get(key) + if args is None: + args = self._invocation_args(invocation, inputs) + cache[key] = args + return args + + def _invocation_args(self, invocation: Any, inputs: tuple[Any, ...]) -> tuple[Any, ...]: + args: list[Any] = [] + for value, prepared_index, expression in zip( + invocation.args_values, invocation.args_prepared_indices, invocation.args_expressions, strict=True + ): + if prepared_index is not None: + source = self.program.prepared_terms[prepared_index] + base = self.program.input_offsets[prepared_index] + paths = {tuple(field.path): base + field.slot_index for field in source.fields} + args.append(_Namespace(paths, (), value, _leaf_getter(self._cache, inputs))) + continue + parsed = _parse_input_expression(expression) + if parsed is not None: + slot, _is_lane = parsed + leaf_value = inputs[slot] + if isinstance(leaf_value, np.ndarray): + args.append(self._cache.tensor(inputs, slot)) + else: + args.append(float(leaf_value)) + continue + args.append(value) + return tuple(args) + + +class _LaneRand: + """SplitMix64 draws on gathered reset rows, updated in place.""" + + def __init__(self, rand_slot: int, base: _BatchContext, gather) -> None: + self._rand_slot = rand_slot + self._base = base + self._gather = gather + self._state: torch.Tensor | None = None + + @property + def state(self) -> torch.Tensor: + if self._state is None: + self._state = self._gather(self._rand_slot, self._base.rand.state) + return self._state + + def next_uniform(self, count: int = 1) -> torch.Tensor: + samples, state = draw_uniform(self.state, count) + self.state.copy_(state) + return samples + + def uniform_range(self, count: int, low: float, high: float) -> torch.Tensor: + return low + (high - low) * np.float32(0.5) * (self.next_uniform(count) + np.float32(1.0)) + + +class _LaneContext: + """Reset-lane context: per-env leaves gathered to the selected rows.""" + + def __init__( + self, + base: _BatchContext, + commands: _CommandsMap, + rand: _LaneRand, + sim: _SimMap, + metrics: _MetricsMap, + ids: torch.Tensor, + actions: _ActionsMap | None = None, + ) -> None: + self._base = base + self.commands = commands + self.rand = rand + self.sim = sim + self.metrics = metrics + self.env_id = ids + self.dt = base.dt + self.sim_reset_requested = base.sim_reset_requested + self.actions = actions if actions is not None else base.actions + + +_EVALUATE_KINDS = ("command_update", "reward", "termination", "command_advance") +_EVALUATE_ORDER = {"command_update": 0, "reward": 1, "termination": 2, "command_advance": 3} + + +__all__ = ["TorchManagerRuntime", "TorchTaskProgram"] diff --git a/motrix_env_core/tests/test_action_state.py b/motrix_env_core/tests/test_action_state.py index 49cd5ab1..e5c3d8bd 100644 --- a/motrix_env_core/tests/test_action_state.py +++ b/motrix_env_core/tests/test_action_state.py @@ -7,9 +7,9 @@ import numpy as np import pytest +from motrix_env_core.manager.actions import ActionState +from motrix_env_core.manager.kernel_data import canonicalize_kernel_data, kernel_data from motrix_env_core.mdp.action import JointPositionActionState, JointPositionActionTerm -from motrix_env_core.numba.kernel_data import canonicalize_kernel_data, kernel_data -from motrix_env_core.numba.manager.actions import ActionState @kernel_data diff --git a/motrix_env_core/tests/test_array_env.py b/motrix_env_core/tests/test_array_env.py index 191be8b5..387a0896 100644 --- a/motrix_env_core/tests/test_array_env.py +++ b/motrix_env_core/tests/test_array_env.py @@ -11,7 +11,7 @@ from motrix_env_core.base import EnvCfg, ObsSpace from motrix_env_core.config.scene import SceneCfg from motrix_env_core.direct.env import DirectEnv -from motrix_env_core.numba.manager.env import ManagerEnv +from motrix_env_core.manager.env import ManagerEnv from motrix_env_core.registry import _infer_sim_backend from motrix_env_core.sim.backend import RenderConfig, SimRenderer diff --git a/motrix_env_core/tests/test_array_wrapper.py b/motrix_env_core/tests/test_array_wrapper.py new file mode 100644 index 00000000..1dbd20bf --- /dev/null +++ b/motrix_env_core/tests/test_array_wrapper.py @@ -0,0 +1,25 @@ +# Copyright Motphys Technology Co., Ltd. 2025, 2026 +# SPDX-License-Identifier: Apache-2.0 + +import numpy as np +import torch + +from motrix_env_core.array_api import Array + + +def test_array_wrapper_preserves_numpy_storage_and_operators() -> None: + source = np.asarray([1.0, 2.0], dtype=np.float32) + value = Array(source) + result = (value + 1.0) * 2.0 + + assert value.native is source + np.testing.assert_allclose(result.native, [4.0, 6.0]) + + +def test_array_wrapper_preserves_torch_storage_and_operators() -> None: + source = torch.tensor([1.0, 2.0]) + value = Array(source) + result = (value + 1.0) * 2.0 + + assert value.native is source + torch.testing.assert_close(result.native, torch.tensor([4.0, 6.0])) diff --git a/motrix_env_core/tests/test_batch_mdp.py b/motrix_env_core/tests/test_batch_mdp.py new file mode 100644 index 00000000..39ce683c --- /dev/null +++ b/motrix_env_core/tests/test_batch_mdp.py @@ -0,0 +1,47 @@ +# Copyright Motphys Technology Co., Ltd. 2025, 2026 +# SPDX-License-Identifier: Apache-2.0 + +"""Shared batch reductions preserve environment independence and native dtype.""" + +import numpy as np +import pytest +import torch + +from motrix_env_core.array import mx +from motrix_env_core.manager.dispatch import dispatch + + +@dispatch +def _batch_pose(position, default, weights): + delta = position - default + return mx.sum(delta * delta * weights, axis=-1) + + +@pytest.mark.parametrize("shape", [(1, 1), (7, 29), (4, 0)]) +def test_shared_batch_pose_preserves_lane_independence(shape): + rng = np.random.default_rng(5) + position = rng.normal(size=shape).astype(np.float32) + default = rng.normal(size=shape[1]).astype(np.float32) + weights = rng.uniform(size=shape[1]).astype(np.float32) + expected = np.sum((position - default) ** 2 * weights, axis=-1) + for native_args in ( + (position, default, weights), + tuple(torch.from_numpy(value) for value in (position, default, weights)), + ): + result = _batch_pose(*native_args) + assert np.asarray(result).shape == (shape[0],) + assert np.asarray(result).dtype == np.float32 + np.testing.assert_allclose(np.asarray(result), expected, rtol=2e-6) + + +def test_sum_supports_runtime_axes_and_strided_input(): + value = np.arange(48, dtype=np.float32).reshape(4, 12)[:, ::2] + for axis in (-2, -1, 0, 1): + np.testing.assert_allclose(mx.sum(value, axis=axis), np.sum(value, axis=axis)) + for axis in (-3, 2): + with pytest.raises(ValueError, match="bounds"): + mx.sum(np.asarray(value), axis=axis) + + +def test_sum_without_axis(): + assert mx.sum(np.ones((2, 3), dtype=np.float32)) == 6 diff --git a/motrix_env_core/tests/test_function_fingerprint.py b/motrix_env_core/tests/test_function_fingerprint.py deleted file mode 100644 index d0e9f4d9..00000000 --- a/motrix_env_core/tests/test_function_fingerprint.py +++ /dev/null @@ -1,113 +0,0 @@ -# Copyright Motphys Technology Co., Ltd. 2025, 2026 -# SPDX-License-Identifier: Apache-2.0 - -import os -import subprocess -import sys - -import numba - -from motrix_env_core.numba.fingerprint import function_fingerprint - - -def test_source_unavailable_constants_change_fingerprint(): - namespace = {"__name__": "fingerprint_test"} - exec("def entry(): return 1", namespace) - first = function_fingerprint(namespace["entry"]) - exec("def entry(): return 2", namespace) - assert function_fingerprint(namespace["entry"]) != first - - -def test_helper_changes_and_cycles(): - namespace = {"__name__": "fingerprint_test"} - exec("def helper(): return entry() + 1\ndef entry(): return helper()", namespace) - entry = namespace["entry"] - first = function_fingerprint(entry) - assert function_fingerprint(entry) == first - exec("def helper(): return entry() + 2", namespace) - assert function_fingerprint(entry) != first - - -def test_closure_constants_and_helpers(): - def make(value): - def entry(): - return value - - return entry - - assert function_fingerprint(make(1)) != function_fingerprint(make(2)) - - def helper(): - return 1 - - entry = make(helper) - first = function_fingerprint(entry) - helper.__code__ = (lambda: 2).__code__ - assert function_fingerprint(entry) != first - - -def test_defaults_keyword_defaults_and_helper_dependencies(): - def entry(value=1, *, scale=2): - return value * scale - - first = function_fingerprint(entry) - entry.__defaults__ = (3,) - second = function_fingerprint(entry) - assert second != first - entry.__kwdefaults__ = {"scale": 4} - assert function_fingerprint(entry) != second - - def helper(): - return 1 - - entry.__defaults__ = (helper,) - first = function_fingerprint(entry) - helper.__code__ = (lambda: 2).__code__ - assert function_fingerprint(entry) != first - - -def test_numba_dispatcher_matches_python_function(): - def helper(value): - return value + 1 - - assert function_fingerprint(numba.njit(helper)) == function_fingerprint(helper) - - -def test_fingerprint_is_stable_across_hash_seeds(tmp_path): - script = tmp_path / "fingerprint.py" - script.write_text( - "from motrix_env_core.numba.fingerprint import function_fingerprint\n" - "def entry(value):\n" - " return value in {'alpha', 'beta', 'gamma'}\n" - "print(function_fingerprint(entry))\n", - encoding="utf-8", - ) - fingerprints = [] - for seed in ("1", "2"): - result = subprocess.run( - [sys.executable, str(script)], - check=True, - capture_output=True, - text=True, - env={**os.environ, "PYTHONHASHSEED": seed, "PYTHONPATH": "motrix_env_core/src"}, - ) - fingerprints.append(result.stdout.strip()) - assert fingerprints[0] == fingerprints[1] - - -def test_fingerprint_is_stable_after_inline_helper_compilation(): - namespace = {"__name__": "fingerprint_test", "numba": numba} - exec( - "@numba.njit(inline='always')\n" - "def helper(value):\n" - " if value > 0:\n" - " return min(value + 1, 10)\n" - " return max(value - 1, -10)\n" - "def entry(value):\n" - " return helper(value)\n", - namespace, - ) - entry = namespace["entry"] - first = function_fingerprint(entry) - assert numba.njit(entry)(2) == 3 - assert function_fingerprint(entry) == first diff --git a/motrix_env_core/tests/test_kernel_clone.py b/motrix_env_core/tests/test_kernel_clone.py index 8fa78381..1d6f502a 100644 --- a/motrix_env_core/tests/test_kernel_clone.py +++ b/motrix_env_core/tests/test_kernel_clone.py @@ -3,16 +3,15 @@ from collections import namedtuple -import numba import numpy as np import pytest -from motrix_env_core.numba.kernel import clone_kernel_value +from motrix_env_core.manager.kernel import clone_kernel_value @pytest.mark.parametrize("layout", ["c", "f", "strided", "reversed", "broadcast", "empty", "empty_strided"]) @pytest.mark.parametrize("readonly", [False, True]) -def test_kernel_clone_preserves_numba_array_type_and_isolates_storage(layout, readonly): +def test_kernel_clone_preserves_array_type_and_isolates_storage(layout, readonly): source = np.arange(48, dtype=np.float32).reshape(8, 6) if layout == "f": source = np.asfortranarray(source) @@ -29,8 +28,8 @@ def test_kernel_clone_preserves_numba_array_type_and_isolates_storage(layout, re source.flags.writeable = not readonly copied = clone_kernel_value(source) - assert numba.typeof(copied) == numba.typeof(source) assert copied.strides == source.strides + assert copied.dtype == source.dtype assert not np.shares_memory(copied, source) np.testing.assert_array_equal(copied, source) if not readonly: diff --git a/motrix_env_core/tests/test_kernel_data.py b/motrix_env_core/tests/test_kernel_data.py index a8a59f8b..bdf8e2f1 100644 --- a/motrix_env_core/tests/test_kernel_data.py +++ b/motrix_env_core/tests/test_kernel_data.py @@ -3,11 +3,10 @@ from dataclasses import FrozenInstanceError -import numba import numpy as np import pytest -from motrix_env_core.numba.kernel_data import ( +from motrix_env_core.manager.kernel_data import ( KernelDataLowering, KernelDataScope, Map, @@ -15,7 +14,6 @@ flatten_kernel_data, iter_layout_leaves, kernel_data, - map_proxy, rebuild_lowered, unflatten_kernel_data, ) @@ -112,62 +110,6 @@ def test_kernel_data_schema_fingerprint_ignores_values() -> None: assert first_layout.lowered_type.__module__ == _Tracking.__module__ -def test_map_supports_python_and_numba_literal_lookup() -> None: - values = Map({"pose": _Pose(np.asarray([[1.0]]), np.asarray([2.0])), "phase": np.float32(0.5)}) - assert values["phase"] == np.float32(0.5) - with pytest.raises(KeyError, match="available entries"): - values["missing"] - - proxy_type = map_proxy("test.values", values.keys_tuple, module_name=__name__) - proxy = proxy_type((np.float32(1.0),), values["phase"]) - - @numba.njit - def read_phase(kernel_values): - return kernel_values["phase"] - - assert read_phase(proxy) == np.float32(0.5) - assert read_phase.nopython_signatures - - -def test_map_participates_in_tree_lowering_and_round_trip() -> None: - position = np.asarray([[1.0]], dtype=np.float32) - orientation = np.asarray([2.0], dtype=np.float32) - value = _MappedValues(Map({"pose": _Pose(position, orientation), "phase": np.float32(0.5)})) - - leaves, tree_def = flatten_kernel_data(value) - assert leaves == (position, orientation, np.float32(0.5)) - rebuilt = unflatten_kernel_data(tree_def, leaves) - assert rebuilt.values.keys_tuple == ("pose", "phase") - assert rebuilt.values["pose"].position is position - - layout = KernelDataLowering().lower(tree_def, context="mapped values") - lowered = rebuild_lowered(layout, leaves) - - @numba.njit - def read_values(kernel_values): - return kernel_values.values["pose"].position[0, 0] + kernel_values.values["phase"] - - assert read_values(lowered) == np.float32(1.5) - assert read_values.nopython_signatures - - -def test_generic_array_map_participates_in_numba_lowering() -> None: - position = np.asarray([[1.0, 2.0]], dtype=np.float32) - velocity = np.asarray([3.0], dtype=np.float32) - value = _ArrayMapValues(Map({"position": position, "velocity": velocity})) - - leaves, tree_def = flatten_kernel_data(value) - layout = KernelDataLowering().lower(tree_def, context="array map") - lowered = rebuild_lowered(layout, leaves) - - @numba.njit - def read_values(kernel_values): - return kernel_values.values["position"][0, 1] + kernel_values.values["velocity"][0] - - assert read_values(lowered) == np.float32(5.0) - assert read_values.nopython_signatures - - def test_map_schema_changes_tree_and_proxy_identity() -> None: first = _MappedValues(Map({"value": np.float32(1.0)})) second = _MappedValues(Map({"value": np.asarray([[1.0]], dtype=np.float32)})) @@ -180,18 +122,6 @@ def test_map_schema_changes_tree_and_proxy_identity() -> None: assert first_layout.fields[0].child.lowered_type is not second_layout.fields[0].child.lowered_type -def test_map_rejects_dynamic_numba_key() -> None: - proxy_type = map_proxy("test.dynamic", ("value",), module_name=__name__) - proxy = proxy_type(np.float32(1.0)) - - @numba.njit - def read_value(kernel_values, name): - return kernel_values[name] - - with pytest.raises(numba.TypingError): - read_value(proxy, "value") - - def test_kernel_data_rejects_unsupported_leaf() -> None: @kernel_data class _Invalid: diff --git a/motrix_env_core/tests/test_kernel_data_methods.py b/motrix_env_core/tests/test_kernel_data_methods.py deleted file mode 100644 index 5403b88a..00000000 --- a/motrix_env_core/tests/test_kernel_data_methods.py +++ /dev/null @@ -1,275 +0,0 @@ -# Copyright Motphys Technology Co., Ltd. 2025, 2026 -# SPDX-License-Identifier: Apache-2.0 - -import json -import os -import subprocess -import sys -from pathlib import Path - -import numba -import numpy as np -import pytest - -from motrix_env_core.numba.kernel_data import ( - KernelDataLowering, - KernelDataScope, - SharedArray, - flatten_kernel_data, - iter_layout_leaves, - kernel_data, - rebuild_lowered, -) -from motrix_env_core.numba.manager.dispatch import dispatch - - -@kernel_data -class _MethodData: - values: np.ndarray - bias: SharedArray - scale: np.float32 - - @dispatch - def total(self): - return self.values.sum() * self.scale + self.bias[0] - - @dispatch - def scaled(self): - return self.values * self.scale - - @dispatch - def combined(self, factor): - return self.total() * factor - - def host_only(self): - return self.values.sum() - - -@kernel_data -class _InheritedMethodData(_MethodData): - extra: np.float32 - - @dispatch - def with_extra(self): - return self.total() + self.extra - - -@kernel_data -class _OverriddenMethodData(_InheritedMethodData): - @dispatch - def total(self): - return self.values.sum() * self.scale - self.bias[0] - - -def _lower_lane(value, env_id): - """Use the same leaf scopes as manager lowering, retaining lane views.""" - leaves, tree_def = flatten_kernel_data(value) - layout = KernelDataLowering().lower(tree_def, context="kernel data methods test") - lane_leaves = tuple( - leaves[leaf.slot_index][env_id] if leaf.scope is KernelDataScope.PER_ENV else leaves[leaf.slot_index] - for leaf in iter_layout_leaves(layout) - ) - return layout, rebuild_lowered(layout, lane_leaves) - - -def _read_total(value): - # Kept at module scope so the subprocess test has a stable disk-cache locator. - return value.total() - - -def _method_value(data_type=_MethodData): - args = ( - np.asarray([[1.0, 2.0], [4.0, 8.0]], dtype=np.float32), - np.asarray([0.5], dtype=np.float32), - np.float32(2.0), - ) - return data_type(*args) if data_type is _MethodData else data_type(*args, np.float32(3.0)) - - -def test_marked_methods_compile_on_lowered_lanes_and_preserve_host_methods() -> None: - value = _method_value() - original_method = _MethodData.scaled - layout, lane = _lower_lane(value, 1) - assert layout.logical_type is _MethodData - assert np.shares_memory(lane.values, value.values) - assert lane.bias is value.bias - - read_total = numba.njit(_read_total) - - @numba.njit - def scaled(kernel_value): - return kernel_value.scaled() - - assert read_total(lane) == value.values[1].sum() * value.scale + value.bias[0] - np.testing.assert_array_equal(scaled(lane), value.values[1] * value.scale) - assert read_total.nopython_signatures - assert scaled.nopython_signatures - - # Lowering must not replace host Python functions or their batched behavior. - assert _MethodData.scaled is original_method - assert not isinstance(_MethodData.scaled, numba.core.registry.CPUDispatcher) - np.testing.assert_array_equal(value.scaled(), value.values * value.scale) - assert value.total() == value.values.sum() * value.scale + value.bias[0] - assert value.host_only() == value.values.sum() - - value.values[1, 0] = 7.0 - assert read_total(lane) == value.values[1].sum() * value.scale + value.bias[0] - - -@pytest.mark.parametrize("data_type", [_InheritedMethodData, _OverriddenMethodData]) -def test_marked_methods_follow_inheritance_and_overrides(data_type) -> None: - value = _method_value(data_type) - _, lane = _lower_lane(value, 0) - - @numba.njit - def read_methods(kernel_value): - return kernel_value.total(), kernel_value.with_extra(), kernel_value.scaled() - - total, with_extra, scaled = read_methods(lane) - expected = value.values[0].sum() * value.scale - expected += -value.bias[0] if data_type is _OverriddenMethodData else value.bias[0] - assert total == expected - assert with_extra == expected + value.extra - np.testing.assert_array_equal(scaled, value.values[0] * value.scale) - assert read_methods.nopython_signatures - - -@pytest.mark.parametrize("data_type", [_MethodData, _OverriddenMethodData]) -def test_marked_method_can_call_another_marked_method(data_type) -> None: - value = _method_value(data_type) - _, lane = _lower_lane(value, 1) - - @numba.njit - def combined(kernel_value, factor): - return kernel_value.combined(factor) - - expected = value.values[1].sum() * value.scale - expected += -value.bias[0] if data_type is _OverriddenMethodData else value.bias[0] - assert combined(lane, np.float32(4.0)) == expected * np.float32(4.0) - assert combined.nopython_signatures - - -def test_unmarked_methods_remain_host_only() -> None: - value = _method_value() - _, lane = _lower_lane(value, 0) - - @numba.njit - def host_only(kernel_value): - return kernel_value.host_only() - - with pytest.raises(numba.TypingError, match="host_only"): - host_only(lane) - assert value.host_only() == value.values.sum() - - -def test_marked_method_cannot_call_an_unmarked_method() -> None: - @kernel_data - class _CallsHostMethod(_MethodData): - @dispatch - def call_host(self): - return self.host_only() - - base = _method_value() - value = _CallsHostMethod(base.values, base.bias, base.scale) - _, lane = _lower_lane(value, 0) - - @numba.njit - def call_host(kernel_value): - return kernel_value.call_host() - - with pytest.raises(numba.TypingError, match="host_only"): - call_host(lane) - assert value.call_host() == value.values.sum() - - -def test_unmarked_override_does_not_inherit_a_marked_base_method() -> None: - @kernel_data - class _HostOverride(_MethodData): - def total(self): - return self.values.sum() - self.bias[0] - - base = _method_value() - value = _HostOverride(base.values, base.bias, base.scale) - _, lane = _lower_lane(value, 0) - read_total = numba.njit(_read_total) - with pytest.raises(numba.TypingError, match="total"): - read_total(lane) - assert value.total() == value.values.sum() - value.bias[0] - - -def test_marked_method_implementation_changes_layout_fingerprint(monkeypatch) -> None: - @kernel_data - class _MutableMethod: - values: np.ndarray - - @dispatch - def total(self): - return self.values.sum() + 1.0 - - value = _MutableMethod(np.asarray([[2.0, 3.0]], dtype=np.float32)) - leaves, tree_def = flatten_kernel_data(value) - lowering = KernelDataLowering() - first = lowering.lower(tree_def, context="before method replacement") - first_lane = rebuild_lowered(first, (leaves[0][0],)) - - @dispatch - def replacement(self): - return self.values.sum() + 2.0 - - monkeypatch.setattr(_MutableMethod, "total", replacement) - _, changed_tree = flatten_kernel_data(value) - second = lowering.lower(changed_tree, context="after method replacement") - second_lane = rebuild_lowered(second, (leaves[0][0],)) - assert first.fingerprint != second.fingerprint - assert first.lowered_type is not second.lowered_type - read_total = numba.njit(_read_total) - assert read_total(first_lane) == 6.0 - assert read_total(second_lane) == 7.0 - assert value.total() == 7.0 - - -def test_marked_methods_load_numba_disk_cache_in_a_second_process(tmp_path) -> None: - # Import the real test module rather than duplicating its classes in a script: - # dynamically generated proxies need a stable importable module for unpickling. - script = """ -import importlib -import json -import numba - -module = importlib.import_module('test_kernel_data_methods') -value = module._method_value() -layout, lane = module._lower_lane(value, 1) -read_total = numba.njit(cache=True)(module._read_total) -result = read_total(lane) -print(json.dumps({ - 'result': result, - 'fingerprint': layout.fingerprint, - 'hits': sum(read_total.stats.cache_hits.values()), - 'misses': sum(read_total.stats.cache_misses.values()), -})) -""" - env = os.environ.copy() - env["NUMBA_CACHE_DIR"] = str(tmp_path / "numba-cache") - env["PYTHONPATH"] = os.pathsep.join(filter(None, (str(Path(__file__).parent), env.get("PYTHONPATH")))) - env["NUMBA_NUM_THREADS"] = "1" - - def run_process(): - completed = subprocess.run( - [sys.executable, "-c", script], - env=env, - capture_output=True, - text=True, - timeout=120, - check=False, - ) - assert completed.returncode == 0, completed.stdout + completed.stderr - return json.loads(completed.stdout.splitlines()[-1]) - - cold = run_process() - warm = run_process() - assert cold["result"] == warm["result"] == 24.5 - assert cold["fingerprint"] == warm["fingerprint"] - assert cold["misses"] == 1 - assert cold["hits"] == 0 - assert warm["misses"] == 0 - assert warm["hits"] == 1 diff --git a/motrix_env_core/tests/test_numba_manager.py b/motrix_env_core/tests/test_manager_env.py similarity index 65% rename from motrix_env_core/tests/test_numba_manager.py rename to motrix_env_core/tests/test_manager_env.py index b2e20e31..4f46aab0 100644 --- a/motrix_env_core/tests/test_numba_manager.py +++ b/motrix_env_core/tests/test_manager_env.py @@ -1,19 +1,22 @@ # Copyright Motphys Technology Co., Ltd. 2025, 2026 # SPDX-License-Identifier: Apache-2.0 -import sys -from concurrent.futures import ThreadPoolExecutor +"""ManagerEnv lifecycle and term-registry contracts on the Torch executor. + +The fixtures build a physics-less manager env (``physics_step`` is a no-op): +action and command terms encode all observable state, so reset semantics, +observation lifetimes, perf phases, and term validation can be asserted +without a simulator backend. +""" + from dataclasses import fields, replace -from pathlib import Path -from types import SimpleNamespace import gymnasium as gym import numpy as np import pytest -from numba import njit +import torch -import motrix_env_core.numba.manager.compiler.compiler as compiler_module -import motrix_env_core.numba.manager.env as manager_env_module +import motrix_env_core.manager.env as manager_env_module from motrix_env_core.base import EnvCfg # noqa: E402 from motrix_env_core.config import configclass # noqa: E402 from motrix_env_core.config.scene import SceneCfg # noqa: E402 @@ -30,13 +33,10 @@ ManagerEnv, ManagerObservationGroupCfg, ManagerObservationsCfg, - ManagerResetCfg, ManagerRewardsCfg, ManagerTerminationsCfg, ObservationTermCfg, ObsTerm, - ResetTerm, - ResetTermCfg, RewardTerm, RewardTermCfg, SharedArray, @@ -45,13 +45,10 @@ kernel_data, metric, ) -from motrix_env_core.numba.kernel_data import Map -from motrix_env_core.numba.manager.compiler import NumbaKernelCompiler # noqa: E402 -from motrix_env_core.numba.manager.dispatch import dispatch -from motrix_env_core.numba.manager.observations import create_observation_groups -from motrix_env_core.numba.manager.rewards import create_reward_terms -from motrix_env_core.numba.manager.terminations import TerminationManager -from motrix_env_core.sim.write import CtrlTargetsWrite +from motrix_env_core.manager.dispatch import dispatch # noqa: E402 +from motrix_env_core.manager.observations import create_observation_groups +from motrix_env_core.manager.rewards import create_reward_terms +from motrix_env_core.manager.terminations import TerminationManager @kernel_data @@ -108,11 +105,10 @@ class _ObservationParams: @dispatch -def _injected_observation(ctx: ManagerContext, out: np.ndarray, params: _ObservationParams) -> None: - action: _TestActionState = ctx.actions["test"] - command: _CounterCommand = ctx.commands["counter"] - out[0] = action.source[0] * params.scale + params.offset[0] - out[1] = command.command[0] +def _injected_observation(ctx: ManagerContext, out: torch.Tensor, params: _ObservationParams) -> None: + action = ctx.actions["test"] + out[:, 0] = action.source[:, 0] * params.scale + params.offset[0] + out[:, 1] = ctx.commands["counter"].command[:, 0] @configclass(kw_only=True) @@ -129,9 +125,8 @@ def __call__(self, env: ManagerEnv) -> ObsTerm: @dispatch -def _derived_observation(ctx: ManagerContext, out: np.ndarray) -> None: - command: _CounterCommand = ctx.commands["counter"] - out[0] = command.double[0] +def _derived_observation(ctx: ManagerContext, out: torch.Tensor) -> None: + out[:, 0] = ctx.commands["counter"].double[:, 0] @configclass(kw_only=True) @@ -146,9 +141,8 @@ def __call__(self, env: ManagerEnv) -> ObsTerm: @dispatch -def _lane_observation(ctx: ManagerContext, out: np.ndarray) -> None: - action: _TestActionState = ctx.actions["test"] - out[0] = ctx.env_id + np.float32(action.reset_flags[0]) * 0.0 +def _lane_observation(ctx: ManagerContext, out: torch.Tensor) -> None: + out[:, 0] = ctx.env_id @configclass(kw_only=True) @@ -173,9 +167,8 @@ class _LaneObservationsCfg(ManagerObservationsCfg): @dispatch -def _injected_reward(ctx: ManagerContext) -> float: - action: _TestActionState = ctx.actions["test"] - return action.source[0] +def _injected_reward(ctx: ManagerContext): + return ctx.actions["test"].source[:, 0] @configclass(kw_only=True) @@ -186,13 +179,10 @@ def __call__(self, env: ManagerEnv) -> RewardTerm: @dispatch -def _injected_termination( - ctx: ManagerContext, - threshold: np.float32, -) -> bool: - action: _TestActionState = ctx.actions["test"] - ctx.metrics["source_at_termination"][0] = action.source[0] - return action.source[0] >= threshold +def _injected_termination(ctx: ManagerContext, threshold) -> bool: + value = ctx.actions["test"].source[:, 0] + ctx.metrics["source_at_termination"][:, 0] = value + return value >= threshold @configclass(kw_only=True) @@ -254,21 +244,23 @@ def __call__(self, env: ManagerEnv) -> "_CounterCommand": class _CounterCommand(CommandTerm): double: np.ndarray = metric() + @staticmethod @dispatch - def update(self, ctx: ManagerContext) -> None: - action: _TestActionState = ctx.actions["test"] - self.double[0] = 2.0 * action.source[0] + def update(receiver, ctx: ManagerContext) -> None: + receiver.double[:, 0] = 2.0 * ctx.actions["test"].source[:, 0] def reset(self, ctx) -> None: self.command[ctx.env_ids] = -1.0 + @staticmethod @dispatch - def advance(self, ctx: ManagerContext) -> None: - self.command[ctx.env_id] += ctx.dt + def advance(receiver, ctx: ManagerContext) -> None: + receiver.command[:, 0] += ctx.dt + @staticmethod @dispatch - def reset_env(self, ctx: ManagerContext) -> None: - self.command[ctx.env_id] = -1.0 + def reset_env(receiver, ctx: ManagerContext) -> None: + receiver.command[:, 0] = -1.0 @configclass @@ -303,42 +295,6 @@ def _manager_groups(*, scale: float = 3.0, threshold: float = 0.5, reward_weight ) -def _compile_manager(env: ManagerEnv): - return NumbaKernelCompiler(env).build() - - -# Share a single on-disk cache directory across the whole test module (hermetic, -# never touches the user cache), but do not clear the in-process kernel caches: -# each distinct plan is compiled once instead of per-test full recompilation -# caused by cache-clearing autouse fixtures (issue #32). -@pytest.fixture(autouse=True, scope="module") -def _shared_numba_cache_dir(tmp_path_factory): - monkeypatch = pytest.MonkeyPatch() - monkeypatch.setenv("NUMBA_CACHE_DIR", str(tmp_path_factory.mktemp("numba-cache"))) - yield - monkeypatch.undo() - - -# Referenced explicitly only by tests that verify cache isolation behavior: -# clear in-process caches and use a separate empty disk cache, restoring the -# shared caches afterwards so later tests are not order-dependent. -@pytest.fixture -def _isolated_numba_cache(tmp_path, monkeypatch): - monkeypatch.setenv("NUMBA_CACHE_DIR", str(tmp_path / "numba-cache")) - saved_kernel_cache = dict(compiler_module._KERNEL_CACHE) - saved_term_cache = dict(compiler_module._TERM_CACHE) - compiler_module._KERNEL_CACHE.clear() - compiler_module._TERM_CACHE.clear() - yield - compiler_module._KERNEL_CACHE.clear() - compiler_module._KERNEL_CACHE.update(saved_kernel_cache) - compiler_module._TERM_CACHE.clear() - compiler_module._TERM_CACHE.update(saved_term_cache) - - -_GROUPS = _manager_groups() - - @configclass class _ManagerCommandsCfg(ManagerCommandsCfg): counter: _CounterCommandCfg = _CounterCommandCfg() @@ -354,9 +310,9 @@ class _ManagerEnvCfg(ManagerBasedEnvCfg): scene: SceneCfg = SceneCfg() commands: _ManagerCommandsCfg = _ManagerCommandsCfg() actions: _ManagerActionsCfg = _ManagerActionsCfg() - observations: _ManagerObservationsCfg = _GROUPS.observations - rewards: _ManagerRewardsCfg = _GROUPS.rewards - terminations: _ManagerTerminationsCfg = _GROUPS.terminations + observations: _ManagerObservationsCfg = _manager_groups().observations + rewards: _ManagerRewardsCfg = _manager_groups().rewards + terminations: _ManagerTerminationsCfg = _manager_groups().terminations class _ManagerEnv(ManagerEnv[EnvCfg]): @@ -373,43 +329,6 @@ def physics_step(self) -> None: pass -@dispatch -def _recording_reset(ctx: ManagerContext, sim_writes: Map[np.ndarray]) -> None: - sim_writes["ctrl"][:] = 0.0 - - -@dispatch -def _noop_reset(ctx: ManagerContext, sim_writes: Map[np.ndarray]) -> None: - _ = ctx - sim_writes["ctrl"][:] = 0.0 - - -@configclass(kw_only=True) -class _DescriptorResetTermCfg(ResetTermCfg): - tag: str - - def __call__(self, env: ManagerEnv) -> ResetTerm: - del env - return ResetTerm(_noop_reset, writes={"ctrl": CtrlTargetsWrite()}) - - -@configclass -class _RecordingResetTermCfg(ResetTermCfg): - def __call__(self, env: ManagerEnv) -> ResetTerm: - del env - return ResetTerm(_recording_reset, writes={"ctrl": CtrlTargetsWrite()}) - - -@configclass -class _RecordingManagerResetCfg(ManagerResetCfg): - term: _RecordingResetTermCfg = _RecordingResetTermCfg() - - -@configclass -class _DescriptorManagerResetCfg(ManagerResetCfg): - term: ResetTermCfg = _DescriptorResetTermCfg(tag="default") - - def _counter_command(env: _ManagerEnv) -> _CounterCommand: command = env.command_terms["counter"] assert isinstance(command, _CounterCommand) @@ -425,19 +344,6 @@ def test_action_and_command_terms_are_environment_owned() -> None: assert _counter_command(env) is env.command_terms["counter"] -def test_reset_kernel_compiles_for_partial_resets() -> None: - env = ManagerEnv(_ManagerEnvCfg(sim_reset=_RecordingManagerResetCfg()), num_envs=3) - state = env.init_state() - reset_term = env.sim_reset_terms["term"] - assert isinstance(reset_term, ResetTerm) - state.terminated[:] = [True, False, True] - - env._reset_done_envs() - - assert env._task_program is not None - assert env._task_program.reset_kernel.nopython_signatures - - def test_command_rematerialization_reuses_transition_kernel_inputs(monkeypatch: pytest.MonkeyPatch) -> None: env = _ManagerEnv(num_envs=3) state = env.init_state() @@ -475,59 +381,6 @@ def record_read() -> None: assert read_count == 1 -def test_reset_descriptor_is_part_of_numba_kernel_cache_key(_isolated_numba_cache) -> None: - first = ManagerEnv( - _ManagerEnvCfg(sim_reset=_DescriptorManagerResetCfg(term=_DescriptorResetTermCfg(tag="first"))), - num_envs=1, - ) - second = ManagerEnv( - _ManagerEnvCfg(sim_reset=_DescriptorManagerResetCfg(term=_DescriptorResetTermCfg(tag="second"))), - num_envs=1, - ) - - first.init_state() - second.init_state() - - assert first.manager_layout.plan_keys == second.manager_layout.plan_keys - - -@njit(inline="always") -def _fingerprint_helper(x): - return x + 1.0 - - -@dispatch -def _fingerprint_entry(ctx: ManagerContext) -> None: - _fingerprint_helper(1.0) - - -def test_dispatch_fingerprint_tracks_inlined_helper_source() -> None: - """Editing a module-level njit helper must change the plan fingerprint (issue #54). - - The fused kernel inlines ``@njit(inline="always")`` helpers called from dispatch - entries, so the fingerprint must cover the helper's source, not only the entry's - own body. The entry body stays byte-identical here; only the global the entry - references is rebound, so a changed fingerprint proves the helper's source is - hashed (the old entry-only fingerprint would be unchanged). - """ - module = sys.modules[__name__] - original_helper = module._fingerprint_helper - before = NumbaKernelCompiler._function_fingerprint(_fingerprint_entry) - - @njit(inline="always") - def _other_helper(x): - return x - 1.0 - - try: - module._fingerprint_helper = _other_helper - rebound = NumbaKernelCompiler._function_fingerprint(_fingerprint_entry) - finally: - module._fingerprint_helper = original_helper - - assert before != rebound - assert before == NumbaKernelCompiler._function_fingerprint(_fingerprint_entry) - - def test_manager_cfg_accepts_dict_groups_and_empty_commands() -> None: cfg = ManagerBasedEnvCfg( scene=SceneCfg(), @@ -548,11 +401,11 @@ def test_manager_cfg_accepts_dict_groups_and_empty_commands() -> None: def test_multiple_action_terms_concatenate_spaces_and_receive_ordered_slices() -> None: @dispatch - def _multiple_action_observation(ctx: ManagerContext, out: np.ndarray) -> None: - first: _TestActionState = ctx.actions["test"] - second: _SecondActionState = ctx.actions["second"] - out[0] = first.source[0] - out[1:] = second.source + def _multiple_action_observation(ctx: ManagerContext, out: torch.Tensor) -> None: + first = ctx.actions["test"] + second = ctx.actions["second"] + out[:, 0] = first.source[:, 0] + out[:, 1:] = second.source @configclass(kw_only=True) class _MultipleActionObservationCfg(ObservationTermCfg): @@ -624,7 +477,7 @@ def test_rand_value_owns_seeded_per_environment_states() -> None: assert not np.array_equal(first, different) -def test_step_perf_records_standard_and_numba_manager_phases() -> None: +def test_step_perf_records_standard_and_manager_phases() -> None: env = _ManagerEnv() state = env.init_state() assert state.episode_steps.shape == (env.num_envs,) @@ -699,7 +552,8 @@ def test_manager_context_is_injected_once_and_reused_across_all_term_kinds() -> assert env.action_terms["test"] is action assert env.command_terms["counter"] is command np.testing.assert_allclose(state.obs.policy[:, 0], [1.75, 3.25]) - np.testing.assert_allclose(state.obs.policy[:, 1], -1.0) + # advance added one ctrl_dt to the reset command value of -1.0. + np.testing.assert_allclose(state.obs.policy[:, 1], -1.0 + env.cfg.ctrl_dt) np.testing.assert_allclose(state.obs.policy[:, 2], [0.5, 1.5]) np.testing.assert_allclose(state.obs.policy[:, 3], [0.0, 1.0]) np.testing.assert_allclose(state.reward, [0.005, 0.015]) @@ -708,7 +562,6 @@ def test_manager_context_is_injected_once_and_reused_across_all_term_kinds() -> np.testing.assert_array_equal(state.metrics["limit"], [False, True]) np.testing.assert_allclose(np.ravel(state.metrics["double"]), [0.5, 1.5]) np.testing.assert_allclose(state.reward_terms["source"], [0.005, 0.015]) - np.testing.assert_array_equal(state.metrics["limit"], [False, True]) assert env._compiled_manager_program is not None warmup = env.warmup() @@ -811,9 +664,10 @@ def test_final_observations_capture_command_before_reset_and_survive_steps(with_ final = state.final_obs assert final is not None np.testing.assert_array_equal(final.env_ids, [0, 2]) - np.testing.assert_allclose(final.obs.policy[:, 1], 2.0) + # Final observations capture the post-transition (post-advance) command. + np.testing.assert_allclose(final.obs.policy[:, 1], 2.0 + env.cfg.ctrl_dt) np.testing.assert_allclose(state.obs.policy[[0, 2], 1], -1.0) - np.testing.assert_allclose(state.obs.policy[1, 1], 2.0) + np.testing.assert_allclose(state.obs.policy[1, 1], 2.0 + env.cfg.ctrl_dt) if with_value: np.testing.assert_array_equal(final.obs.value[:, 0], [0, 2]) else: @@ -869,168 +723,6 @@ def test_post_reset_observation_does_not_run_command_evaluation() -> None: np.testing.assert_allclose(state.obs.policy[:, 2], [0.5, 1.0, 1.5]) -def test_materialize_source_is_safe_for_concurrent_same_plan_writers( - _isolated_numba_cache, tmp_path, monkeypatch -) -> None: - monkeypatch.setenv("NUMBA_CACHE_DIR", str(tmp_path)) - source = "generated = True\\n" - - def materialize(_: int) -> str: - return compiler_module.NumbaKernelCompiler._materialize_source(source, "plan") - - with ThreadPoolExecutor(max_workers=8) as executor: - paths = list(executor.map(materialize, range(32))) - - assert len(set(paths)) == 1 - assert Path(paths[0]).read_text(encoding="utf-8") == source - assert not list((tmp_path / "generated").glob("*.tmp")) - - -def test_build_rematerializes_source_after_cache_invalidation(_isolated_numba_cache, monkeypatch, tmp_path) -> None: - monkeypatch.setenv("NUMBA_CACHE_DIR", str(tmp_path)) - compiler_module._KERNEL_CACHE.clear() - env = _ManagerEnv(num_envs=1) - compiler = NumbaKernelCompiler(env) - calls = [] - - def compile_kernel(kind, source, filename): - del source - calls.append((kind, filename, Path(filename).exists())) - if len(calls) == 1: - raise ImportError("corrupt generated source") - return object() - - monkeypatch.setattr(compiler, "_compile_kernel", compile_kernel) - compiled = compiler.build() - # The failing evaluate kernel is rematerialized and retried; observe and - # reset kernels compile once each. - assert [call[0] for call in calls] == ["evaluate", "evaluate", "observe", "reset"] - assert compiled.layout.generated_filenames == (calls[1][1], calls[2][1], calls[3][1]) - assert all(call[2] is True for call in calls) - for plan_key in compiled.layout.plan_keys: - compiler_module._KERNEL_CACHE.pop(plan_key, None) - - -def test_specialization_cache_failure_rebuilds_callable_dispatchers_before_retry( - _isolated_numba_cache, monkeypatch -) -> None: - env = _ManagerEnv(num_envs=1) - invalidated = [] - compile_attempts = [] - - class FakeDispatcher: - def __init__(self, fail: bool = False): - self.fail = fail - - def compile(self, signature): - compile_attempts.append((self, signature)) - if self.fail: - self.fail = False - raise EOFError("truncated cache") - - class FakeTask: - def __init__(self, fail: bool = False): - self.evaluate_kernel = FakeDispatcher(fail) - self.observe_kernel = FakeDispatcher() - self.reset_kernel = FakeDispatcher() - self.reward_weights = np.empty(0, dtype=np.float32) - - class FakeCompiled: - layout = SimpleNamespace(plan_keys=("plan", "plan", "plan")) - - first_task = FakeTask(fail=True) - rebuilt_task = FakeTask() - env._task_program = first_task - env._compiled_manager_program = FakeCompiled() - compiler_module._KERNEL_CACHE["plan"] = object() - term_dispatcher = SimpleNamespace(_cache=SimpleNamespace(flush=lambda: None)) - compiler_module._TERM_CACHE[lambda: None] = term_dispatcher - - def rebuild(self): - del self - env._compiled_manager_program = FakeCompiled() - env._task_program = rebuilt_task - return SimpleNamespace(task=rebuilt_task) - - monkeypatch.setattr(compiler_module.NumbaKernelCompiler, "build", rebuild) - monkeypatch.setattr(compiler_module.NumbaKernelCompiler, "_invalidate_generated_cache", invalidated.append) - - env._kernel_buffers = () - env._kernel_outputs = () - env._state = SimpleNamespace() - env._compile_manager_specializations(()) - - assert invalidated == ["plan", "plan", "plan"] - assert env._task_program is rebuilt_task - assert first_task.evaluate_kernel in {dispatcher for dispatcher, _ in compile_attempts} - assert rebuilt_task.evaluate_kernel in {dispatcher for dispatcher, _ in compile_attempts} - assert "plan" not in compiler_module._KERNEL_CACHE - assert not compiler_module._TERM_CACHE - - -def test_numeric_values_reuse_compiled_plan_and_remain_environment_local() -> None: - env = _ManagerEnv(num_envs=1) - state = env.init_state() - assert env._compiled_manager_program is not None - compiled = _compile_manager(env) - - assert compiled.layout.plan_keys == env.manager_layout.plan_keys - assert compiled.task.evaluate_kernel is env._compiled_manager_program.task.evaluate_kernel - np.testing.assert_allclose(compiled.task.reward_weights, [2.0]) - action = env.action_terms["test"] - assert isinstance(action, ActionTerm) - action.state.source[:] = 0.25 - assert env._kernel_outputs is not None - inputs = compiled.read_plan.read(env) - compiled.task.observe_kernel(inputs, env._kernel_outputs) - np.testing.assert_allclose(state.obs.policy[:, 0], 1.75) - - second = _ManagerEnv(num_envs=1) - second.init_state() - assert env._reward_terms["source"] is not second._reward_terms["source"] - - -def _variant_groups_cfg(scale: float, threshold: float) -> _ManagerEnvCfg: - groups = _manager_groups(scale=scale, threshold=threshold) - return _ManagerEnvCfg( - observations=groups.observations, - rewards=groups.rewards, - terminations=groups.terminations, - ) - - -def test_scalar_term_args_and_ctrl_dt_share_one_compiled_plan() -> None: - """Numeric tuning values stay runtime buffers: variants reuse one plan.""" - first = ManagerEnv(_variant_groups_cfg(scale=3.0, threshold=0.5), num_envs=1) - second = ManagerEnv(_variant_groups_cfg(scale=5.0, threshold=0.1), num_envs=1) - first.init_state() - second.init_state() - - assert first.manager_layout.plan_keys == second.manager_layout.plan_keys - - for env, scale, threshold in ((first, 3.0, 0.5), (second, 5.0, 0.1)): - action = env.action_terms["test"] - assert isinstance(action, ActionTerm) - action.state.source[:] = 0.25 - state = env._state - assert state is not None - env.compute_observation(state) - np.testing.assert_allclose(state.obs.policy[:, 0], 0.25 * scale + 1.0) - env.compute_transition(state) - assert state.terminated[0] == (0.25 >= threshold) - - third_cfg = _variant_groups_cfg(scale=3.0, threshold=0.5) - third_cfg.ctrl_dt = 0.02 - third = ManagerEnv(third_cfg, num_envs=1) - third.init_state() - assert third.manager_layout.plan_keys == first.manager_layout.plan_keys - action = third.action_terms["test"] - assert isinstance(action, ActionTerm) - action.state.source[:] = 0.5 - third.compute_transition(third._state) - np.testing.assert_allclose(third._state.reward, 2.0 * 0.5 * 0.02) - - def test_duplicate_termination_config_type_is_rejected() -> None: @configclass class DuplicateTerminationsCfg(ManagerTerminationsCfg): @@ -1053,7 +745,7 @@ def test_shared_arrays_in_terms_are_writable() -> None: def test_observation_term_must_be_kernel_data_with_instance_compute() -> None: class _InvalidObservation: - def compute(self, env_id: int, out: np.ndarray) -> None: + def compute(self, env_id: int, out: torch.Tensor) -> None: out[0] = env_id @configclass(kw_only=True) @@ -1080,8 +772,8 @@ class _InvalidObservationsCfg(ManagerObservationsCfg): create_observation_groups(cfg, env) @dispatch - def _static_observation(env_id: int, out: np.ndarray) -> None: - out[0] = env_id + def _static_observation(ctx: ManagerContext, out: torch.Tensor) -> None: + out[:, 0] = ctx.env_id @configclass(kw_only=True) class _StaticObservationCfg(ObservationTermCfg): @@ -1136,9 +828,8 @@ def __call__(self, env: ManagerEnv) -> RewardTerm: def test_dispatch_term_parameter_names_are_not_constrained() -> None: @dispatch - def _renamed_reward(context: ManagerContext) -> float: - action: _TestActionState = context.actions["test"] - return action.source[0] * 2.0 + def _renamed_reward(context: ManagerContext): + return context.actions["test"].source[:, 0] * 2.0 @configclass(kw_only=True) class _RenamedRewardCfg(RewardTermCfg): diff --git a/motrix_env_core/tests/test_numba_env.py b/motrix_env_core/tests/test_manager_env_config_guard.py similarity index 87% rename from motrix_env_core/tests/test_numba_env.py rename to motrix_env_core/tests/test_manager_env_config_guard.py index d6f7b572..3a5c6779 100644 --- a/motrix_env_core/tests/test_numba_env.py +++ b/motrix_env_core/tests/test_manager_env_config_guard.py @@ -8,6 +8,6 @@ from motrix_env_core.manager import ManagerEnv # noqa: E402 -def test_numba_env_requires_manager_config() -> None: +def test_manager_env_requires_manager_config() -> None: with pytest.raises(TypeError, match="must inherit ManagerBasedEnvCfg"): ManagerEnv(EnvCfg(scene=SceneCfg())) diff --git a/motrix_env_core/tests/test_manager_sim_backend.py b/motrix_env_core/tests/test_manager_sim_backend.py index 11a831bd..05152dd7 100644 --- a/motrix_env_core/tests/test_manager_sim_backend.py +++ b/motrix_env_core/tests/test_manager_sim_backend.py @@ -9,22 +9,23 @@ import gymnasium as gym import numpy as np +import torch from motrix_env_core.config import configclass from motrix_env_core.config.scene import SceneCfg -from motrix_env_core.numba.kernel_data import kernel_data -from motrix_env_core.numba.manager.actions import ActionCfg, ActionState, ActionTerm, ManagerActionsCfg -from motrix_env_core.numba.manager.context import ManagerContext -from motrix_env_core.numba.manager.dispatch import dispatch -from motrix_env_core.numba.manager.env import ManagerBasedEnvCfg, ManagerEnv -from motrix_env_core.numba.manager.observations import ( +from motrix_env_core.manager.actions import ActionCfg, ActionState, ActionTerm, ManagerActionsCfg +from motrix_env_core.manager.context import ManagerContext +from motrix_env_core.manager.dispatch import dispatch +from motrix_env_core.manager.env import ManagerBasedEnvCfg, ManagerEnv +from motrix_env_core.manager.kernel_data import kernel_data +from motrix_env_core.manager.observations import ( ManagerObservationGroupCfg, ManagerObservationsCfg, ObservationTermCfg, ObsTerm, ) -from motrix_env_core.numba.manager.rewards import ManagerRewardsCfg, RewardTerm, RewardTermCfg -from motrix_env_core.numba.manager.terminations import ( +from motrix_env_core.manager.rewards import ManagerRewardsCfg, RewardTerm, RewardTermCfg +from motrix_env_core.manager.terminations import ( ManagerTerminationsCfg, TerminationTerm, TerminationTermCfg, @@ -237,16 +238,13 @@ def __call__(self, env, actuator_indices): @dispatch -def _sim_reading_observation(ctx: ManagerContext, out: np.ndarray) -> None: - dof_pos = ctx.sim["dof_pos"] - out[0] = dof_pos[0] - out[1] = dof_pos[1] +def _sim_reading_observation(ctx: ManagerContext, out: torch.Tensor) -> None: + out.copy_(ctx.sim["dof_pos"]) @configclass(kw_only=True) class _SimReadingObservationCfg(ObservationTermCfg): def __call__(self, env): - del env return ObsTerm(2, _sim_reading_observation) @@ -261,8 +259,8 @@ class _FakeObservationsCfg(ManagerObservationsCfg): @dispatch -def _unit_reward(ctx: ManagerContext) -> float: - return 1.0 +def _unit_reward(ctx: ManagerContext) -> torch.Tensor: + return torch.ones_like(ctx.env_id, dtype=torch.float32) @configclass(kw_only=True) @@ -270,7 +268,6 @@ class _UnitRewardCfg(RewardTermCfg): weight: float = 1.0 def __call__(self, env): - del env return RewardTerm(_unit_reward) @@ -280,14 +277,13 @@ class _FakeRewardsCfg(ManagerRewardsCfg): @dispatch -def _never_termination(ctx: ManagerContext) -> bool: - return False +def _never_termination(ctx: ManagerContext) -> torch.Tensor: + return torch.zeros_like(ctx.env_id, dtype=torch.bool) @configclass(kw_only=True) class _NeverTerminationCfg(TerminationTermCfg): def __call__(self, env): - del env return TerminationTerm(_never_termination) @@ -358,8 +354,6 @@ def test_partial_reset_hands_canonical_rows_to_runtime_and_clears_ctrl() -> None def test_manager_frontend_does_not_import_motrixsim() -> None: - code = ( - "import sys\nimport motrix_env_core.numba.manager.env\nsys.exit(0 if 'motrixsim' not in sys.modules else 1)\n" - ) + code = "import sys\nimport motrix_env_core.manager.env\nsys.exit(0 if 'motrixsim' not in sys.modules else 1)\n" result = subprocess.run([sys.executable, "-c", code], capture_output=True, text=True) assert result.returncode == 0, result.stderr diff --git a/motrix_env_core/tests/test_mdp_noise.py b/motrix_env_core/tests/test_mdp_noise.py deleted file mode 100644 index 4e626489..00000000 --- a/motrix_env_core/tests/test_mdp_noise.py +++ /dev/null @@ -1,21 +0,0 @@ -# Copyright Motphys Technology Co., Ltd. 2025, 2026 -# SPDX-License-Identifier: Apache-2.0 - -import numpy as np - -from motrix_env_core.mdp.noise import add_uniform_noise - - -def test_uniform_noise_consumes_one_sequence_across_consecutive_calls() -> None: - combined_state = np.asarray([123], dtype=np.uint64) - combined = np.zeros((6,), dtype=np.float32) - add_uniform_noise(combined, 1.0, combined_state) - - split_state = np.asarray([123], dtype=np.uint64) - first = np.zeros((2,), dtype=np.float32) - second = np.zeros((4,), dtype=np.float32) - add_uniform_noise(first, 1.0, split_state) - add_uniform_noise(second, 1.0, split_state) - - np.testing.assert_array_equal(np.concatenate((first, second)), combined) - np.testing.assert_array_equal(split_state, combined_state) diff --git a/motrix_env_core/tests/test_numba_math.py b/motrix_env_core/tests/test_numba_math.py deleted file mode 100644 index 5c198f80..00000000 --- a/motrix_env_core/tests/test_numba_math.py +++ /dev/null @@ -1,144 +0,0 @@ -# Copyright Motphys Technology Co., Ltd. 2025, 2026 -# SPDX-License-Identifier: Apache-2.0 - -import numba -import numpy as np - -from motrix_env_core.math import quaternion -from motrix_env_core.numba.math.quaternion import ( - from_euler, - inverse, - mul, - rotate_inverse, - rotate_vector, - rotation_distance, - to_matrix_first_two_rows, -) - - -@numba.njit -def _from_euler(roll: float, pitch: float, yaw: float) -> np.ndarray: - return from_euler(roll, pitch, yaw) - - -@numba.njit -def _inverse(value: np.ndarray) -> np.ndarray: - return inverse(value) - - -@numba.njit -def _mul(lhs: np.ndarray, rhs: np.ndarray) -> np.ndarray: - return mul(lhs, rhs) - - -@numba.njit -def _rotate_inverse_array(rotation: np.ndarray, vector: np.ndarray) -> np.ndarray: - return rotate_inverse(rotation, vector) - - -@numba.njit -def _rotate_vector_array(rotation: np.ndarray, vector: np.ndarray) -> np.ndarray: - return rotate_vector(rotation, vector) - - -@numba.njit -def _rotate_inverse_tuple(rotation: np.ndarray, vector: tuple[float, float, float]) -> np.ndarray: - buffer = np.full(5, np.nan, dtype=np.float32) - rotate_inverse(rotation, vector, buffer[1:4]) - return buffer - - -@numba.njit -def _rotation_distance(lhs: np.ndarray, rhs: np.ndarray) -> float: - return rotation_distance(lhs, rhs) - - -@numba.njit -def _to_matrix_first_two_rows(value: np.ndarray) -> np.ndarray: - return to_matrix_first_two_rows(value) - - -def test_numba_quaternion_from_euler_matches_numpy_and_supports_caller_owned_out() -> None: - roll, pitch, yaw = np.float32(0.2), np.float32(-0.3), np.float32(0.4) - expected = quaternion.from_euler(roll, pitch, yaw) - - np.testing.assert_allclose(_from_euler(roll, pitch, yaw), expected, atol=1e-6) - out = np.empty((4,), dtype=np.float32) - returned = from_euler(roll, pitch, yaw, out) - assert returned is out - np.testing.assert_allclose(returned, expected, atol=1e-6) - - -def test_numba_quaternion_inverse_matches_numpy_and_supports_aliased_out() -> None: - value = np.asarray([0.1, -0.2, 0.3, 0.9], dtype=np.float32) - expected = quaternion.inverse(value) - - np.testing.assert_allclose(_inverse(value), expected, atol=1e-6) - out = value.copy() - returned = inverse(out, out) - assert returned is out - np.testing.assert_allclose(returned, expected, atol=1e-6) - - -def test_numba_rotate_vector_matches_numpy_and_supports_aliased_out() -> None: - rotation = np.asarray([0.1, -0.2, 0.3, 0.9], dtype=np.float32) - rotation /= np.linalg.norm(rotation) - vector = np.asarray([1.0, -2.0, 0.5], dtype=np.float32) - expected = quaternion.rotate_vector(rotation, vector) - - np.testing.assert_allclose(_rotate_vector_array(rotation, vector), expected, atol=1e-6) - out = vector.copy() - returned = rotate_vector(rotation, out, out) - assert returned is out - np.testing.assert_allclose(returned, expected, atol=1e-6) - - -def test_numba_rotate_inverse_matches_numpy_for_array_and_tuple_vectors() -> None: - rotation = np.asarray([0.1, -0.2, 0.3, 0.9], dtype=np.float32) - rotation /= np.linalg.norm(rotation) - vector = np.asarray([1.0, -2.0, 0.5], dtype=np.float32) - expected = quaternion.rotate_inverse(rotation, vector) - - np.testing.assert_allclose(_rotate_inverse_array(rotation, vector), expected, atol=1e-6) - tuple_result = _rotate_inverse_tuple(rotation, tuple(vector)) - assert np.isnan(tuple_result[0]) and np.isnan(tuple_result[4]) - np.testing.assert_allclose(tuple_result[1:4], expected, atol=1e-6) - - out = np.empty(3, dtype=np.float32) - returned = rotate_inverse(rotation, vector, out) - assert returned is out - np.testing.assert_allclose(returned, expected, atol=1e-6) - - -def test_numba_rotation_distance_matches_numpy() -> None: - lhs = np.asarray([0.1, -0.2, 0.3, 0.9], dtype=np.float32) - rhs = np.asarray([-0.3, 0.4, 0.2, 0.8], dtype=np.float32) - lhs /= np.linalg.norm(lhs) - rhs /= np.linalg.norm(rhs) - - expected = quaternion.rotation_distance(lhs, rhs) - np.testing.assert_allclose(_rotation_distance(lhs, rhs), expected, atol=1e-6) - - -def test_numba_to_matrix_first_two_rows_matches_numpy_and_supports_overlapping_out() -> None: - value = np.asarray([0.1, -0.2, 0.3, 0.9], dtype=np.float32) - expected = quaternion.to_matrix(value)[:2].reshape(6) - - np.testing.assert_allclose(_to_matrix_first_two_rows(value), expected, atol=1e-6) - out = np.full(6, np.nan, dtype=np.float32) - out[:4] = value - returned = to_matrix_first_two_rows(out[:4], out) - assert returned is out - np.testing.assert_allclose(returned, expected, atol=1e-6) - - -def test_numba_quaternion_mul_matches_numpy_and_supports_aliased_out() -> None: - lhs = np.asarray([0.1, -0.2, 0.3, 0.9], dtype=np.float32) - rhs = np.asarray([-0.3, 0.4, 0.2, 0.8], dtype=np.float32) - expected = quaternion.mul(lhs, rhs) - - np.testing.assert_allclose(_mul(lhs, rhs), expected, atol=1e-6) - out = lhs.copy() - returned = mul(out, rhs, out) - assert returned is out - np.testing.assert_allclose(returned, expected, atol=1e-6) diff --git a/motrix_env_core/tests/test_numba_threading.py b/motrix_env_core/tests/test_numba_threading.py deleted file mode 100644 index 4c2cd30b..00000000 --- a/motrix_env_core/tests/test_numba_threading.py +++ /dev/null @@ -1,10 +0,0 @@ -# Copyright Motphys Technology Co., Ltd. 2025, 2026 -# SPDX-License-Identifier: Apache-2.0 - -from motrix_env_core.numba.threading import preload_tbb - - -def test_preload_tbb_returns_bool_without_raising() -> None: - # Whether the library exists depends on the tbb extra; the contract is - # only that discovery is side-effect-safe and reports success as bool. - assert isinstance(preload_tbb(), bool) diff --git a/motrix_env_core/tests/test_registry.py b/motrix_env_core/tests/test_registry.py index 8ba097eb..bc022dc0 100644 --- a/motrix_env_core/tests/test_registry.py +++ b/motrix_env_core/tests/test_registry.py @@ -16,7 +16,7 @@ from motrix_env_core.config import configclass from motrix_env_core.config.scene import MjcfFileCfg, RobotCfg, SceneCfg from motrix_env_core.direct.env import DirectEnv, DirectEnvCfg -from motrix_env_core.numba.manager.env import ManagerBasedEnvCfg, ManagerEnv +from motrix_env_core.manager.env import ManagerBasedEnvCfg, ManagerEnv from motrix_env_core.sim import ModelQuery from motrix_env_core.sim.backend import SimBackend from motrix_env_core.sim.model import SimModel diff --git a/motrix_env_core/tests/test_torch_manager_actions.py b/motrix_env_core/tests/test_torch_manager_actions.py new file mode 100644 index 00000000..e63d19e6 --- /dev/null +++ b/motrix_env_core/tests/test_torch_manager_actions.py @@ -0,0 +1,139 @@ +# Copyright Motphys Technology Co., Ltd. 2025, 2026 +# SPDX-License-Identifier: Apache-2.0 + +"""Tensor action processing agrees with the native manager's history lifecycle.""" + +from collections.abc import Callable + +import gymnasium as gym +import numpy as np +import pytest +import torch + +from motrix_env_core.mdp.action import JointPositionActionState, JointPositionActionTerm +from motrix_env_core.torch_manager.actions import process_joint_position_actions, reset_action_history + + +def _make_term(width: int, delayed: bool) -> JointPositionActionTerm: + num_envs, num_actions = width + 1, 3 + # Nonzero history catches off-by-one reads before fresh actions fill the ring. + queue = np.arange(num_envs * width * num_actions, dtype=np.float32).reshape(num_envs, width, num_actions) / 8 + state = JointPositionActionState( + action_queue=queue, + action_ptr=np.asarray([width - 1], dtype=np.int64), + default_angles=np.asarray([0.25, -2.0, 0.75], dtype=np.float32), + joint_lower=np.full(num_actions, -np.inf, dtype=np.float32), + joint_upper=np.full(num_actions, np.inf, dtype=np.float32), + action_scales=np.asarray([0.5, 2.0, 0.0], dtype=np.float32), + delay_steps=np.arange(num_envs, dtype=np.int64) % width if delayed else np.zeros(num_envs, dtype=np.int64), + # Keep native reset resampling deterministic without mocking its path. + delay_lo=width - 1 if delayed else 0, + delay_hi=width - 1 if delayed else 0, + ) + return JointPositionActionTerm(gym.spaces.Box(-np.inf, np.inf, shape=(num_actions,)), state) + + +def _exercise_process(process: Callable, *, width: int, delayed: bool) -> None: + term = _make_term(width, delayed) + state = term.state + queue = torch.from_numpy(state.action_queue.copy()) + ptr = torch.from_numpy(state.action_ptr.copy()) + delays = torch.from_numpy(state.delay_steps.copy()) + scales = torch.from_numpy(state.action_scales.copy()) + defaults = torch.from_numpy(state.default_angles.copy()) + rng = np.random.default_rng(17) + for _ in range(2 * width + 1): + raw = rng.uniform(-5, 5, size=state.current().shape).astype(np.float32) + actions = torch.from_numpy(raw.copy()) + queue_before, ptr_before = queue.clone(), ptr.clone() + targets, new_queue, new_ptr, current, previous = process(queue, ptr, actions, delays, scales, defaults) + expected = term.process(raw) + torch.testing.assert_close(targets, torch.from_numpy(expected)) + torch.testing.assert_close(new_queue, torch.from_numpy(state.action_queue), rtol=0, atol=0) + torch.testing.assert_close(new_ptr, torch.from_numpy(state.action_ptr), rtol=0, atol=0) + torch.testing.assert_close(current, torch.from_numpy(state.current()), rtol=0, atol=0) + torch.testing.assert_close(previous, torch.from_numpy(state.previous()), rtol=0, atol=0) + torch.testing.assert_close(queue, queue_before, rtol=0, atol=0) + torch.testing.assert_close(ptr, ptr_before, rtol=0, atol=0) + torch.testing.assert_close(actions, torch.from_numpy(raw), rtol=0, atol=0) + assert new_ptr.shape == (1,) and new_ptr.dtype == torch.int64 + queue, ptr = new_queue, new_ptr + + +@pytest.mark.parametrize("width", [2, 5]) +@pytest.mark.parametrize("delayed", [False, True]) +def test_joint_position_process_matches_native(width: int, delayed: bool) -> None: + _exercise_process(process_joint_position_actions, width=width, delayed=delayed) + + +def _exercise_reset(reset: Callable, process: Callable) -> None: + term = _make_term(4, True) + state = term.state + queue = torch.from_numpy(state.action_queue.copy()) + ptr = torch.from_numpy(state.action_ptr.copy()) + for mask in ( + np.asarray([False, True, False, True, False]), + np.zeros(5, dtype=bool), + np.ones(5, dtype=bool), + ): + before = queue.clone() + pointer_before = ptr.clone() + queue_after = reset(queue, torch.from_numpy(mask)) + term.reset(np.flatnonzero(mask)) + torch.testing.assert_close(queue_after, torch.from_numpy(state.action_queue), rtol=0, atol=0) + torch.testing.assert_close(queue, before, rtol=0, atol=0) + torch.testing.assert_close(ptr, pointer_before, rtol=0, atol=0) + torch.testing.assert_close(ptr, torch.from_numpy(state.action_ptr), rtol=0, atol=0) + # Reset randomization belongs to the host, not the pure history helper. + delays = torch.from_numpy(state.delay_steps.copy()) + raw = np.full((5, 3), 3.5, dtype=np.float32) + targets, queue, ptr, current, previous = process( + queue_after, + ptr, + torch.from_numpy(raw), + delays, + torch.from_numpy(state.action_scales), + torch.from_numpy(state.default_angles), + ) + torch.testing.assert_close(targets, torch.from_numpy(term.process(raw))) + torch.testing.assert_close(queue, torch.from_numpy(state.action_queue), rtol=0, atol=0) + torch.testing.assert_close(current, torch.from_numpy(state.current()), rtol=0, atol=0) + torch.testing.assert_close(previous, torch.from_numpy(state.previous()), rtol=0, atol=0) + + +def test_reset_action_history_matches_native_selected_lanes() -> None: + _exercise_reset(reset_action_history, process_joint_position_actions) + + +@pytest.mark.slow +def test_action_helpers_compile_fullgraph_cpu() -> None: + process = torch.compile(process_joint_position_actions, backend="inductor", fullgraph=True) + reset = torch.compile(reset_action_history, backend="inductor", fullgraph=True) + _exercise_process(process, width=4, delayed=False) + _exercise_process(process, width=4, delayed=True) + _exercise_reset(reset, process) + + +def test_batch_action_views_read_each_lanes_own_slot() -> None: + """The shared (1,) cursor must gather every lane's row, not lane 0's. + + Regression: a (1,1,A) gather index silently broadcast lane 0's action to + all lanes, corrupting the per-lane action-echo observation and the + action-rate reward. + """ + from motrix_env_core.torch_manager.runtime import _BatchActions + + num_envs, width, num_actions = 8, 4, 3 + queue = torch.arange(num_envs * width * num_actions, dtype=torch.float32).reshape(num_envs, width, num_actions) + for ptr in (0, 1, width - 1): + namespace = _FixedActionNamespace(queue, torch.tensor([ptr], dtype=torch.int64)) + actions = _BatchActions(namespace) + torch.testing.assert_close(actions.current(), queue[:, ptr, :]) + previous = queue[:, (ptr - 1) % width, :] + torch.testing.assert_close(actions.previous(), previous) + + +class _FixedActionNamespace: + def __init__(self, queue: torch.Tensor, ptr: torch.Tensor) -> None: + self.action_queue = queue + self.action_ptr = ptr diff --git a/motrix_env_core/tests/test_torch_manager_rand.py b/motrix_env_core/tests/test_torch_manager_rand.py new file mode 100644 index 00000000..c77c9498 --- /dev/null +++ b/motrix_env_core/tests/test_torch_manager_rand.py @@ -0,0 +1,61 @@ +# Copyright Motphys Technology Co., Ltd. 2025, 2026 +# SPDX-License-Identifier: Apache-2.0 + +import numpy as np +import pytest +import torch + +from motrix_env_core.manager.rand import initialize_rand_states, next_uniform +from motrix_env_core.torch_manager.rand import draw_uniform + + +def _oracle(states: np.ndarray, count: int) -> tuple[np.ndarray, np.ndarray]: + samples = np.empty((len(states), count), dtype=np.float32) + updated = states.copy() + for row, state in enumerate(updated): + state_ref = np.array([state], dtype=np.uint64) + for column in range(count): + samples[row, column] = next_uniform(state_ref) + updated[row] = state_ref[0] + return samples, updated + + +def _states() -> torch.Tensor: + boundaries = np.array([0, 1, 2**64 - 1, 2**63, 2**63 - 1], dtype=np.uint64) + seeded = initialize_rand_states(11, 12345).ravel() + bits = np.concatenate((boundaries, seeded)) + return torch.from_numpy(bits.view(np.int64))[:, None] + + +@pytest.mark.parametrize("count", [1, 7, 67]) +def test_draw_uniform_matches_host_oracle(count: int) -> None: + state = _states() + original = state.clone() + expected_samples, expected_state = _oracle(state.numpy().ravel().view(np.uint64), count) + samples, updated = draw_uniform(state, count) + np.testing.assert_array_equal(samples.numpy().view(np.uint32), expected_samples.view(np.uint32)) + np.testing.assert_array_equal(updated.numpy().ravel().view(np.uint64), expected_state) + torch.testing.assert_close(state, original) + following_samples, following_state = draw_uniform(updated, count) + expected_samples, expected_state = _oracle(expected_state, count) + np.testing.assert_array_equal(following_samples.numpy().view(np.uint32), expected_samples.view(np.uint32)) + np.testing.assert_array_equal(following_state.numpy().ravel().view(np.uint64), expected_state) + + +def test_draw_uniform_zero_count_preserves_state() -> None: + state = _states() + samples, updated = draw_uniform(state, 0) + assert samples.shape == (state.shape[0], 0) + assert samples.dtype is torch.float32 + torch.testing.assert_close(updated, state) + + +@pytest.mark.slow +@pytest.mark.parametrize("count", [0, 7, 67]) +def test_draw_uniform_compiles_with_fullgraph_inductor(count: int) -> None: + state = _states() + compiled = torch.compile(draw_uniform, backend="inductor", fullgraph=True) + actual_samples, actual_state = compiled(state, count) + expected_samples, expected_state = _oracle(state.numpy().ravel().view(np.uint64), count) + np.testing.assert_array_equal(actual_samples.numpy().view(np.uint32), expected_samples.view(np.uint32)) + np.testing.assert_array_equal(actual_state.numpy().ravel().view(np.uint64), expected_state) diff --git a/motrix_env_core/tests/typecheck/numba_env.py b/motrix_env_core/tests/typecheck/numba_env.py deleted file mode 100644 index 2ef3c34b..00000000 --- a/motrix_env_core/tests/typecheck/numba_env.py +++ /dev/null @@ -1,18 +0,0 @@ -# Copyright Motphys Technology Co., Ltd. 2025, 2026 -# SPDX-License-Identifier: Apache-2.0 - -from typing_extensions import assert_type - -from motrix_env_core.array.env import ArrayEnvState -from motrix_env_core.base import EnvCfg -from motrix_env_core.manager import ManagerEnv -from motrix_env_core.numba.kernel import NumbaKernelOutputs, clone_kernel_value - - -def check_numba_env_types(env: ManagerEnv[EnvCfg], state: ArrayEnvState) -> None: - assert_type(env.compute_transition(state), ArrayEnvState) - assert_type(env.compute_observation(state), ArrayEnvState) - - -def check_clone_kernel_value_type(outputs: NumbaKernelOutputs) -> None: - assert_type(clone_kernel_value(outputs), NumbaKernelOutputs) diff --git a/motrix_envs/src/motrix_envs/locomotion/ball_balance/mdp/observations.py b/motrix_envs/src/motrix_envs/locomotion/ball_balance/mdp/observations.py index 33490a6e..0f83e90d 100644 --- a/motrix_envs/src/motrix_envs/locomotion/ball_balance/mdp/observations.py +++ b/motrix_envs/src/motrix_envs/locomotion/ball_balance/mdp/observations.py @@ -4,6 +4,7 @@ """Observation terms for the manager-based ball-balance task.""" import numpy as np +import torch from motrix_env_core.config import configclass from motrix_env_core.manager import ( @@ -11,23 +12,17 @@ ObservationTermCfg, ObsTerm, ) -from motrix_env_core.manager.math.quaternion import rotate_inverse -from motrix_env_core.mdp.noise import add_uniform_noise +from motrix_env_core.manager.dispatch import dispatch +from motrix_env_core.math.quaternion_array import rotate_inverse as rotate_inverse_batch from motrix_env_core.mdp.observations import UniformNoiseCfg -from motrix_env_core.numba.manager.dispatch import dispatch @dispatch -def ball_relative_position_obs(ctx: ManagerContext, out: np.ndarray, noise_amplitude: np.float32) -> None: - ball_pos = ctx.sim["ball_pos"] - base_pos = ctx.sim["robot_base_pos"] - base_quat = ctx.sim["robot_base_quat"] - rotate_inverse( - base_quat, - (ball_pos[0] - base_pos[0], ball_pos[1] - base_pos[1], ball_pos[2] - base_pos[2]), - out, - ) - add_uniform_noise(out, noise_amplitude, ctx.rand.state) +def ball_relative_position_obs(ctx: ManagerContext, out: torch.Tensor, noise_amplitude: float) -> None: + relative = ctx.sim["ball_pos"] - ctx.sim["robot_base_pos"] + out.copy_(rotate_inverse_batch(ctx.sim["robot_base_quat"], relative)) + if noise_amplitude != 0.0: + out += ctx.rand.next_uniform(out.shape[1]) * noise_amplitude @configclass(kw_only=True) @@ -37,16 +32,14 @@ class BallRelativePositionObsCfg(ObservationTermCfg): noise: UniformNoiseCfg = UniformNoiseCfg() def __call__(self, ctx) -> ObsTerm: - del ctx return ObsTerm(3, ball_relative_position_obs, np.float32(self.noise.amplitude)) @dispatch -def ball_relative_velocity_obs(ctx: ManagerContext, out: np.ndarray, noise_amplitude: np.float32) -> None: - ball_lin_vel = ctx.sim["ball_lin_vel"] - base_quat = ctx.sim["robot_base_quat"] - rotate_inverse(base_quat, (ball_lin_vel[0], ball_lin_vel[1], ball_lin_vel[2]), out) - add_uniform_noise(out, noise_amplitude, ctx.rand.state) +def ball_relative_velocity_obs(ctx: ManagerContext, out: torch.Tensor, noise_amplitude: float) -> None: + out.copy_(rotate_inverse_batch(ctx.sim["robot_base_quat"], ctx.sim["ball_lin_vel"])) + if noise_amplitude != 0.0: + out += ctx.rand.next_uniform(out.shape[1]) * noise_amplitude @configclass(kw_only=True) @@ -56,13 +49,12 @@ class BallRelativeVelocityObsCfg(ObservationTermCfg): noise: UniformNoiseCfg = UniformNoiseCfg() def __call__(self, ctx) -> ObsTerm: - del ctx return ObsTerm(3, ball_relative_velocity_obs, np.float32(self.noise.amplitude)) @dispatch -def ball_position_obs(ctx: ManagerContext, out: np.ndarray) -> None: - out[:] = ctx.sim["ball_pos"] +def ball_position_obs(ctx: ManagerContext, out: torch.Tensor) -> None: + out.copy_(ctx.sim["ball_pos"]) @configclass(kw_only=True) @@ -70,13 +62,12 @@ class BallPositionObsCfg(ObservationTermCfg): """Ball position in the world frame (value observations only).""" def __call__(self, ctx) -> ObsTerm: - del ctx return ObsTerm(3, ball_position_obs) @dispatch -def ball_velocity_obs(ctx: ManagerContext, out: np.ndarray) -> None: - out[:] = ctx.sim["ball_lin_vel"] +def ball_velocity_obs(ctx: ManagerContext, out: torch.Tensor) -> None: + out.copy_(ctx.sim["ball_lin_vel"]) @configclass(kw_only=True) @@ -84,5 +75,4 @@ class BallVelocityObsCfg(ObservationTermCfg): """Ball linear velocity in the world frame (value observations only).""" def __call__(self, ctx) -> ObsTerm: - del ctx return ObsTerm(3, ball_velocity_obs) diff --git a/motrix_envs/src/motrix_envs/locomotion/ball_balance/mdp/reset.py b/motrix_envs/src/motrix_envs/locomotion/ball_balance/mdp/reset.py index ac832762..915a15ed 100644 --- a/motrix_envs/src/motrix_envs/locomotion/ball_balance/mdp/reset.py +++ b/motrix_envs/src/motrix_envs/locomotion/ball_balance/mdp/reset.py @@ -5,9 +5,8 @@ from __future__ import annotations -import math - import numpy as np +import torch from motrix_env_core.config import configclass from motrix_env_core.manager import ( @@ -15,9 +14,9 @@ ResetTerm, ResetTermCfg, ) -from motrix_env_core.numba.kernel_data import Map, SharedArray, kernel_data -from motrix_env_core.numba.manager.dispatch import dispatch -from motrix_env_core.numba.math import quaternion as numba_quaternion +from motrix_env_core.manager.dispatch import dispatch +from motrix_env_core.manager.kernel_data import Map, SharedArray, kernel_data +from motrix_env_core.math.quaternion_array import from_euler from motrix_env_core.sim import ( BodyAngularVelocityWrite, BodyLinearVelocityWrite, @@ -30,14 +29,10 @@ @dispatch def _reset_body_pos( - ctx: ManagerContext, - sim_writes: Map[np.ndarray], - spawn: tuple[np.float32, np.float32, np.float32], - noise_scale: tuple[np.float32, np.float32, np.float32], + ctx: ManagerContext, sim_writes: Map[torch.Tensor], spawn: torch.Tensor, noise_amplitude: torch.Tensor ) -> None: - position = sim_writes["position"] - for index in range(3): - position[0, index] = spawn[index] + ctx.rand.next_uniform() * noise_scale[index] + noise = ctx.rand.next_uniform(3) + sim_writes["position"][:, 0].copy_(spawn + noise * noise_amplitude) @configclass(kw_only=True) @@ -52,34 +47,24 @@ def __call__(self, ctx) -> ResetTerm: body = ctx.cfg.scene.objs.robot.resolved_base_link_name return ResetTerm( _reset_body_pos, - tuple(np.asarray(self.spawn, dtype=np.float32)), - tuple(np.asarray(self.noise, dtype=np.float32) * np.float32(self.noise_scale)), + torch.from_numpy(np.ascontiguousarray(self.spawn, dtype=np.float32)), + torch.from_numpy(np.ascontiguousarray(self.noise, dtype=np.float32) * np.float32(self.noise_scale)), writes={"position": BodyPositionWrite((body,))}, ) @dispatch -def _reset_body_rot( - ctx: ManagerContext, - sim_writes: Map[np.ndarray], - noise_scale: tuple[np.float32, np.float32, np.float32], -) -> None: - rotation = sim_writes["rotation"] - noisy_quat = np.empty(4, dtype=np.float32) - numba_quaternion.from_euler( - ctx.rand.next_uniform() * noise_scale[0], - ctx.rand.next_uniform() * noise_scale[1], - ctx.rand.next_uniform() * noise_scale[2], - noisy_quat, - ) - rotation[0] = noisy_quat - norm = math.sqrt( - rotation[0, 0] * rotation[0, 0] - + rotation[0, 1] * rotation[0, 1] - + rotation[0, 2] * rotation[0, 2] - + rotation[0, 3] * rotation[0, 3] - ) - rotation[0] /= norm +def _reset_body_rot(ctx: ManagerContext, sim_writes: Map[torch.Tensor], noise_amplitude: torch.Tensor) -> None: + # Fixed promotion order: f32 angles, f64 trig, f32 quaternion, + # sequential f32 sum of squares, sqrt via f64 rounded back to f32, and an + # f32 division. + angles = ctx.rand.next_uniform(3) * noise_amplitude + quat64 = from_euler(angles.to(torch.float64)) + quat = quat64.to(torch.float32) + squares = quat * quat + total = squares[:, 0] + squares[:, 1] + squares[:, 2] + squares[:, 3] + norm = total.to(torch.float64).sqrt().to(torch.float32) + sim_writes["rotation"][:, 0].copy_(quat / norm[:, None]) @configclass(kw_only=True) @@ -94,20 +79,15 @@ def __call__(self, ctx) -> ResetTerm: amplitude = np.asarray(self.noise, dtype=np.float32) * np.float32(self.noise_scale) return ResetTerm( _reset_body_rot, - tuple(amplitude), + torch.from_numpy(amplitude), writes={"rotation": BodyRotationWrite((body,))}, ) @dispatch -def _reset_body_lin_vel( - ctx: ManagerContext, - sim_writes: Map[np.ndarray], - noise_scale: tuple[np.float32, np.float32, np.float32], -) -> None: - linear_velocity = sim_writes["linear_velocity"] - for index in range(3): - linear_velocity[0, index] = ctx.rand.next_uniform() * noise_scale[index] +def _reset_body_lin_vel(ctx: ManagerContext, sim_writes: Map[torch.Tensor], noise_amplitude: torch.Tensor) -> None: + noise = ctx.rand.next_uniform(3) + sim_writes["linear_velocity"][:, 0].copy_(noise * noise_amplitude) @configclass(kw_only=True) @@ -122,20 +102,15 @@ def __call__(self, ctx) -> ResetTerm: amplitude = np.asarray(self.noise, dtype=np.float32) * np.float32(self.noise_scale) return ResetTerm( _reset_body_lin_vel, - tuple(amplitude), + torch.from_numpy(amplitude), writes={"linear_velocity": BodyLinearVelocityWrite((body,))}, ) @dispatch -def _reset_body_rot_vel( - ctx: ManagerContext, - sim_writes: Map[np.ndarray], - noise_scale: tuple[np.float32, np.float32, np.float32], -) -> None: - angular_velocity = sim_writes["angular_velocity"] - for index in range(3): - angular_velocity[0, index] = ctx.rand.next_uniform() * noise_scale[index] +def _reset_body_rot_vel(ctx: ManagerContext, sim_writes: Map[torch.Tensor], noise_amplitude: torch.Tensor) -> None: + noise = ctx.rand.next_uniform(3) + sim_writes["angular_velocity"][:, 0].copy_(noise * noise_amplitude) @configclass(kw_only=True) @@ -150,7 +125,7 @@ def __call__(self, ctx) -> ResetTerm: amplitude = np.asarray(self.noise, dtype=np.float32) * np.float32(self.noise_scale) return ResetTerm( _reset_body_rot_vel, - tuple(amplitude), + torch.from_numpy(amplitude), writes={"angular_velocity": BodyAngularVelocityWrite((body,))}, ) @@ -164,14 +139,11 @@ class DofResetParams: @dispatch -def _reset_body_dof_pos(ctx: ManagerContext, sim_writes: Map[np.ndarray], params: DofResetParams) -> None: - position = sim_writes["position"] - velocity = sim_writes["velocity"] - velocity[:] = 0.0 - rng = ctx.rand - for index in range(position.shape[0]): - value = params.default_pos[index] + rng.next_uniform() * params.noise_scale - position[index] = min(max(value, params.joint_lower[index]), params.joint_upper[index]) +def _reset_body_dof_pos(ctx: ManagerContext, sim_writes: Map[torch.Tensor], params: DofResetParams) -> None: + noise = ctx.rand.next_uniform(params.default_pos.shape[0]) + value = params.default_pos + noise * params.noise_scale + sim_writes["position"].copy_(value.clamp(min=params.joint_lower, max=params.joint_upper)) + sim_writes["velocity"].zero_() @configclass(kw_only=True) @@ -214,18 +186,17 @@ def __call__(self, ctx) -> ResetTerm: @dispatch def _reset_ball( ctx: ManagerContext, - sim_writes: Map[np.ndarray], - spawn: tuple[np.float32, np.float32, np.float32], - noise_scale: tuple[np.float32, np.float32, np.float32], - velocity_noise_scale: tuple[np.float32, np.float32, np.float32], + sim_writes: Map[torch.Tensor], + spawn: torch.Tensor, + noise_amplitude: torch.Tensor, + velocity_noise_amplitude: torch.Tensor, ) -> None: - position = sim_writes["position"] - linear_velocity = sim_writes["linear_velocity"] - angular_velocity = sim_writes["angular_velocity"] - for index in range(3): - position[0, index] = spawn[index] + ctx.rand.next_uniform() * noise_scale[index] - linear_velocity[0, index] = ctx.rand.next_uniform() * velocity_noise_scale[index] - angular_velocity[0, index] = 0.0 + # Position and velocity draws interleave per axis, so one 6-wide draw is + # split into even (position) and odd (velocity) slots. + samples = ctx.rand.next_uniform(6) + sim_writes["position"][:, 0].copy_(spawn + samples[:, 0::2] * noise_amplitude) + sim_writes["linear_velocity"][:, 0].copy_(samples[:, 1::2] * velocity_noise_amplitude) + sim_writes["angular_velocity"].zero_() @configclass(kw_only=True) @@ -241,9 +212,11 @@ class BallResetCfg(ResetTermCfg): def __call__(self, ctx) -> ResetTerm: return ResetTerm( _reset_ball, - tuple(np.asarray(self.spawn, dtype=np.float32)), - tuple(np.asarray(self.noise, dtype=np.float32) * np.float32(self.noise_scale)), - tuple(np.asarray(self.velocity_noise, dtype=np.float32) * np.float32(self.noise_scale)), + torch.from_numpy(np.ascontiguousarray(self.spawn, dtype=np.float32)), + torch.from_numpy(np.ascontiguousarray(self.noise, dtype=np.float32) * np.float32(self.noise_scale)), + torch.from_numpy( + np.ascontiguousarray(self.velocity_noise, dtype=np.float32) * np.float32(self.noise_scale) + ), writes={ "position": BodyPositionWrite((self.ball_link_name,)), "linear_velocity": BodyLinearVelocityWrite((self.ball_link_name,)), diff --git a/motrix_envs/src/motrix_envs/locomotion/ball_balance/mdp/rewards.py b/motrix_envs/src/motrix_envs/locomotion/ball_balance/mdp/rewards.py index a800e5d8..4696c7db 100644 --- a/motrix_envs/src/motrix_envs/locomotion/ball_balance/mdp/rewards.py +++ b/motrix_envs/src/motrix_envs/locomotion/ball_balance/mdp/rewards.py @@ -3,26 +3,26 @@ """Reward terms for the manager-based ball-balance task.""" -import math - import numpy as np +import torch from motrix_env_core.config import configclass from motrix_env_core.manager import ManagerContext, RewardTerm, RewardTermCfg -from motrix_env_core.manager.math.quaternion import rotate_inverse -from motrix_env_core.numba.kernel_data import SharedArray, kernel_data -from motrix_env_core.numba.manager.dispatch import dispatch +from motrix_env_core.manager.dispatch import dispatch +from motrix_env_core.manager.kernel_data import SharedArray, kernel_data +from motrix_env_core.math.quaternion_array import rotate_inverse as rotate_inverse_batch @dispatch -def upright_reward(ctx: ManagerContext, sigma: np.float32) -> float: +def upright_reward(ctx: ManagerContext, sigma: float) -> torch.Tensor: base_quat = ctx.sim["robot_base_quat"] - gravity = np.empty(3, dtype=np.float32) - rotate_inverse(base_quat, (0.0, 0.0, -1.0), gravity) - error_sq = gravity[0] * gravity[0] + gravity[1] * gravity[1] - aligned = gravity[2] + 1.0 + gravity = base_quat.new_zeros(base_quat.shape[0], 3) + gravity[:, 2] = -1.0 + gravity = rotate_inverse_batch(base_quat, gravity) + error_sq = gravity[:, 0] * gravity[:, 0] + gravity[:, 1] * gravity[:, 1] + aligned = gravity[:, 2] + 1.0 error_sq += aligned * aligned - return math.exp(-error_sq / (sigma * sigma)) + return torch.exp(-error_sq / (sigma * sigma)) @configclass(kw_only=True) @@ -32,15 +32,13 @@ class UprightRewardCfg(RewardTermCfg): sigma: float def __call__(self, ctx) -> RewardTerm: - del ctx return RewardTerm(upright_reward, np.float32(self.sigma)) @dispatch -def base_height_reward(ctx: ManagerContext, target_z: np.float32, sigma: np.float32) -> float: - base_pos = ctx.sim["robot_base_pos"] - error = base_pos[2] - target_z - return math.exp(-(error * error) / (sigma * sigma)) +def base_height_reward(ctx: ManagerContext, target_z: float, sigma: float) -> torch.Tensor: + error = ctx.sim["robot_base_pos"][:, 2] - target_z + return torch.exp(-(error * error) / (sigma * sigma)) @configclass(kw_only=True) @@ -51,20 +49,19 @@ class BaseHeightRewardCfg(RewardTermCfg): sigma: float def __call__(self, ctx) -> RewardTerm: - del ctx return RewardTerm(base_height_reward, np.float32(self.target_z), np.float32(self.sigma)) @dispatch -def ball_under_feet_reward(ctx: ManagerContext, sigma: np.float32) -> float: +def ball_under_feet_reward(ctx: ManagerContext, sigma: float) -> torch.Tensor: feet_pos = ctx.sim["feet_pos"] ball_pos = ctx.sim["ball_pos"] - mid_x = 0.5 * (feet_pos[0, 0] + feet_pos[1, 0]) - mid_y = 0.5 * (feet_pos[0, 1] + feet_pos[1, 1]) - error_x = ball_pos[0] - mid_x - error_y = ball_pos[1] - mid_y + mid_x = 0.5 * (feet_pos[:, 0, 0] + feet_pos[:, 1, 0]) + mid_y = 0.5 * (feet_pos[:, 0, 1] + feet_pos[:, 1, 1]) + error_x = ball_pos[:, 0] - mid_x + error_y = ball_pos[:, 1] - mid_y error_sq = error_x * error_x + error_y * error_y - return math.exp(-error_sq / (sigma * sigma)) + return torch.exp(-error_sq / (sigma * sigma)) @configclass(kw_only=True) @@ -74,7 +71,6 @@ class BallUnderFeetRewardCfg(RewardTermCfg): sigma: float def __call__(self, ctx) -> RewardTerm: - del ctx return RewardTerm(ball_under_feet_reward, np.float32(self.sigma)) @@ -85,14 +81,10 @@ class DofDefaultParams: @dispatch -def dof_default_reward(ctx: ManagerContext, params: DofDefaultParams) -> float: - dof_pos = ctx.sim["robot_dof_pos"] - error_sq = 0.0 - for joint_id in range(dof_pos.shape[0]): - error = dof_pos[joint_id] - params.reference[joint_id] - error_sq += error * error - error_sq /= dof_pos.shape[0] - return math.exp(-error_sq / (params.sigma * params.sigma)) +def dof_default_reward(ctx: ManagerContext, params: DofDefaultParams) -> torch.Tensor: + error = ctx.sim["robot_dof_pos"] - params.reference + error_sq = (error * error).sum(dim=1) / error.shape[1] + return torch.exp(-error_sq / (params.sigma * params.sigma)) @configclass(kw_only=True) diff --git a/motrix_envs/src/motrix_envs/locomotion/ball_balance/mdp/terminations.py b/motrix_envs/src/motrix_envs/locomotion/ball_balance/mdp/terminations.py index ffff6207..793f506e 100644 --- a/motrix_envs/src/motrix_envs/locomotion/ball_balance/mdp/terminations.py +++ b/motrix_envs/src/motrix_envs/locomotion/ball_balance/mdp/terminations.py @@ -4,11 +4,12 @@ """Termination terms for the manager-based ball-balance task.""" import numpy as np +import torch from motrix_env_core.config import configclass from motrix_env_core.manager import ManagerContext, TerminationTerm, TerminationTermCfg -from motrix_env_core.manager.math.quaternion import rotate_inverse -from motrix_env_core.numba.manager.dispatch import dispatch +from motrix_env_core.manager.dispatch import dispatch +from motrix_env_core.math.quaternion_array import rotate_inverse as rotate_inverse_batch @configclass(kw_only=True) @@ -17,10 +18,9 @@ class _BallBalanceTerminationCfg(TerminationTermCfg): @dispatch -def bad_base_z_termination(ctx: ManagerContext, threshold: np.float32) -> bool: - base_pos = ctx.sim["robot_base_pos"] - error = base_pos[2] - ctx.metrics["base_z"][0] = error +def bad_base_z_termination(ctx: ManagerContext, threshold: float) -> torch.Tensor: + error = ctx.sim["robot_base_pos"][:, 2] + ctx.metrics["base_z"][:, 0] = error return error < threshold @@ -29,7 +29,6 @@ class BadBaseZTerminationCfg(_BallBalanceTerminationCfg): """Terminate when the base falls below its balanced on-ball height.""" def __call__(self, ctx) -> TerminationTerm: - del ctx return TerminationTerm( bad_base_z_termination, np.float32(self.threshold), @@ -38,12 +37,13 @@ def __call__(self, ctx) -> TerminationTerm: @dispatch -def bad_orientation_termination(ctx: ManagerContext, threshold: np.float32) -> bool: +def bad_orientation_termination(ctx: ManagerContext, threshold: float) -> torch.Tensor: base_quat = ctx.sim["robot_base_quat"] - gravity = np.empty(3, dtype=np.float32) - rotate_inverse(base_quat, (0.0, 0.0, -1.0), gravity) - tilt_sq = gravity[0] * gravity[0] + gravity[1] * gravity[1] - ctx.metrics["base_tilt"][0] = tilt_sq + gravity = base_quat.new_zeros(base_quat.shape[0], 3) + gravity[:, 2] = -1.0 + gravity = rotate_inverse_batch(base_quat, gravity) + tilt_sq = gravity[:, 0] * gravity[:, 0] + gravity[:, 1] * gravity[:, 1] + ctx.metrics["base_tilt"][:, 0] = tilt_sq return tilt_sq > threshold * threshold @@ -52,7 +52,6 @@ class BadOrientationTerminationCfg(_BallBalanceTerminationCfg): """Terminate when the base tilts farther than ``threshold`` from upright.""" def __call__(self, ctx) -> TerminationTerm: - del ctx return TerminationTerm( bad_orientation_termination, np.float32(self.threshold), @@ -61,15 +60,15 @@ def __call__(self, ctx) -> TerminationTerm: @dispatch -def ball_escaped_termination(ctx: ManagerContext, threshold: np.float32) -> bool: +def ball_escaped_termination(ctx: ManagerContext, threshold: float) -> torch.Tensor: feet_pos = ctx.sim["feet_pos"] ball_pos = ctx.sim["ball_pos"] - mid_x = 0.5 * (feet_pos[0, 0] + feet_pos[1, 0]) - mid_y = 0.5 * (feet_pos[0, 1] + feet_pos[1, 1]) - error_x = ball_pos[0] - mid_x - error_y = ball_pos[1] - mid_y + mid_x = 0.5 * (feet_pos[:, 0, 0] + feet_pos[:, 1, 0]) + mid_y = 0.5 * (feet_pos[:, 0, 1] + feet_pos[:, 1, 1]) + error_x = ball_pos[:, 0] - mid_x + error_y = ball_pos[:, 1] - mid_y distance_sq = error_x * error_x + error_y * error_y - ctx.metrics["ball_feet_distance"][0] = distance_sq + ctx.metrics["ball_feet_distance"][:, 0] = distance_sq return distance_sq > threshold * threshold @@ -78,7 +77,6 @@ class BallEscapedTerminationCfg(_BallBalanceTerminationCfg): """Terminate when the ball rolls out from under the feet.""" def __call__(self, ctx) -> TerminationTerm: - del ctx return TerminationTerm( ball_escaped_termination, np.float32(self.threshold), diff --git a/motrix_envs/src/motrix_envs/locomotion/humanoid/walk_manager_mdp/command.py b/motrix_envs/src/motrix_envs/locomotion/humanoid/walk_manager_mdp/command.py index 9d1c0c97..1f4060f9 100644 --- a/motrix_envs/src/motrix_envs/locomotion/humanoid/walk_manager_mdp/command.py +++ b/motrix_envs/src/motrix_envs/locomotion/humanoid/walk_manager_mdp/command.py @@ -6,7 +6,7 @@ import math import numpy as np -from numba import njit +import torch from motrix_env_core.config import configclass from motrix_env_core.manager import ( @@ -18,49 +18,71 @@ kernel_data, metric, ) -from motrix_env_core.numba.manager.commands import ResetContext -from motrix_env_core.numba.manager.dispatch import dispatch - - -@njit(inline="always") -def _lane_phase( - cmd: np.ndarray, - phase_out: np.ndarray, - sin_cos_out: np.ndarray, - phase_offset: np.ndarray, - steps: float, - phase_step: float, +from motrix_env_core.manager.commands import ResetContext +from motrix_env_core.manager.dispatch import dispatch + + +def _phase( + command: torch.Tensor, + phase_out: torch.Tensor, + sin_cos_out: torch.Tensor, + phase_offset: torch.Tensor, + steps: torch.Tensor, + phase_step: torch.Tensor, ) -> None: - """Refresh both lanes' phase clocks, pinning standing commands to ``pi``.""" + """Refresh both lanes' phase clocks over whole batches, pinning standing commands to ``pi``. + + The clock arithmetic runs in float64 (the kernel's ``math.pi`` + promotion) and is stored back as float32; standing commands pin both + lanes to ``pi`` per lane. + """ tau = 2.0 * math.pi - for i in range(2): - phase_out[i] = (steps * phase_step + phase_offset[i] + math.pi) % tau - math.pi - speed = math.sqrt(cmd[0] * cmd[0] + cmd[1] * cmd[1]) - if speed < 0.01 and abs(cmd[2]) < 0.01: - # Standing commands pin both lanes, matching the direct env's - # ``phase[stand] = pi``: the gait clock must freeze for both feet. - phase_out[:] = math.pi - for i in range(2): - sin_cos_out[i] = math.sin(phase_out[i]) - sin_cos_out[2 + i] = math.cos(phase_out[i]) - - -@njit(inline="always") -def _lane_resample_phase_offset(ctx: ManagerContext, phase_offset: np.ndarray) -> None: - first = ctx.rand.uniform_range(np.float32(-math.pi), np.float32(math.pi)) - phase_offset[0] = first - phase_offset[1] = (first + 2.0 * math.pi) % (2.0 * math.pi) - math.pi - - -@njit(inline="always") -def _lane_resample_command( - ctx: ManagerContext, commands: np.ndarray, low: np.ndarray, high: np.ndarray, stand_prob: np.float32 + steps = steps.to(torch.float64)[:, None] + phase_step = phase_step.to(torch.float64)[:, None] + raw = (steps * phase_step + phase_offset.to(torch.float64) + math.pi) % tau - math.pi + speed = torch.sqrt(command[:, 0] * command[:, 0] + command[:, 1] * command[:, 1]) + stand = (speed.to(torch.float64) < 0.01) & (command[:, 2].abs().to(torch.float64) < 0.01) + phase = raw.to(torch.float32) + phase = torch.where(stand[:, None], torch.full_like(phase, np.float32(math.pi)), phase) + sin = torch.sin(phase.to(torch.float64)).to(torch.float32) + cos = torch.cos(phase.to(torch.float64)).to(torch.float32) + sin_cos_out[:, 0] = sin[:, 0] + sin_cos_out[:, 1] = sin[:, 1] + sin_cos_out[:, 2] = cos[:, 0] + sin_cos_out[:, 3] = cos[:, 1] + phase_out.copy_(phase) + + +def _uniform_range(u: torch.Tensor, low, high) -> torch.Tensor: + """Float32 ``low + (high - low) * 0.5 * (u + 1)`` matching the kernel convention.""" + return low + (high - low) * np.float32(0.5) * (u + np.float32(1.0)) + + +def _resample_command( + ctx: ManagerContext, + commands: torch.Tensor, + low: torch.Tensor, + high: torch.Tensor, + stand_prob, + resample: torch.Tensor, ) -> None: - rand = ctx.rand - for index in range(3): - commands[index] = rand.uniform_range(low[index], high[index]) - if (rand.next_uniform() + 1.0) * 0.5 < stand_prob: - commands[:] = 0.0 + """Resample velocity commands on ``resample`` lanes, preserving other lanes' RNG rows. + + The draw counts and per-lane order are fixed: one 3-draw uniform batch + then one stand-probability draw; lanes that do not resample keep their + pre-draw RNG state. + """ + if not bool(resample.any()): + return + previous_state = ctx.rand.state.clone() + u = ctx.rand.next_uniform(3) + sampled = _uniform_range(u, low, high) + stand_u = ctx.rand.next_uniform(1)[:, 0] + stand = (stand_u.to(torch.float64) + 1.0) * 0.5 < float(stand_prob) + sampled = torch.where(stand[:, None], torch.zeros_like(sampled), sampled) + keep_state = ~resample[:, None] + ctx.rand.state.copy_(torch.where(keep_state, previous_state, ctx.rand.state)) + commands.copy_(torch.where(resample[:, None], sampled, commands)) @kernel_data @@ -129,40 +151,63 @@ class WalkCommand(CommandTerm): phase_step: np.ndarray steps: np.ndarray = metric(name="command_steps", dtype=np.float32) + @staticmethod @dispatch - def update(self, ctx: ManagerContext) -> None: - _lane_phase( - self.command, - self.phase, - self.sin_cos, - self.phase_offset, - self.steps[0], - self.phase_step[0], + def update(receiver, ctx: ManagerContext) -> None: + _phase( + receiver.command, + receiver.phase, + receiver.sin_cos, + receiver.phase_offset, + receiver.steps[:, 0], + receiver.phase_step[:, 0], ) + @staticmethod @dispatch - def advance(self, ctx: ManagerContext) -> None: - self.steps[0] += 1.0 - if self.steps[0] % self.resample_steps == 0.0: - _lane_resample_command(ctx, self.command, self.vel_limit_low, self.vel_limit_high, self.stand_prob) + def advance(receiver, ctx: ManagerContext) -> None: + steps = receiver.steps + steps[:, 0] += np.float32(1.0) + resample = torch.remainder(steps[:, 0], float(receiver.resample_steps)) == 0.0 + _resample_command( + ctx, + receiver.command, + receiver.vel_limit_low, + receiver.vel_limit_high, + receiver.stand_prob, + resample, + ) + @staticmethod @dispatch - def reset_env(self, ctx: ManagerContext) -> None: - self.steps[0] = 0.0 - if self.gait_period_width > np.float32(0.0): - period = float(self.gait_period_base) + ctx.rand.uniform_range( - -self.gait_period_width, self.gait_period_width - ) - self.phase_step[0] = np.float32(2.0 * math.pi * ctx.dt / max(period, 0.1)) - _lane_resample_phase_offset(ctx, self.phase_offset) - _lane_resample_command(ctx, self.command, self.vel_limit_low, self.vel_limit_high, self.stand_prob) - _lane_phase( - self.command, - self.phase, - self.sin_cos, - self.phase_offset, - self.steps[0], - self.phase_step[0], + def reset_env(receiver, ctx: ManagerContext) -> None: + receiver.steps[:, 0] = 0.0 + width = float(receiver.gait_period_width) + if width > 0.0: + u = ctx.rand.next_uniform(1)[:, 0] + jitter = _uniform_range(u, np.float32(-width), np.float32(width)) + period = float(receiver.gait_period_base) + jitter.to(torch.float64) + receiver.phase_step[:, 0] = (2.0 * math.pi * ctx.dt / period.clamp(min=0.1)).to(torch.float32) + first = _uniform_range(ctx.rand.next_uniform(1)[:, 0], np.float32(-math.pi), np.float32(math.pi)) + offsets = torch.empty_like(receiver.phase_offset) + offsets[:, 0] = first + offsets[:, 1] = ((first.to(torch.float64) + 2.0 * math.pi) % (2.0 * math.pi) - math.pi).to(torch.float32) + receiver.phase_offset.copy_(offsets) + _resample_command( + ctx, + receiver.command, + receiver.vel_limit_low, + receiver.vel_limit_high, + receiver.stand_prob, + torch.ones_like(first, dtype=torch.bool), + ) + _phase( + receiver.command, + receiver.phase, + receiver.sin_cos, + receiver.phase_offset, + receiver.steps[:, 0], + receiver.phase_step[:, 0], ) def reset(self, ctx: ResetContext) -> None: diff --git a/motrix_envs/src/motrix_envs/locomotion/humanoid/walk_manager_mdp/observations.py b/motrix_envs/src/motrix_envs/locomotion/humanoid/walk_manager_mdp/observations.py index 28417197..db958408 100644 --- a/motrix_envs/src/motrix_envs/locomotion/humanoid/walk_manager_mdp/observations.py +++ b/motrix_envs/src/motrix_envs/locomotion/humanoid/walk_manager_mdp/observations.py @@ -4,19 +4,18 @@ """Observation terms for the humanoid velocity-tracking task.""" import numpy as np +import torch from motrix_env_core.config import configclass -from motrix_env_core.numba.manager.context import BuildContext, ManagerContext -from motrix_env_core.numba.manager.dispatch import dispatch -from motrix_env_core.numba.manager.observations import ObservationTermCfg, ObsTerm -from motrix_envs.locomotion.humanoid.walk_manager_mdp.command import WalkCommand +from motrix_env_core.manager.context import BuildContext, ManagerContext +from motrix_env_core.manager.dispatch import dispatch +from motrix_env_core.manager.observations import ObservationTermCfg, ObsTerm @dispatch -def gait_phase_obs(ctx: ManagerContext, out: np.ndarray, offset: np.int64, size: np.int64) -> None: - walk: WalkCommand = ctx.commands["walk"] - for index in range(size): - out[index] = walk.sin_cos[offset + index] +def gait_phase_obs(ctx: ManagerContext, out: torch.Tensor, offset: np.int64, size: np.int64) -> None: + walk = ctx.commands["walk"] + out.copy_(walk.sin_cos[:, offset : offset + size]) @configclass(kw_only=True) diff --git a/motrix_envs/src/motrix_envs/locomotion/humanoid/walk_manager_mdp/reset.py b/motrix_envs/src/motrix_envs/locomotion/humanoid/walk_manager_mdp/reset.py index 0e19b3b7..9cdc24ee 100644 --- a/motrix_envs/src/motrix_envs/locomotion/humanoid/walk_manager_mdp/reset.py +++ b/motrix_envs/src/motrix_envs/locomotion/humanoid/walk_manager_mdp/reset.py @@ -3,10 +3,8 @@ """Simulator reset terms for the humanoid velocity-tracking task.""" -import math - import numpy as np -from numba import njit +import torch from motrix_env_core.config import configclass from motrix_env_core.manager import ( @@ -16,10 +14,10 @@ SharedArray, kernel_data, ) +from motrix_env_core.manager.context import BuildContext +from motrix_env_core.manager.dispatch import dispatch +from motrix_env_core.manager.kernel_data.map import Map from motrix_env_core.mdp.terrain import HeightFieldGrid, heightfield_lookup -from motrix_env_core.numba.kernel_data.map import Map -from motrix_env_core.numba.manager.context import BuildContext -from motrix_env_core.numba.manager.dispatch import dispatch from motrix_env_core.sim.write import ( ActuatorDampingWrite, ActuatorKpWrite, @@ -33,6 +31,7 @@ LinkComWrite, LinkMassWrite, ) +from motrix_envs.locomotion.humanoid.walk_manager_mdp.command import _uniform_range from motrix_envs.locomotion.humanoid.walk_manager_mdp.randomization import WalkRandomizationCfg @@ -94,97 +93,90 @@ class WalkResetParams: com_noise: np.ndarray -@njit(inline="always") -def _sample_randomized_dynamics(ctx: ManagerContext, sim_writes: Map[np.ndarray], params: WalkResetParams) -> None: - """Sample one lane's randomized dynamics into the reset write buffers. +def _sample_spawn_xy(ctx: ManagerContext, params: WalkResetParams) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Sample world xy and lifted z, mirroring the per-lane spawn draws.""" + pose = params.init_pose + low = np.float32(-float(params.spawn_range)) + high = np.float32(float(params.spawn_range)) + x = _uniform_range(ctx.rand.next_uniform(1)[:, 0], low, high) + y = _uniform_range(ctx.rand.next_uniform(1)[:, 0], low, high) + offsets = torch.tensor([(i - 1) * 0.15 for i in range(3)], dtype=torch.float64) + xs = x.to(torch.float64)[:, None] + offsets[None, :] + ys = y.to(torch.float64)[:, None] + offsets[None, :] + xs9 = xs[:, :, None].expand(-1, -1, 3).reshape(-1) + ys9 = ys[:, None, :].expand(-1, 3, -1).reshape(-1) + samples = heightfield_lookup(params.heightfield, xs9, ys9).view(x.shape[0], 9) + z = pose[2].to(torch.float64) + samples.max(dim=1).values + return x, y, z + - Row-scoped buffers: ``(N, A)``-style writes arrive as 1-D rows, while - ``(N, 1, 3)``-style writes keep their leading lane dimension. - """ +def _sample_randomized_dynamics(ctx: ManagerContext, sim_writes: Map[torch.Tensor], params: WalkResetParams) -> None: + """Sample randomized dynamics into the reset write buffers on compact rows.""" rand = ctx.rand - if params.kp_range[1] > params.kp_range[0]: - kp = sim_writes["kp"] - for i in range(kp.shape[0]): - kp[i] = params.kp_default[i] * rand.uniform_range(params.kp_range[0], params.kp_range[1]) - if params.damping_range[1] > params.damping_range[0]: - damping = sim_writes["damping"] - for i in range(damping.shape[0]): - damping[i] = params.damping_default[i] * rand.uniform_range( - params.damping_range[0], params.damping_range[1] - ) - if params.friction_range[1] > params.friction_range[0]: - ratio = rand.uniform_range(params.friction_range[0], params.friction_range[1]) - friction = sim_writes["friction"] - for i in range(3): - friction[0, i] = params.friction_default[i] * ratio - mass_randomized = params.mass_scale_range[1] > params.mass_scale_range[0] - base_mass_randomized = params.base_mass_off_range[1] > params.base_mass_off_range[0] + if float(params.kp_range[1]) > float(params.kp_range[0]): + u = rand.next_uniform(int(params.kp_default.numel())) + scale = _uniform_range(u, params.kp_range[0], params.kp_range[1]) + sim_writes["kp"][:] = params.kp_default * scale + if float(params.damping_range[1]) > float(params.damping_range[0]): + u = rand.next_uniform(int(params.damping_default.numel())) + scale = _uniform_range(u, params.damping_range[0], params.damping_range[1]) + sim_writes["damping"][:] = params.damping_default * scale + if float(params.friction_range[1]) > float(params.friction_range[0]): + ratio = _uniform_range(rand.next_uniform(1), params.friction_range[0], params.friction_range[1]) + sim_writes["friction"][:, 0, :] = params.friction_default * ratio + mass_randomized = float(params.mass_scale_range[1]) > float(params.mass_scale_range[0]) + base_mass_randomized = float(params.base_mass_off_range[1]) > float(params.base_mass_off_range[0]) if mass_randomized or base_mass_randomized: - mass = sim_writes["mass"] - for i in range(mass.shape[0]): - if i == params.base_link_index: - mass[i] = params.base_mass_default + rand.uniform_range( - params.base_mass_off_range[0], params.base_mass_off_range[1] - ) - else: - mass[i] = params.link_mass_default[i] * rand.uniform_range( - params.mass_scale_range[0], params.mass_scale_range[1] - ) - if params.com_noise[0] > 0.0 or params.com_noise[1] > 0.0 or params.com_noise[2] > 0.0: - com = sim_writes["com"] - com[0, 0] = params.base_com_default[0] + rand.uniform_range(-params.com_noise[0], params.com_noise[0]) - com[0, 1] = params.base_com_default[1] + rand.uniform_range(-params.com_noise[1], params.com_noise[1]) - com[0, 2] = params.base_com_default[2] + rand.uniform_range(-params.com_noise[2], params.com_noise[2]) - - -@njit(inline="always") -def _spawn_ground_height(grid: HeightFieldGrid, x: float, y: float) -> float: - """Highest terrain height over the +/-0.15 m 9-point spawn grid.""" - best = -math.inf - for i in range(3): - for j in range(3): - best = max(best, heightfield_lookup(grid, x + (i - 1) * 0.15, y + (j - 1) * 0.15)) - return best + u = rand.next_uniform(int(params.link_mass_default.numel())) + offset = _uniform_range(u, params.base_mass_off_range[0], params.base_mass_off_range[1]) + scale = _uniform_range(u, params.mass_scale_range[0], params.mass_scale_range[1]) + is_base = torch.arange(int(params.link_mass_default.numel()), device=u.device) == int(params.base_link_index) + values = torch.where( + is_base[None, :], + params.base_mass_default + offset, + params.link_mass_default * scale, + ) + sim_writes["mass"][:] = values + noise = params.com_noise + if float(noise[0]) > 0.0 or float(noise[1]) > 0.0 or float(noise[2]) > 0.0: + u = rand.next_uniform(3) + offsets = _uniform_range(u, -noise, noise) + sim_writes["com"][:, 0, :] = params.base_com_default + offsets -@njit(inline="always") def _write_spawn_state( - ctx: ManagerContext, sim_writes: Map[np.ndarray], params: WalkResetParams, x: float, y: float, z: float + sim_writes: Map[torch.Tensor], params: WalkResetParams, x: torch.Tensor, y: torch.Tensor, z: torch.Tensor ) -> None: - # Homogeneous float32 tuple: z picks up float64 from the terrain-height add. - sim_writes["position"][0, :3] = np.float32(x), np.float32(y), np.float32(z) - sim_writes["rotation"][0] = params.init_pose[3:] - sim_writes["linear_velocity"][0, :] = 0.0 - sim_writes["angular_velocity"][0, :] = 0.0 + position = sim_writes["position"][:, 0, :3] + position[:, 0] = x.to(torch.float32) + position[:, 1] = y.to(torch.float32) + position[:, 2] = z.to(torch.float32) + sim_writes["rotation"][:, 0] = params.init_pose[3:] + sim_writes["linear_velocity"][:, 0] = 0.0 + sim_writes["angular_velocity"][:, 0] = 0.0 sim_writes["joints_position"][:] = params.default_joint_angles sim_writes["joints_velocity"][:] = 0.0 @dispatch -def reset_walk_state(ctx: ManagerContext, sim_writes: Map[np.ndarray], params: WalkResetParams) -> None: +def reset_walk_state(ctx: ManagerContext, sim_writes: Map[torch.Tensor], params: WalkResetParams) -> None: """Spawn-state-only reset for configs without dynamics randomization.""" pose = params.init_pose x, y, z = pose[0], pose[1], pose[2] - if params.spawn_range > 0.0: - rand = ctx.rand - x = rand.uniform_range(-params.spawn_range, params.spawn_range) - y = rand.uniform_range(-params.spawn_range, params.spawn_range) - z += _spawn_ground_height(params.heightfield, x, y) - _write_spawn_state(ctx, sim_writes, params, x, y, z) + if float(params.spawn_range) > 0.0: + x, y, z = _sample_spawn_xy(ctx, params) + _write_spawn_state(sim_writes, params, x, y, z) @dispatch -def reset_walk_state_randomized(ctx: ManagerContext, sim_writes: Map[np.ndarray], params: WalkResetParams) -> None: +def reset_walk_state_randomized(ctx: ManagerContext, sim_writes: Map[torch.Tensor], params: WalkResetParams) -> None: """Reset plus in-kernel dynamics randomization (kp/damping/friction/mass/com).""" pose = params.init_pose x, y, z = pose[0], pose[1], pose[2] - if params.spawn_range > 0.0: - rand = ctx.rand - x = rand.uniform_range(-params.spawn_range, params.spawn_range) - y = rand.uniform_range(-params.spawn_range, params.spawn_range) - z += _spawn_ground_height(params.heightfield, x, y) + if float(params.spawn_range) > 0.0: + x, y, z = _sample_spawn_xy(ctx, params) _sample_randomized_dynamics(ctx, sim_writes, params) - _write_spawn_state(ctx, sim_writes, params, x, y, z) + _write_spawn_state(sim_writes, params, x, y, z) @configclass(kw_only=True) @@ -217,8 +209,8 @@ def __call__(self, ctx: BuildContext) -> ResetTerm: if not joint_names or len(set(joint_names)) != len(joint_names): raise ValueError("humanoid walk requires unique actuator target joints") # The dof_pos query (and every consumer aligned to it) uses the - # key-pose declaration order; require it to match the body's - # joint-DOF order so per-joint arrays stay aligned. + # key-pose declaration order; require it to match the body's joint-DOF + # order so per-joint arrays stay aligned. key_pose_names = tuple(robot.resolve_name(name) for name in robot.key_pose.joint_names) if joint_names != key_pose_names: raise ValueError( @@ -249,8 +241,8 @@ def __call__(self, ctx: BuildContext) -> ResetTerm: actuator_names = tuple(spec.name for spec in ctx.model.actuators) link_names = body.link_names base_index = link_names.index(body.base_link_name) - # Keep the write-map schema stable for Numba: Map keys are part of - # the compiled type even when a particular range is degenerate. + # Keep the write-map schema stable: Map keys are part of the + # compiled type even when a particular range is degenerate. writes.update( kp=ActuatorKpWrite(actuator_names), damping=ActuatorDampingWrite(actuator_names), diff --git a/motrix_envs/src/motrix_envs/locomotion/humanoid/walk_manager_mdp/rewards.py b/motrix_envs/src/motrix_envs/locomotion/humanoid/walk_manager_mdp/rewards.py index b8a7370c..628e58fd 100644 --- a/motrix_envs/src/motrix_envs/locomotion/humanoid/walk_manager_mdp/rewards.py +++ b/motrix_envs/src/motrix_envs/locomotion/humanoid/walk_manager_mdp/rewards.py @@ -6,7 +6,7 @@ import math import numpy as np -from numba import njit +import torch from motrix_env_core.config import configclass from motrix_env_core.manager import ( @@ -14,9 +14,10 @@ RewardTerm, RewardTermCfg, ) +from motrix_env_core.manager.context import BuildContext +from motrix_env_core.manager.dispatch import dispatch +from motrix_env_core.math.quaternion_array import rotate_inverse_components from motrix_env_core.mdp.terrain import HeightFieldGrid, heightfield_lookup -from motrix_env_core.numba.manager.dispatch import dispatch -from motrix_env_core.numba.math.quaternion import rotate_inverse_components from motrix_env_core.sim import ( BatchLinkPositionQuery, BatchLinkQuaternionQuery, @@ -25,27 +26,23 @@ LinkQuaternionQuery, SitePositionQuery, ) -from motrix_envs.locomotion.humanoid.walk_manager_mdp.command import WalkCommand from motrix_envs.locomotion.humanoid.walk_manager_mdp.terrain import ground_height_grid -@njit(inline="always") -def _expected_foot_height(phi: float, swing_height: float) -> float: - """Expected foot height from gait phase using the direct env's Bezier profile.""" - x = (phi + math.pi) / (2.0 * math.pi) - - def bezier(y_start, y_end, t): - return y_start + (y_end - y_start) * (t**3 + 3.0 * (t**2 * (1.0 - t))) - - if x <= 0.5: - return bezier(0.0, swing_height, 2.0 * x) - return bezier(swing_height, 0.0, 2.0 * x - 1.0) +def _expected_foot_height(phi: torch.Tensor, swing_height: float) -> torch.Tensor: + """Expected foot height from gait phase using the direct env's Bezier profile (float64, ``(N,)``).""" + x = (phi.to(torch.float64) + math.pi) / (2.0 * math.pi) + swing = float(swing_height) + t = torch.where(x <= 0.5, 2.0 * x, 2.0 * x - 1.0) + y_start = torch.where(x <= 0.5, 0.0, swing) + y_end = torch.where(x <= 0.5, swing, 0.0) + return y_start + (y_end - y_start) * (t**3 + 3.0 * (t**2 * (1.0 - t))) @dispatch -def penalty_ang_vel_xy_reward(ctx: ManagerContext, base_quat: np.ndarray, base_ang_vel: np.ndarray) -> float: +def penalty_ang_vel_xy_reward(ctx: ManagerContext, base_quat: torch.Tensor, base_ang_vel: torch.Tensor) -> torch.Tensor: vx, vy, _ = rotate_inverse_components(base_quat, base_ang_vel) - walk: WalkCommand = ctx.commands["walk"] + walk = ctx.commands["walk"] return (vx * vx + vy * vy) * walk.penalty_scale[0] @@ -61,9 +58,11 @@ def __call__(self, ctx) -> RewardTerm: @dispatch -def penalty_orientation_reward(ctx: ManagerContext, base_quat: np.ndarray) -> float: - gx, gy, _ = rotate_inverse_components(base_quat, (0.0, 0.0, -1.0)) - walk: WalkCommand = ctx.commands["walk"] +def penalty_orientation_reward(ctx: ManagerContext, base_quat: torch.Tensor) -> torch.Tensor: + gravity = base_quat.new_zeros(base_quat.shape[0], 3) + gravity[:, 2] = -1.0 + gx, gy, _ = rotate_inverse_components(base_quat, gravity) + walk = ctx.commands["walk"] return (gx * gx + gy * gy) * walk.penalty_scale[0] @@ -78,38 +77,38 @@ def __call__(self, ctx) -> RewardTerm: @dispatch -def penalty_action_rate_reward(ctx: ManagerContext) -> float: +def penalty_action_rate_reward(ctx: ManagerContext) -> torch.Tensor: action = ctx.actions["joint_position"] delta = action.current() - action.previous() - walk: WalkCommand = ctx.commands["walk"] - return float(np.dot(delta, delta)) * walk.penalty_scale[0] + walk = ctx.commands["walk"] + return (delta * delta).sum(dim=1) * walk.penalty_scale[0] @configclass(kw_only=True) class PenaltyActionRateRewardCfg(RewardTermCfg): def __call__(self, ctx) -> RewardTerm: - del ctx return RewardTerm(penalty_action_rate_reward) @dispatch def feet_phase_reward( ctx: ManagerContext, - swing_height: np.float32, - feet_phase_sigma: np.float32, + swing_height: float, + feet_phase_sigma: float, heightfield: HeightFieldGrid, - sole_l: np.ndarray, - sole_r: np.ndarray, -) -> float: - walk: WalkCommand = ctx.commands["walk"] - error = 0.0 + sole_l: torch.Tensor, + sole_r: torch.Tensor, +) -> torch.Tensor: + walk = ctx.commands["walk"] + soles = (sole_l, sole_r) + error = None for foot in range(2): - expected = _expected_foot_height(walk.phase[foot], swing_height) - sole = sole_l if foot == 0 else sole_r - ground_z = heightfield_lookup(heightfield, sole[0], sole[1]) - delta = sole[2] - ground_z - expected - error += delta * delta - return math.exp(-error / feet_phase_sigma) + expected = _expected_foot_height(walk.phase[:, foot], swing_height) + sole = soles[foot] + ground_z = heightfield_lookup(heightfield, sole[:, 0], sole[:, 1]) + delta = sole[:, 2].to(torch.float64) - ground_z - expected + error = delta * delta if error is None else error + delta * delta + return torch.exp(-error / float(feet_phase_sigma)).to(torch.float32) @configclass(kw_only=True) @@ -141,11 +140,11 @@ def __call__(self, ctx) -> RewardTerm: @dispatch def pose_reward( - ctx: ManagerContext, default_joint_angles: np.ndarray, pose_weights: np.ndarray, joint_pos: np.ndarray -) -> float: + ctx: ManagerContext, default_joint_angles: torch.Tensor, pose_weights: torch.Tensor, joint_pos: torch.Tensor +) -> torch.Tensor: delta = joint_pos - default_joint_angles - walk: WalkCommand = ctx.commands["walk"] - return float(np.dot(delta * delta, pose_weights)) * walk.penalty_scale[0] + walk = ctx.commands["walk"] + return ((delta * delta) * pose_weights).sum(dim=1) * walk.penalty_scale[0] @configclass(kw_only=True) @@ -178,22 +177,19 @@ def __call__(self, ctx) -> RewardTerm: @dispatch def penalty_close_feet_xy_reward( - ctx: ManagerContext, close_feet_threshold: np.float32, base_quat: np.ndarray, foot_pos: np.ndarray -) -> float: - left = foot_pos[0] - right = foot_pos[1] + ctx: ManagerContext, close_feet_threshold: float, base_quat: torch.Tensor, foot_pos: torch.Tensor +) -> torch.Tensor: # Lateral foot separation in the base frame: rotate the world-frame foot # delta's XY components (z forced to zero — vertical separation is # intentionally ignored) back into the base frame and take its y # magnitude. Rotating the world x-axis instead (as a yaw proxy) flips the # yaw sign under the inverse rotation and misjudges staggered feet as # crossing. - _, lateral, _ = rotate_inverse_components(base_quat, (left[0] - right[0], left[1] - right[1], 0.0)) - distance = abs(lateral) - walk: WalkCommand = ctx.commands["walk"] - if distance < close_feet_threshold: - return 1.0 * walk.penalty_scale[0] - return 0.0 + delta = foot_pos.new_zeros(foot_pos.shape[0], 3) + delta[:, :2] = foot_pos[:, 0, :2] - foot_pos[:, 1, :2] + _, lateral, _ = rotate_inverse_components(base_quat, delta) + walk = ctx.commands["walk"] + return torch.where(lateral.abs() < close_feet_threshold, walk.penalty_scale[0], torch.zeros_like(lateral)) @configclass(kw_only=True) @@ -213,17 +209,28 @@ def __call__(self, ctx) -> RewardTerm: @dispatch -def penalty_feet_ori_reward(ctx: ManagerContext, default_foot_gravity: np.ndarray, foot_quat: np.ndarray) -> float: - total = 0.0 - for foot in range(2): - gx, gy, gz = rotate_inverse_components(foot_quat[foot], (0.0, 0.0, -1.0)) - reference = default_foot_gravity[foot] - # |cross(foot_gravity, reference)| - cx = gy * reference[2] - gz * reference[1] - cy = gz * reference[0] - gx * reference[2] - cz = gx * reference[1] - gy * reference[0] - total += math.sqrt(cx * cx + cy * cy + cz * cz) - walk: WalkCommand = ctx.commands["walk"] +def penalty_feet_ori_reward( + ctx: ManagerContext, default_foot_gravity: torch.Tensor, foot_quat: torch.Tensor +) -> torch.Tensor: + # foot_quat is (N, F, 4); rotate gravity by each foot's inverse quaternion. + num_feet = foot_quat.shape[1] + gravity = foot_quat.new_zeros(*foot_quat.shape[:2], 3) + gravity[..., 2] = -1.0 + gx, gy, gz = rotate_inverse_components(foot_quat.reshape(-1, 4), gravity.reshape(-1, 3)) + reference = torch.as_tensor(np.asarray(default_foot_gravity)) + total = None + for foot in range(num_feet): + # Flat layout is lane-major (lane * F + foot): strided slices pick one foot. + fx = gx[foot::num_feet] + fy = gy[foot::num_feet] + fz = gz[foot::num_feet] + ref = reference[foot] + cx = fy * ref[2] - fz * ref[1] + cy = fz * ref[0] - fx * ref[2] + cz = fx * ref[1] - fy * ref[0] + norm = torch.sqrt(cx * cx + cy * cy + cz * cz) + total = norm if total is None else total + norm + walk = ctx.commands["walk"] return total * walk.penalty_scale[0] @@ -231,18 +238,20 @@ def penalty_feet_ori_reward(ctx: ManagerContext, default_foot_gravity: np.ndarra class PenaltyFeetOriRewardCfg(RewardTermCfg): def __call__(self, ctx) -> RewardTerm: robot = ctx.cfg.scene.objs.robot - body = ctx.model.bodies["robot"] # Default foot gravity at the init key pose, from BodyModel's # compile-time init-pose FK snapshot. - gravity_vec = (0.0, 0.0, -1.0) - default_foot_gravity = np.stack( - [ - rotate_inverse_components(body.init_link_quats[body.link_names.index(link)], gravity_vec) - for link in robot.resolved_foot_link_names - ] - ).astype(np.float32) + default_foot_gravity = _default_foot_gravity(ctx, robot) return RewardTerm( penalty_feet_ori_reward, default_foot_gravity, BatchLinkQuaternionQuery(links=robot.resolved_foot_link_names), ) + + +def _default_foot_gravity(ctx: BuildContext, robot) -> np.ndarray: + """Rotate world gravity into each default foot frame (host-side numpy).""" + body = ctx.model.bodies["robot"] + quats = np.stack([body.init_link_quats[body.link_names.index(link)] for link in robot.resolved_foot_link_names]) + gravity = np.broadcast_to(np.array((0.0, 0.0, -1.0)), quats.shape[:-1] + (3,)) + gx, gy, gz = rotate_inverse_components(quats, gravity) + return np.stack((gx, gy, gz), axis=-1).astype(np.float32) diff --git a/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/command.py b/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/command.py index 00dd1618..ef72ac33 100644 --- a/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/command.py +++ b/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/command.py @@ -3,10 +3,8 @@ """Command terms for manager-based whole-body tracking.""" -import math - import numpy as np -from numba import njit +import torch from omegaconf import MISSING from motrix_env_core.config import configclass @@ -20,38 +18,40 @@ kernel_data, metric, ) -from motrix_env_core.manager.math.quaternion import inverse as quat_inverse -from motrix_env_core.manager.math.quaternion import mul as quat_mul -from motrix_env_core.manager.math.quaternion import rotate_vector -from motrix_env_core.numba.manager.commands import ResetContext -from motrix_env_core.numba.manager.dispatch import dispatch +from motrix_env_core.manager.commands import ResetContext +from motrix_env_core.manager.dispatch import dispatch +from motrix_env_core.math.quaternion_array import align_targets +from motrix_env_core.torch_manager.rand import draw_uniform from motrix_envs.motion import MotionLibrary, WbtMotionClip -@njit(inline="always") def _sample_motion_step( - rand, - sampling_cdf, - num_frames: np.int64, - start_at_timestep_zero_prob: np.float32, -): - """Draw one start frame from the adaptive-bin CDF (or uniformly when disabled).""" - unit = (rand.next_uniform() + np.float32(1.0)) * np.float32(0.5) - if sampling_cdf.size == 0: - step_value = unit * np.float32(num_frames - 1) + draws: torch.Tensor, + sampling_cdf: torch.Tensor, + num_frames: int, + start_zero_prob: torch.Tensor, +) -> torch.Tensor: + """Batched start-frame sampling from pre-drawn uniform samples. + + ``draws`` is ``(N, 2)`` without adaptive sampling or ``(N, 3)`` with it — + one unit draw, optionally one within-bin draw, and one start-at-zero draw + per lane, matching the fixed draw order. + """ + unit = (draws[:, :1] + 1.0) * 0.5 + adaptive = sampling_cdf.numel() != 0 + if adaptive: + # ``sum(cdf < unit)`` is search-left: equality selects the equal bin, + # matching the per-lane linear scan that stops at ``unit <= cdf[bin]``. + bin_id = (sampling_cdf[None, :] < unit).sum(dim=1, dtype=torch.int64) + bin_unit = (draws[:, 1:2] + 1.0) * 0.5 + phase = (bin_id.to(torch.float32) + bin_unit[:, 0]) / sampling_cdf.numel() + step_value = phase * num_frames else: - bin_id = np.int64(0) - while bin_id + 1 < sampling_cdf.size and unit > sampling_cdf[bin_id]: - bin_id += 1 - bin_unit = (rand.next_uniform() + np.float32(1.0)) * np.float32(0.5) - phase = (np.float32(bin_id) + bin_unit) / np.float32(sampling_cdf.size) - step_value = phase * np.float32(num_frames) - max_value = np.float32(num_frames - 2) - step_value = min(max(step_value, np.float32(0.0)), max_value) - step = np.int64(step_value) - if (rand.next_uniform() + np.float32(1.0)) * np.float32(0.5) < start_at_timestep_zero_prob: - step = np.int64(0) - return step + step_value = unit[:, 0] * (num_frames - 1) + step_value = torch.clamp(step_value, min=0.0, max=float(num_frames - 2)) + steps = step_value.to(torch.int64) + zero_unit = (draws[:, -1:] + 1.0) * 0.5 + return torch.where(zero_unit[:, 0] < start_zero_prob, torch.zeros_like(steps), steps) @kernel_data @@ -59,8 +59,8 @@ class WbtMotionCommand(CommandTerm): """Drive each environment through a shared whole-body reference-motion clip. Clip arrays are shared and indexed by the per-environment ``steps`` vector. - :meth:`evaluate` aligns the selected frame to the simulated robot's current - reference body, the fused-kernel hooks :meth:`advance` / :meth:`reset_env` + :meth:`update` aligns the selected frame to the simulated robot's current + reference body, the kernel hooks :meth:`advance` / :meth:`reset_env` advance and resample the per-lane frame index, and host-side lifecycle methods maintain adaptive-sampling statistics and the sampling distribution. @@ -119,57 +119,6 @@ class WbtMotionCommand(CommandTerm): steps: np.ndarray = metric(name="motion_step", dtype=np.float32) clip_ended: np.ndarray = metric() - @dispatch - def update(self, ctx: ManagerContext) -> None: - """Align the selected motion frame to one simulated robot lane.""" - # Select the current motion frame for this environment lane. - step = self.steps[0] - motion_reference_body_pos_w = self.clip.reference_body_pos_w[step] - motion_reference_body_quat_w = self.clip.reference_body_quat_w[step] - motion_tracked_bodies_pos_w = self.clip.tracked_bodies_pos_w[step] - motion_tracked_bodies_quat_w = self.clip.tracked_bodies_quat_w[step] - # Anchor the motion to the current simulated pose of the configured - # reference body, rather than to the clip's absolute world transform. - reference_index = self.reference_index - robot_reference_body_pos_w = ctx.sim["tracked_body_pos"][reference_index] - robot_reference_body_quat_w = ctx.sim["tracked_body_quat"][reference_index] - out_body_pos = self.target_body_position_relative - out_body_quat = self.target_body_orientation_relative - - # Keep only the relative yaw between motion and robot. Roll and pitch - # remain those of each body in the source motion after composition. - delta_quat = out_body_quat[0] - quat_inverse(motion_reference_body_quat_w, delta_quat) - quat_mul(robot_reference_body_quat_w, delta_quat, delta_quat) - dx, dy, dz, dw = delta_quat - half_yaw = np.float32(0.5 * math.atan2(2.0 * (dw * dz + dx * dy), 1.0 - 2.0 * (dy * dy + dz * dz))) - yaw_z = math.sin(half_yaw) - yaw_w = math.cos(half_yaw) - height_delta = motion_reference_body_pos_w[2] - robot_reference_body_pos_w[2] - - # Apply the same planar rigid transform to every tracked body, then - # compose the yaw alignment with its source-motion orientation. - for body_id in range(motion_tracked_bodies_pos_w.shape[0]): - out_quat = out_body_quat[body_id] - out_quat[0] = 0.0 - out_quat[1] = 0.0 - out_quat[2] = yaw_z - out_quat[3] = yaw_w - out_pos = out_body_pos[body_id] - rotate_vector( - out_quat, - ( - motion_tracked_bodies_pos_w[body_id, 0] - motion_reference_body_pos_w[0], - motion_tracked_bodies_pos_w[body_id, 1] - motion_reference_body_pos_w[1], - motion_tracked_bodies_pos_w[body_id, 2] - motion_reference_body_pos_w[2], - ), - out_pos, - ) - out_pos[0] += robot_reference_body_pos_w[0] - out_pos[1] += robot_reference_body_pos_w[1] - out_pos[2] += robot_reference_body_pos_w[2] + height_delta - quat_mul(out_quat, motion_tracked_bodies_quat_w[body_id], out_quat) - def reset(self, ctx: ResetContext) -> None: """Update adaptive statistics and prepare sampling for selected environments.""" env_ids = ctx.env_ids @@ -213,82 +162,6 @@ def on_transition(self) -> None: ) self.adaptive_current_bin_failed_count.fill(0.0) - @dispatch - def reset_env(self, ctx: ManagerContext) -> None: - """Sample the starting frame for one reset environment lane.""" - if self.sequential_clips: - # Sequential playback keeps the timeline unbroken: every reset - # (episode terminations included) rematerializes the lane onto - # its current frame instead of resampling a new start. - return - num_frames = self.clip.joint_pos.shape[0] - self.steps[0] = _sample_motion_step( - ctx.rand, - self.sampling_cdf, - np.int64(num_frames), - self.start_at_timestep_zero_prob, - ) - # ``clip_ended`` is intentionally left untouched: advance recomputes it - # every transition, so a lane that just wrapped keeps its flag for the - # metrics/observation pass that follows rematerialization. - - @dispatch - def advance(self, ctx: ManagerContext) -> None: - """Advance one frame for the current environment lane. - - The lane leaves its clip when the step index exceeds the per-frame - ``frame_clip_end`` boundary of the clip it just tracked (the clip's - final frame is tracked one step before the index goes out of range). - Without ``sequential_clips``, ``hold_at_clip_end`` clamps the lane - onto its own clip's final frame and the wrap branch otherwise - resamples from the whole corpus. With ``sequential_clips``, the lane - crosses intermediate boundaries into the next clip's head frame and - the corpus end either loops back to the head or, with - ``hold_at_clip_end``, holds on the corpus final frame. - - The lowering already binds this lane's writable row view to - ``self.steps`` / ``self.clip_ended``; only the wrap branch needs ``ctx`` - to draw the replacement frame. - """ - num_frames = self.clip.joint_pos.shape[0] - clip_end = self.clip.frame_clip_end[self.steps[0]] - step = self.steps[0] + 1 - self.steps[0] = step - self.clip_ended[0] = step > clip_end - if self.hold_at_clip_end and not self.sequential_clips: - # Freeze on the lane's own clip final frame. - self.steps[0] = min(step, clip_end) - elif self.clip_ended[0]: - # The lane passed its clip's final frame. Crossing requests - # sim-only rematerialization: the reset pipeline reruns this - # lane's reset_env hook and the configured sim reset terms - # teleport the robot to the resulting frame. Episode bookkeeping - # and action-term state are untouched. - if self.sequential_clips: - # Clips are concatenated, so the frame past a clip boundary - # is the next clip's head. Intermediate boundaries cross in - # order; the corpus end loops back to the head, or holds on - # the final frame when hold_at_clip_end is set. reset_env - # keeps the chosen frame. - if step < num_frames: - self.steps[0] = step - ctx.sim_reset_requested[0] = True - elif self.hold_at_clip_end: - self.steps[0] = num_frames - 1 - else: - self.steps[0] = 0 - ctx.sim_reset_requested[0] = True - else: - # The inline resample keeps steps valid and consistently - # distributed between this kernel and the reset pipeline. - self.steps[0] = _sample_motion_step( - ctx.rand, - self.sampling_cdf, - np.int64(num_frames), - self.start_at_timestep_zero_prob, - ) - ctx.sim_reset_requested[0] = True - def _sampling_probabilities(self) -> np.ndarray: """Build normalized frame-bin probabilities from failure history.""" num_bins = self.adaptive_bin_failed_count.size @@ -307,6 +180,91 @@ def _sampling_probabilities(self) -> np.ndarray: return np.full((num_bins,), 1.0 / num_bins, dtype=np.float32) return (probability / total).astype(np.float32) + # -- Kernel hooks ---------------------------------------------------------- + # These static hooks receive a batched context (all tensors whole-batch) + # and a receiver namespace over this command's kernel-data fields. Per-lane + # state is written in place through the context tensors, so host-visible + # buffers stay authoritative. + + @staticmethod + @dispatch + def update(receiver, ctx: ManagerContext) -> None: + """Yaw-align the clip to each lane.""" + steps = receiver.steps[:, 0] + reference_index = receiver.reference_index + robot_position = ctx.sim["tracked_body_pos"][:, reference_index] + robot_quat = ctx.sim["tracked_body_quat"][:, reference_index] + target_position, target_quat = align_targets( + robot_position, + robot_quat, + receiver.clip.reference_body_pos_w[steps], + receiver.clip.reference_body_quat_w[steps], + receiver.clip.tracked_bodies_pos_w[steps], + receiver.clip.tracked_bodies_quat_w[steps], + ) + receiver.target_body_position_relative.copy_(target_position) + receiver.target_body_orientation_relative.copy_(target_quat) + + @staticmethod + @dispatch + def advance(receiver, ctx: ManagerContext) -> None: + """Advance every lane one frame.""" + steps = receiver.steps + clip_end = receiver.clip.frame_clip_end[steps[:, 0]] + num_frames = receiver.clip.joint_pos.shape[0] + advanced = steps + 1 + ended = advanced[:, 0] > clip_end + receiver.clip_ended[:, 0] = ended + if bool(receiver.sequential_clips): + if bool(receiver.hold_at_clip_end): + next_steps = torch.clamp(advanced, max=num_frames - 1) + crossing = ended & (advanced[:, 0] < num_frames) + else: + next_steps = torch.where(advanced >= num_frames, torch.zeros_like(advanced), advanced) + crossing = ended + steps.copy_(next_steps) + ctx.sim_reset_requested |= crossing[:, None] + return + if bool(receiver.hold_at_clip_end): + # Freeze each lane on its own clip's final frame; no reset request. + steps.copy_(torch.minimum(advanced, clip_end[:, None])) + return + # Non-sequential wrap: resample each ended lane inline. Only ended + # lanes consume sampling draws, so lanes that did not wrap keep their + # RNG stream: mask the batched state advance back for the rest. + sampling_cdf = receiver.sampling_cdf + count = 3 if sampling_cdf.numel() != 0 else 2 + draws, next_state = draw_uniform(ctx.rand.state, count) + sampled = _sample_motion_step( + draws, + sampling_cdf, + num_frames, + receiver.start_at_timestep_zero_prob, + ) + steps.copy_(torch.where(ended[:, None], sampled[:, None], advanced)) + ctx.rand.state.copy_(torch.where(ended[:, None], next_state, ctx.rand.state)) + ctx.sim_reset_requested |= ended[:, None] + + @staticmethod + @dispatch + def reset_env(receiver, ctx: ManagerContext) -> None: + """Sample start frames for reset lanes.""" + if bool(receiver.sequential_clips): + # Sequential playback rematerializes onto the current frame. + return + sampling_cdf = receiver.sampling_cdf + count = 3 if sampling_cdf.numel() != 0 else 2 + draws = ctx.rand.next_uniform(count) + steps = _sample_motion_step( + draws, + sampling_cdf, + receiver.clip.joint_pos.shape[0], + receiver.start_at_timestep_zero_prob, + ) + # ``clip_ended`` is intentionally left untouched; advance recomputes it + # on the following transition. + receiver.steps[:, 0] = steps + @configclass(kw_only=True) class WbtMotionCommandCfg(CommandCfg): diff --git a/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/observations.py b/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/observations.py index 5ed93bc8..1f6addef 100644 --- a/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/observations.py +++ b/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/observations.py @@ -4,6 +4,7 @@ """WBT command-backed observation terms.""" import numpy as np +import torch from motrix_env_core.config import configclass from motrix_env_core.config.scene import RobotCfg @@ -12,38 +13,25 @@ ObservationTermCfg, ObsTerm, kernel_data, - njit, ) -from motrix_env_core.manager.math.quaternion import inverse as quat_inverse -from motrix_env_core.manager.math.quaternion import mul as quat_mul -from motrix_env_core.manager.math.quaternion import rotate_inverse, to_matrix_first_two_rows -from motrix_env_core.mdp.noise import add_uniform_noise +from motrix_env_core.manager.dispatch import dispatch +from motrix_env_core.manager.kernel_data import SharedArray +from motrix_env_core.math import quaternion_array as batch_math +from motrix_env_core.math.quaternion_array import relative_orientation, rotate_inverse from motrix_env_core.mdp.observations import ( UniformNoiseCfg, + add_uniform_noise, ) -from motrix_env_core.numba.kernel_data import SharedArray -from motrix_env_core.numba.manager.dispatch import dispatch from motrix_envs.locomotion.wbt.mdp.command import WbtMotionCommand -@njit(inline="always") -def _write_relative_orientation_6d( - reference_quat_w: np.ndarray, - target_quat_w: np.ndarray, - out: np.ndarray, -) -> None: - relative_quat = out[:4] - quat_inverse(reference_quat_w, relative_quat) - quat_mul(relative_quat, target_quat_w, relative_quat) - to_matrix_first_two_rows(relative_quat, out[:6]) - - @dispatch -def motion_joint_obs(ctx: ManagerContext, out: np.ndarray) -> None: - motion: WbtMotionCommand = ctx.commands["motion"] +def motion_joint_obs(ctx: ManagerContext, out: torch.Tensor) -> None: + motion = ctx.commands["motion"] + steps = motion.steps[:, 0] joint_count = motion.clip.joint_pos.shape[1] - out[:joint_count] = motion.clip.joint_pos[motion.steps[0]] - out[joint_count:] = motion.clip.joint_vel[motion.steps[0]] + out[:, :joint_count] = motion.clip.joint_pos[steps] + out[:, joint_count:] = motion.clip.joint_vel[steps] @configclass(kw_only=True) @@ -56,39 +44,35 @@ def __call__(self, ctx) -> ObsTerm: @dispatch -def motion_reference_position_obs(ctx: ManagerContext, out: np.ndarray) -> None: - tracked_body_pos = ctx.sim["tracked_body_pos"] - tracked_body_quat = ctx.sim["tracked_body_quat"] - motion: WbtMotionCommand = ctx.commands["motion"] - reference_pos = tracked_body_pos[motion.reference_index] - reference_quat = tracked_body_quat[motion.reference_index] - target_ref_pos = motion.clip.reference_body_pos_w[motion.steps[0]] - rotate_inverse( +def motion_reference_position_obs(ctx: ManagerContext, out: torch.Tensor) -> None: + motion = ctx.commands["motion"] + reference_index = motion.reference_index + reference_pos = ctx.sim["tracked_body_pos"][:, reference_index] + reference_quat = ctx.sim["tracked_body_quat"][:, reference_index] + local = rotate_inverse( reference_quat, - ( - target_ref_pos[0] - reference_pos[0], - target_ref_pos[1] - reference_pos[1], - target_ref_pos[2] - reference_pos[2], - ), - out, + motion.clip.reference_body_pos_w[motion.steps[:, 0]] - reference_pos, ) + out.copy_(local) @configclass(kw_only=True) class MotionReferencePositionObsCfg(ObservationTermCfg): def __call__(self, ctx) -> ObsTerm: - del ctx return ObsTerm(3, motion_reference_position_obs) @dispatch -def motion_reference_orientation_obs(ctx: ManagerContext, out: np.ndarray, noise_amplitude: np.float32) -> None: - tracked_body_quat = ctx.sim["tracked_body_quat"] - motion: WbtMotionCommand = ctx.commands["motion"] - rand = ctx.rand - reference_quat = tracked_body_quat[motion.reference_index] - _write_relative_orientation_6d(reference_quat, motion.clip.reference_body_quat_w[motion.steps[0]], out) - add_uniform_noise(out, noise_amplitude, rand.state) +def motion_reference_orientation_obs( + ctx: ManagerContext, + out: torch.Tensor, + noise_amplitude: float, +) -> None: + motion = ctx.commands["motion"] + reference_quat = ctx.sim["tracked_body_quat"][:, motion.reference_index] + local = relative_orientation(reference_quat, motion.clip.reference_body_quat_w[motion.steps[:, 0]]) + out.copy_(local) + add_uniform_noise(ctx, out, noise_amplitude) @configclass(kw_only=True) @@ -96,28 +80,19 @@ class MotionReferenceOrientationObsCfg(ObservationTermCfg): noise: UniformNoiseCfg = UniformNoiseCfg() def __call__(self, ctx) -> ObsTerm: - del ctx return ObsTerm(6, motion_reference_orientation_obs, np.float32(self.noise.amplitude)) @dispatch -def robot_body_position_in_reference_frame_obs(ctx: ManagerContext, out: np.ndarray) -> None: - tracked_body_pos = ctx.sim["tracked_body_pos"] - tracked_body_quat = ctx.sim["tracked_body_quat"] - motion: WbtMotionCommand = ctx.commands["motion"] - reference_pos = tracked_body_pos[motion.reference_index] - reference_quat = tracked_body_quat[motion.reference_index] - for body_id in range(tracked_body_pos.shape[0]): - offset = body_id * 3 - rotate_inverse( - reference_quat, - ( - tracked_body_pos[body_id, 0] - reference_pos[0], - tracked_body_pos[body_id, 1] - reference_pos[1], - tracked_body_pos[body_id, 2] - reference_pos[2], - ), - out[offset : offset + 3], - ) +def robot_body_position_in_reference_frame_obs(ctx: ManagerContext, out: torch.Tensor) -> None: + motion = ctx.commands["motion"] + reference_index = motion.reference_index + reference_pos = ctx.sim["tracked_body_pos"][:, reference_index] + local = rotate_inverse( + ctx.sim["tracked_body_quat"][:, reference_index][:, None, :], + ctx.sim["tracked_body_pos"] - reference_pos[:, None, :], + ) + out.copy_(local.reshape(local.shape[0], -1)) @configclass(kw_only=True) @@ -128,16 +103,17 @@ def __call__(self, ctx) -> ObsTerm: @dispatch -def robot_body_orientation_obs(ctx: ManagerContext, out: np.ndarray) -> None: +def robot_body_orientation_obs(ctx: ManagerContext, out: torch.Tensor) -> None: + # Bug-compatible with the historical per-lane oracle: body zero writes its + # 6D matrix over the shared out[:4] inverse-quaternion scratch, and every + # later body reads those four matrix entries as its reference quaternion. + # Do not silently repair this alias; migrate the observation contract + # explicitly. tracked_body_quat = ctx.sim["tracked_body_quat"] - motion: WbtMotionCommand = ctx.commands["motion"] - quat_inverse(tracked_body_quat[motion.reference_index], out[:4]) - for body_id in range(tracked_body_quat.shape[0]): - offset = body_id * 6 - relative_quat = out[offset : offset + 4] - relative_quat[:] = out[:4] - quat_mul(relative_quat, tracked_body_quat[body_id], relative_quat) - to_matrix_first_two_rows(relative_quat, out[offset : offset + 6]) + reference_quat = tracked_body_quat[:, ctx.commands["motion"].reference_index] + first = batch_math.matrix6d(batch_math.mul(batch_math.inverse(reference_quat), tracked_body_quat[:, 0, :])) + rest = batch_math.matrix6d(batch_math.mul(first[:, None, :4], tracked_body_quat[:, 1:, :])) + out.copy_(torch.cat((first[:, None, :], rest), dim=1).reshape(tracked_body_quat.shape[0], -1)) @configclass(kw_only=True) @@ -154,11 +130,9 @@ class RelativePositionParams: @dispatch -def dof_pos_rel_obs(ctx: ManagerContext, out: np.ndarray, params: RelativePositionParams) -> None: - dof_pos = ctx.sim["robot_dof_pos"] - rand = ctx.rand - np.subtract(dof_pos, params.reference, out) - add_uniform_noise(out, params.noise_amplitude, rand.state) +def dof_pos_rel_obs(ctx: ManagerContext, out: torch.Tensor, params: RelativePositionParams) -> None: + out.copy_(ctx.sim["robot_dof_pos"] - params.reference) + add_uniform_noise(ctx, out, params.noise_amplitude) @configclass(kw_only=True) @@ -199,11 +173,9 @@ def __call__(self, ctx) -> ObsTerm: @dispatch -def dof_vel_obs(ctx: ManagerContext, out: np.ndarray, noise_amplitude: np.float32) -> None: - dof_vel = ctx.sim["robot_dof_vel"] - rand = ctx.rand - out[:] = dof_vel - add_uniform_noise(out, noise_amplitude, rand.state) +def dof_vel_obs(ctx: ManagerContext, out: torch.Tensor, noise_amplitude: float) -> None: + out.copy_(ctx.sim["robot_dof_vel"]) + add_uniform_noise(ctx, out, noise_amplitude) @configclass(kw_only=True) diff --git a/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/reset.py b/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/reset.py index 425eb1c2..50b2dd94 100644 --- a/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/reset.py +++ b/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/reset.py @@ -5,9 +5,8 @@ from __future__ import annotations -import math - import numpy as np +import torch from motrix_env_core.config import configclass from motrix_env_core.manager import ( @@ -15,10 +14,9 @@ ResetTerm, ResetTermCfg, ) -from motrix_env_core.mdp.action import JointPositionActionState -from motrix_env_core.numba.kernel_data import Map -from motrix_env_core.numba.manager.dispatch import dispatch -from motrix_env_core.numba.math import quaternion as numba_quaternion +from motrix_env_core.manager.dispatch import dispatch +from motrix_env_core.manager.kernel_data import Map +from motrix_env_core.math.quaternion_array import mul from motrix_env_core.sim import ( BodyAngularVelocityWrite, BodyLinearVelocityWrite, @@ -27,23 +25,20 @@ JointPositionWrite, JointVelocityWrite, ) -from motrix_envs.locomotion.wbt.mdp.command import WbtMotionCommand @dispatch def _reset_body_pos( ctx: ManagerContext, - sim_writes: Map[np.ndarray], - noise_scale: tuple[np.float32, np.float32, np.float32], + sim_writes: Map[torch.Tensor], + noise_amplitude: torch.Tensor, ) -> None: + motion = ctx.commands["motion"] position = sim_writes["position"] - motion: WbtMotionCommand = ctx.commands["motion"] - position[0] = motion.clip.root_body_pos_w[motion.steps[0]] - for index in range(3): - # next_uniform() returns [-1, 1), so this is centered uniform noise - # in [-scale, +scale) per axis. (An earlier "centering fix" subtracted - # 0.5 from an already-centered sample, producing a [-3s, +1s) bias.) - position[0, index] += ctx.rand.next_uniform() * noise_scale[index] + position[:, 0, :] = motion.clip.root_body_pos_w[motion.steps[:, 0]] + # next_uniform() returns [-1, 1), so this is centered uniform noise in + # [-scale, +scale) per axis. + position[:, 0, :] += ctx.rand.next_uniform(3) * noise_amplitude @configclass(kw_only=True) @@ -58,7 +53,7 @@ def __call__(self, ctx) -> ResetTerm: amplitude = np.asarray(self.noise, dtype=np.float32) * np.float32(self.noise_scale) return ResetTerm( _reset_body_pos, - tuple(amplitude), + torch.from_numpy(amplitude), writes={"position": BodyPositionWrite((body,))}, ) @@ -66,24 +61,38 @@ def __call__(self, ctx) -> ResetTerm: @dispatch def _reset_body_rot( ctx: ManagerContext, - sim_writes: Map[np.ndarray], - noise_scale: tuple[np.float32, np.float32, np.float32], + sim_writes: Map[torch.Tensor], + noise_amplitude: torch.Tensor, ) -> None: + motion = ctx.commands["motion"] rotation = sim_writes["rotation"] - motion: WbtMotionCommand = ctx.commands["motion"] - rotation[0] = motion.clip.root_body_quat_w[motion.steps[0]] - noisy_quat = np.empty((4,), dtype=np.float32) - numba_quaternion.from_euler( - ctx.rand.next_uniform() * noise_scale[0], - ctx.rand.next_uniform() * noise_scale[1], - ctx.rand.next_uniform() * noise_scale[2], - noisy_quat, - ) - base_quat = np.empty((4,), dtype=np.float32) - base_quat[:] = rotation[0] - numba_quaternion.mul(noisy_quat, base_quat, rotation[0]) - norm = math.sqrt(float(np.dot(rotation[0], rotation[0]))) - rotation[0] /= norm + base_quat = motion.clip.root_body_quat_w[motion.steps[:, 0]] + samples = ctx.rand.next_uniform(3) * noise_amplitude + # Euler roll/pitch/yaw noise left-multiplied onto the frame quaternion, + # renormalized. Trig runs in float64 and rounds to float32 on store; the + # squared-norm reduction accumulates sequentially in float32 and the + # square root goes through float64 before rounding back — torch's direct + # float32 sqrt is 1 ulp off IEEE on CPU for values just below 1. The + # final divide stays float32. + roll, pitch, yaw = (samples.double() * 0.5).unbind(dim=-1) + cr, cp, cy = torch.cos(roll), torch.cos(pitch), torch.cos(yaw) + sr, sp, sy = torch.sin(roll), torch.sin(pitch), torch.sin(yaw) + noise_quat = torch.stack( + ( + sr * cp * cy - cr * sp * sy, + cr * sp * cy + sr * cp * sy, + cr * cp * sy - sr * sp * cy, + cr * cp * cy + sr * sp * sy, + ), + dim=-1, + ).to(torch.float32) + noisy = mul(noise_quat, base_quat) + squared = noisy * noisy + norm_squared = squared[:, 0:1] + for index in range(1, squared.shape[1]): + norm_squared = norm_squared + squared[:, index : index + 1] + norm = torch.sqrt(norm_squared.double()).to(torch.float32) + rotation[:, 0, :] = noisy / norm @configclass(kw_only=True) @@ -98,7 +107,7 @@ def __call__(self, ctx) -> ResetTerm: amplitude = np.asarray(self.noise, dtype=np.float32) * np.float32(self.noise_scale) return ResetTerm( _reset_body_rot, - tuple(amplitude), + torch.from_numpy(amplitude), writes={"rotation": BodyRotationWrite((body,))}, ) @@ -106,15 +115,13 @@ def __call__(self, ctx) -> ResetTerm: @dispatch def _reset_body_lin_vel( ctx: ManagerContext, - sim_writes: Map[np.ndarray], - noise_scale: tuple[np.float32, np.float32, np.float32], + sim_writes: Map[torch.Tensor], + noise_amplitude: torch.Tensor, ) -> None: + motion = ctx.commands["motion"] linear_velocity = sim_writes["linear_velocity"] - motion: WbtMotionCommand = ctx.commands["motion"] - linear_velocity[0] = motion.clip.root_body_lin_vel_w[motion.steps[0]] - for index in range(3): - # Centered noise: see _reset_body_pos. - linear_velocity[0, index] += ctx.rand.next_uniform() * noise_scale[index] + linear_velocity[:, 0, :] = motion.clip.root_body_lin_vel_w[motion.steps[:, 0]] + linear_velocity[:, 0, :] += ctx.rand.next_uniform(3) * noise_amplitude @configclass(kw_only=True) @@ -129,7 +136,7 @@ def __call__(self, ctx) -> ResetTerm: amplitude = np.asarray(self.noise, dtype=np.float32) * np.float32(self.noise_scale) return ResetTerm( _reset_body_lin_vel, - tuple(amplitude), + torch.from_numpy(amplitude), writes={"linear_velocity": BodyLinearVelocityWrite((body,))}, ) @@ -137,15 +144,13 @@ def __call__(self, ctx) -> ResetTerm: @dispatch def _reset_body_rot_vel( ctx: ManagerContext, - sim_writes: Map[np.ndarray], - noise_scale: tuple[np.float32, np.float32, np.float32], + sim_writes: Map[torch.Tensor], + noise_amplitude: torch.Tensor, ) -> None: + motion = ctx.commands["motion"] angular_velocity = sim_writes["angular_velocity"] - motion: WbtMotionCommand = ctx.commands["motion"] - angular_velocity[0] = motion.clip.root_body_ang_vel_w[motion.steps[0]] - for index in range(3): - # Centered noise: see _reset_body_pos. - angular_velocity[0, index] += ctx.rand.next_uniform() * noise_scale[index] + angular_velocity[:, 0, :] = motion.clip.root_body_ang_vel_w[motion.steps[:, 0]] + angular_velocity[:, 0, :] += ctx.rand.next_uniform(3) * noise_amplitude @configclass(kw_only=True) @@ -160,24 +165,24 @@ def __call__(self, ctx) -> ResetTerm: amplitude = np.asarray(self.noise, dtype=np.float32) * np.float32(self.noise_scale) return ResetTerm( _reset_body_rot_vel, - tuple(amplitude), + torch.from_numpy(amplitude), writes={"angular_velocity": BodyAngularVelocityWrite((body,))}, ) @dispatch -def _reset_body_dof_pos(ctx: ManagerContext, sim_writes: Map[np.ndarray], noise_scale: np.float32) -> None: +def _reset_body_dof_pos(ctx: ManagerContext, sim_writes: Map[torch.Tensor], noise_scale: float) -> None: + motion = ctx.commands["motion"] + steps = motion.steps[:, 0] position = sim_writes["position"] velocity = sim_writes["velocity"] - motion: WbtMotionCommand = ctx.commands["motion"] - action: JointPositionActionState = ctx.actions["joint_position"] - step = motion.steps[0] - position[:] = motion.clip.joint_pos[step] - velocity[:] = motion.clip.joint_vel[step] - rng = ctx.rand - for index in range(position.shape[0]): - value = position[index] + rng.next_uniform() * noise_scale - position[index] = min(max(value, action.joint_lower[index]), action.joint_upper[index]) + position[:] = motion.clip.joint_pos[steps] + velocity[:] = motion.clip.joint_vel[steps] + action = ctx.actions["joint_position"] + position += ctx.rand.next_uniform(position.shape[1]) * noise_scale + # clamp has no out= form Dynamo accepts on the arena-backed + # (non-contiguous) write view; copy_ back instead. + position.copy_(torch.clamp(position, min=action.joint_lower, max=action.joint_upper)) @configclass(kw_only=True) diff --git a/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/rewards.py b/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/rewards.py index c2c25688..17047fd5 100644 --- a/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/rewards.py +++ b/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/rewards.py @@ -3,27 +3,25 @@ """WBT command-backed reward terms.""" -import math from typing import cast import numpy as np +import torch from motrix_env_core.config import configclass from motrix_env_core.manager import ManagerContext, RewardTerm, RewardTermCfg -from motrix_env_core.manager.math.quaternion import rotation_distance +from motrix_env_core.manager.dispatch import dispatch +from motrix_env_core.manager.kernel_data import SharedArray, kernel_data +from motrix_env_core.math.quaternion_array import rotation_distance from motrix_env_core.mdp.action import JointPositionActionState -from motrix_env_core.numba.kernel_data import SharedArray, kernel_data -from motrix_env_core.numba.manager.dispatch import dispatch -from motrix_envs.locomotion.wbt.mdp.command import WbtMotionCommand @dispatch -def global_ref_position_reward(ctx: ManagerContext, sigma: np.float32) -> float: - tracked_body_pos = ctx.sim["tracked_body_pos"] - motion: WbtMotionCommand = ctx.commands["motion"] - error = motion.clip.reference_body_pos_w[motion.steps[0]] - tracked_body_pos[motion.reference_index] - error_sq = float(np.dot(error, error)) - return math.exp(-error_sq / (sigma * sigma)) +def global_ref_position_reward(ctx: ManagerContext, sigma: float) -> torch.Tensor: + motion = ctx.commands["motion"] + steps = motion.steps[:, 0] + error = motion.clip.reference_body_pos_w[steps] - ctx.sim["tracked_body_pos"][:, motion.reference_index] + return torch.exp(-torch.sum(error * error, dim=-1) / (sigma * sigma)) @configclass(kw_only=True) @@ -31,19 +29,17 @@ class GlobalRefPositionRewardCfg(RewardTermCfg): sigma: float def __call__(self, ctx) -> RewardTerm: - del ctx return RewardTerm(global_ref_position_reward, np.float32(self.sigma)) @dispatch -def global_ref_orientation_reward(ctx: ManagerContext, sigma: np.float32) -> float: - tracked_body_quat = ctx.sim["tracked_body_quat"] - motion: WbtMotionCommand = ctx.commands["motion"] - distance = rotation_distance( - motion.clip.reference_body_quat_w[motion.steps[0]], - tracked_body_quat[motion.reference_index], - ) - return math.exp(-(distance * distance) / (sigma * sigma)) +def global_ref_orientation_reward(ctx: ManagerContext, sigma: float) -> torch.Tensor: + motion = ctx.commands["motion"] + steps = motion.steps[:, 0] + target = motion.clip.reference_body_quat_w[steps] + quaternion = ctx.sim["tracked_body_quat"][:, motion.reference_index] + distance = rotation_distance(target, quaternion) + return torch.exp(-(distance * distance) / (sigma * sigma)) @configclass(kw_only=True) @@ -51,17 +47,15 @@ class GlobalRefOrientationRewardCfg(RewardTermCfg): sigma: float def __call__(self, ctx) -> RewardTerm: - del ctx return RewardTerm(global_ref_orientation_reward, np.float32(self.sigma)) @dispatch -def relative_body_position_reward(ctx: ManagerContext, sigma: np.float32) -> float: - tracked_body_pos = ctx.sim["tracked_body_pos"] - motion: WbtMotionCommand = ctx.commands["motion"] - diff = motion.target_body_position_relative - tracked_body_pos - error_sq = float(np.sum(diff * diff)) - return math.exp(-(error_sq / tracked_body_pos.shape[0]) / (sigma * sigma)) +def relative_body_position_reward(ctx: ManagerContext, sigma: float) -> torch.Tensor: + motion = ctx.commands["motion"] + diff = motion.target_body_position_relative - ctx.sim["tracked_body_pos"] + error_sq = torch.sum(diff * diff, dim=(-2, -1)) / diff.shape[-2] + return torch.exp(-error_sq / (sigma * sigma)) @configclass(kw_only=True) @@ -69,19 +63,16 @@ class RelativeBodyPositionRewardCfg(RewardTermCfg): sigma: float def __call__(self, ctx) -> RewardTerm: - del ctx return RewardTerm(relative_body_position_reward, np.float32(self.sigma)) @dispatch -def ee_body_pos_reward(ctx: ManagerContext, body_indices: tuple[int, ...], sigma: np.float32) -> float: - tracked_body_pos = ctx.sim["tracked_body_pos"] - motion: WbtMotionCommand = ctx.commands["motion"] - error_sq = 0.0 - for body_id in body_indices: - diff = motion.target_body_position_relative[body_id] - tracked_body_pos[body_id] - error_sq += float(np.dot(diff, diff)) - return math.exp(-(error_sq / len(body_indices)) / (sigma * sigma)) +def ee_body_pos_reward(ctx: ManagerContext, body_indices: tuple[int, ...], sigma: float) -> torch.Tensor: + motion = ctx.commands["motion"] + position = ctx.sim["tracked_body_pos"] + error = motion.target_body_position_relative[:, body_indices, :] - position[:, body_indices, :] + error_sq = torch.sum(error * error, dim=(-2, -1)) / len(body_indices) + return torch.exp(-error_sq / (sigma * sigma)) @configclass(kw_only=True) @@ -111,17 +102,12 @@ def __call__(self, ctx) -> RewardTerm: @dispatch -def relative_body_orientation_reward(ctx: ManagerContext, sigma: np.float32) -> float: - tracked_body_quat = ctx.sim["tracked_body_quat"] - motion: WbtMotionCommand = ctx.commands["motion"] - error_sq = 0.0 - for body_id in range(tracked_body_quat.shape[0]): - distance = rotation_distance( - motion.target_body_orientation_relative[body_id], - tracked_body_quat[body_id], - ) - error_sq += distance * distance - return math.exp(-(error_sq / tracked_body_quat.shape[0]) / (sigma * sigma)) +def relative_body_orientation_reward(ctx: ManagerContext, sigma: float) -> torch.Tensor: + motion = ctx.commands["motion"] + quaternion = ctx.sim["tracked_body_quat"] + distance = rotation_distance(motion.target_body_orientation_relative, quaternion) + error_sq = torch.sum(distance * distance, dim=-1) / quaternion.shape[-2] + return torch.exp(-error_sq / (sigma * sigma)) @configclass(kw_only=True) @@ -129,17 +115,17 @@ class RelativeBodyOrientationRewardCfg(RewardTermCfg): sigma: float def __call__(self, ctx) -> RewardTerm: - del ctx return RewardTerm(relative_body_orientation_reward, np.float32(self.sigma)) @dispatch -def global_body_linear_velocity_reward(ctx: ManagerContext, sigma: np.float32) -> float: - tracked_body_linear_velocity = ctx.sim["tracked_body_linear_velocity"] - motion: WbtMotionCommand = ctx.commands["motion"] - diff = motion.clip.tracked_bodies_lin_vel_w[motion.steps[0]] - tracked_body_linear_velocity - error_sq = float(np.sum(diff * diff)) - return math.exp(-(error_sq / tracked_body_linear_velocity.shape[0]) / (sigma * sigma)) +def global_body_linear_velocity_reward(ctx: ManagerContext, sigma: float) -> torch.Tensor: + motion = ctx.commands["motion"] + steps = motion.steps[:, 0] + velocity = ctx.sim["tracked_body_linear_velocity"] + error = motion.clip.tracked_bodies_lin_vel_w[steps] - velocity + error_sq = torch.sum(error * error, dim=(-2, -1)) / velocity.shape[-2] + return torch.exp(-error_sq / (sigma * sigma)) @configclass(kw_only=True) @@ -147,17 +133,17 @@ class GlobalBodyLinearVelocityRewardCfg(RewardTermCfg): sigma: float def __call__(self, ctx) -> RewardTerm: - del ctx return RewardTerm(global_body_linear_velocity_reward, np.float32(self.sigma)) @dispatch -def global_body_angular_velocity_reward(ctx: ManagerContext, sigma: np.float32) -> float: - tracked_body_angular_velocity = ctx.sim["tracked_body_angular_velocity"] - motion: WbtMotionCommand = ctx.commands["motion"] - diff = motion.clip.tracked_bodies_ang_vel_w[motion.steps[0]] - tracked_body_angular_velocity - error_sq = float(np.sum(diff * diff)) - return math.exp(-(error_sq / tracked_body_angular_velocity.shape[0]) / (sigma * sigma)) +def global_body_angular_velocity_reward(ctx: ManagerContext, sigma: float) -> torch.Tensor: + motion = ctx.commands["motion"] + steps = motion.steps[:, 0] + velocity = ctx.sim["tracked_body_angular_velocity"] + error = motion.clip.tracked_bodies_ang_vel_w[steps] - velocity + error_sq = torch.sum(error * error, dim=(-2, -1)) / velocity.shape[-2] + return torch.exp(-error_sq / (sigma * sigma)) @configclass(kw_only=True) @@ -165,7 +151,6 @@ class GlobalBodyAngularVelocityRewardCfg(RewardTermCfg): sigma: float def __call__(self, ctx) -> RewardTerm: - del ctx return RewardTerm(global_body_angular_velocity_reward, np.float32(self.sigma)) @@ -178,10 +163,11 @@ class DofLimitParams: @dispatch -def dof_limit_reward(ctx: ManagerContext, params: DofLimitParams) -> float: - dof_pos = ctx.sim["robot_dof_pos"] - violation = np.abs(dof_pos - params.midpoint) - params.half_range * params.soft_limit - return min(float(np.sum(np.maximum(violation, 0.0))), params.cap) +def dof_limit_reward(ctx: ManagerContext, params: DofLimitParams) -> torch.Tensor: + dof_position = ctx.sim["robot_dof_pos"] + violation = torch.abs(dof_position - params.midpoint) - params.half_range * params.soft_limit + total = torch.sum(torch.where(violation > 0.0, violation, 0.0), dim=-1) + return torch.where(total > params.cap, params.cap, total) @configclass(kw_only=True) @@ -201,14 +187,11 @@ def __call__(self, ctx) -> RewardTerm: @dispatch -def undesired_contacts_reward(ctx: ManagerContext, threshold: np.float32) -> float: +def undesired_contacts_reward(ctx: ManagerContext, threshold: float) -> torch.Tensor: contact_forces = ctx.sim["undesired_contact_forces"] - count = 0.0 - for link_id in range(contact_forces.shape[0]): - fx, fy, fz = contact_forces[link_id, :3] - if math.sqrt(fx * fx + fy * fy + fz * fz) > threshold: - count += 1.0 - return count + forces = contact_forces[..., :3] + norm = torch.sqrt(torch.sum(forces * forces, dim=-1)) + return torch.sum((norm > threshold).to(contact_forces.dtype), dim=-1) @configclass(kw_only=True) @@ -216,5 +199,4 @@ class UndesiredContactsRewardCfg(RewardTermCfg): threshold: float def __call__(self, ctx) -> RewardTerm: - del ctx return RewardTerm(undesired_contacts_reward, np.float32(self.threshold)) diff --git a/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/terminations.py b/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/terminations.py index 2c68d2a8..03b7194f 100644 --- a/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/terminations.py +++ b/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/terminations.py @@ -3,15 +3,12 @@ """WBT termination terms and raw diagnostic metrics.""" -import math - import numpy as np +import torch from motrix_env_core.config import configclass from motrix_env_core.manager import ManagerContext, TerminationTerm, TerminationTermCfg -from motrix_env_core.mdp.action import JointPositionActionState -from motrix_env_core.numba.manager.dispatch import dispatch -from motrix_envs.locomotion.wbt.mdp.command import WbtMotionCommand +from motrix_env_core.manager.dispatch import dispatch @configclass(kw_only=True) @@ -20,18 +17,19 @@ class _WbtTerminationCfg(TerminationTermCfg): @dispatch -def bad_ref_z_termination(ctx: ManagerContext, threshold: np.float32) -> bool: - tracked_body_pos = ctx.sim["tracked_body_pos"] - motion: WbtMotionCommand = ctx.commands["motion"] - error = abs(motion.clip.reference_body_pos_w[motion.steps[0], 2] - tracked_body_pos[motion.reference_index, 2]) - ctx.metrics["ref_z_abs_err"][0] = error +def bad_ref_z_termination(ctx: ManagerContext, threshold: float) -> torch.Tensor: + motion = ctx.commands["motion"] + steps = motion.steps[:, 0] + error = torch.abs( + motion.clip.reference_body_pos_w[steps][..., 2] - ctx.sim["tracked_body_pos"][:, motion.reference_index][..., 2] + ) + ctx.metrics["ref_z_abs_err"][:, 0] = error return error > threshold @configclass(kw_only=True) class BadRefZTerminationCfg(_WbtTerminationCfg): def __call__(self, ctx) -> TerminationTerm: - del ctx return TerminationTerm( bad_ref_z_termination, np.float32(self.threshold), @@ -40,22 +38,19 @@ def __call__(self, ctx) -> TerminationTerm: @dispatch -def bad_ref_orientation_termination(ctx: ManagerContext, threshold: np.float32) -> bool: - tracked_body_quat = ctx.sim["tracked_body_quat"] - motion: WbtMotionCommand = ctx.commands["motion"] - target = motion.clip.reference_body_quat_w[motion.steps[0]] - motion_gravity_z = 2.0 * (target[0] * target[0] + target[1] * target[1]) - 1.0 - robot_quat = tracked_body_quat[motion.reference_index] - robot_gravity_z = 2.0 * (robot_quat[0] * robot_quat[0] + robot_quat[1] * robot_quat[1]) - 1.0 - error = abs(motion_gravity_z - robot_gravity_z) - ctx.metrics["ref_ori_abs_err"][0] = error +def bad_ref_orientation_termination(ctx: ManagerContext, threshold: float) -> torch.Tensor: + motion = ctx.commands["motion"] + steps = motion.steps[:, 0] + target = motion.clip.reference_body_quat_w[steps] + actual = ctx.sim["tracked_body_quat"][:, motion.reference_index] + error = torch.abs(2 * (target[..., 0] ** 2 + target[..., 1] ** 2 - actual[..., 0] ** 2 - actual[..., 1] ** 2)) + ctx.metrics["ref_ori_abs_err"][:, 0] = error return error > threshold @configclass(kw_only=True) class BadRefOrientationTerminationCfg(_WbtTerminationCfg): def __call__(self, ctx) -> TerminationTerm: - del ctx return TerminationTerm( bad_ref_orientation_termination, np.float32(self.threshold), @@ -67,8 +62,8 @@ def __call__(self, ctx) -> TerminationTerm: def bad_motion_body_position_termination( ctx: ManagerContext, body_indices: tuple[int, ...], - threshold: np.float32, -) -> bool: + threshold: float, +) -> torch.Tensor: """Holosoma-aligned tracked-body tracking termination. Terminates when the mean position error of the configured tracked bodies @@ -76,19 +71,15 @@ def bad_motion_body_position_termination( their reference targets (the yaw/height-anchored body targets the relative rewards track) exceeds ``threshold``. """ - tracked_body_pos = ctx.sim["tracked_body_pos"] - motion: WbtMotionCommand = ctx.commands["motion"] - count = len(body_indices) - if count == 0: - ctx.metrics["motion_body_pos_err"][0] = 0.0 - return False - error_sq_sum = 0.0 - for index in range(count): - body_id = body_indices[index] - diff = motion.target_body_position_relative[body_id] - tracked_body_pos[body_id] - error_sq_sum += float(np.dot(diff, diff)) - mean_err = math.sqrt(error_sq_sum / count) - ctx.metrics["motion_body_pos_err"][0] = mean_err + motion = ctx.commands["motion"] + if not body_indices: + error = torch.zeros_like(motion.steps[:, 0], dtype=torch.float32) + ctx.metrics["motion_body_pos_err"][:, 0] = error + return torch.zeros_like(error, dtype=torch.bool) + diff = motion.target_body_position_relative[:, body_indices, :] - ctx.sim["tracked_body_pos"][:, body_indices, :] + error_sq_sum = torch.sum(diff * diff, dim=(-2, -1)) + mean_err = torch.sqrt(error_sq_sum / len(body_indices)) + ctx.metrics["motion_body_pos_err"][:, 0] = mean_err return mean_err > threshold @@ -111,17 +102,16 @@ def __call__(self, ctx) -> TerminationTerm: def bad_body_z_termination( ctx: ManagerContext, body_indices: tuple[int, ...], - threshold: np.float32, -) -> bool: - tracked_body_pos = ctx.sim["tracked_body_pos"] - motion: WbtMotionCommand = ctx.commands["motion"] - error = 0.0 - for body_id in body_indices: - error = max( - error, - abs(motion.target_body_position_relative[body_id, 2] - tracked_body_pos[body_id, 2]), - ) - ctx.metrics["body_z_max_err"][0] = error + threshold: float, +) -> torch.Tensor: + motion = ctx.commands["motion"] + target = motion.target_body_position_relative[:, body_indices, :] + actual = ctx.sim["tracked_body_pos"][:, body_indices, :] + if target.shape[-2] == 0: + error = torch.zeros_like(torch.sum(target, dim=(-2, -1))) + else: + error = torch.max(torch.abs(target[..., 2] - actual[..., 2]), dim=-1).values + ctx.metrics["body_z_max_err"][:, 0] = error return error > threshold @@ -141,28 +131,21 @@ def __call__(self, ctx) -> TerminationTerm: @dispatch -def bad_dof_position_termination(ctx: ManagerContext, threshold: np.float32) -> bool: - dof_pos = ctx.sim["robot_dof_pos"] - action: JointPositionActionState = ctx.actions["joint_position"] - error = 0.0 - finite = True - for joint_id in range(dof_pos.shape[0]): - position = dof_pos[joint_id] - if math.isfinite(position): - violation = max(action.joint_lower[joint_id] - position, 0.0) - violation += max(position - action.joint_upper[joint_id], 0.0) - error = max(error, violation) - else: - finite = False - error = math.inf - ctx.metrics["dof_limit_violation_max"][0] = error - return (not finite) or error > threshold +def bad_dof_position_termination(ctx: ManagerContext, threshold: float) -> torch.Tensor: + action = ctx.actions["joint_position"] + position = ctx.sim["robot_dof_pos"] + zero = torch.zeros_like(position) + violation = torch.maximum(action.joint_lower - position, zero) + torch.maximum(position - action.joint_upper, zero) + finite = torch.all(torch.isfinite(position), dim=-1) + error = torch.max(violation, dim=-1).values + error = torch.where(finite, error, torch.full_like(error, float("inf"))) + ctx.metrics["dof_limit_violation_max"][:, 0] = error + return error > threshold @configclass(kw_only=True) class BadDofPositionTerminationCfg(_WbtTerminationCfg): def __call__(self, ctx) -> TerminationTerm: - del ctx return TerminationTerm( bad_dof_position_termination, np.float32(self.threshold), diff --git a/motrix_envs/src/motrix_envs/motion/library.py b/motrix_envs/src/motrix_envs/motion/library.py index 06f06a0a..afa147c5 100644 --- a/motrix_envs/src/motrix_envs/motion/library.py +++ b/motrix_envs/src/motrix_envs/motion/library.py @@ -16,7 +16,7 @@ import numpy as np -from motrix_env_core.numba.kernel_data import Map +from motrix_env_core.manager.kernel_data import Map from motrix_envs.motion.loader import MotrixMotion from motrix_envs.motion.tracked import MotionChannel, WbtMotionClip diff --git a/motrix_envs/src/motrix_envs/motion/tracked.py b/motrix_envs/src/motrix_envs/motion/tracked.py index d3c482e6..74b1c481 100644 --- a/motrix_envs/src/motrix_envs/motion/tracked.py +++ b/motrix_envs/src/motrix_envs/motion/tracked.py @@ -7,7 +7,7 @@ import numpy as np -from motrix_env_core.numba.kernel_data import Map, SharedArray, kernel_data +from motrix_env_core.manager.kernel_data import Map, SharedArray, kernel_data from motrix_envs.motion.loader import MotionFormatError, MotrixMotion diff --git a/motrix_envs/tests/test_mdp_obs.py b/motrix_envs/tests/test_mdp_obs.py index 09257300..0dcae084 100644 --- a/motrix_envs/tests/test_mdp_obs.py +++ b/motrix_envs/tests/test_mdp_obs.py @@ -5,8 +5,11 @@ import numpy as np import pytest +import torch from motrix_env_core.config.scene import KeyPoseCfg, ModelFileCfg, RobotCfg +from motrix_env_core.manager.context import BuildContext +from motrix_env_core.manager.rand import initialize_rand_states from motrix_env_core.mdp.action import JointPositionActionState from motrix_env_core.mdp.observations import ( ActionsObsCfg, @@ -15,9 +18,6 @@ BodyLinearVelocityObsCfg, UniformNoiseCfg, ) -from motrix_env_core.mdp.state import RandValue -from motrix_env_core.numba.manager.context import BuildContext -from motrix_env_core.numba.manager.rand import initialize_rand_states from motrix_env_core.sim import ( BodyJointPositionQuery, BodyJointVelocityQuery, @@ -26,6 +26,7 @@ LinkLinearVelocityQuery, LinkQuaternionQuery, ) +from motrix_env_core.torch_manager.runtime import _BatchRand from motrix_envs.locomotion.wbt.mdp.observations import ( DofPosRelObsCfg, ) @@ -117,21 +118,22 @@ def test_actions_observation_reads_current_actions() -> None: env = SimpleNamespace(action_terms={"joint_position": SimpleNamespace(state=action)}) cfg = ActionsObsCfg() term = cfg.__call__(env) - current = np.asarray([1.0, 2.0, 3.0], dtype=np.float32) - out = np.empty(term.size, dtype=np.float32) + current = np.asarray([[1.0, 2.0, 3.0]], dtype=np.float32) + out = torch.empty((1, term.size), dtype=torch.float32) assert term.size == 3 - # Manager kernels receive a lane view: queue shape is (W, A), not (N, W, A). - queue = np.zeros((2, 3), dtype=np.float32) - queue[1] = current - lane_state = SimpleNamespace( - current=lambda: queue[1], + # The batched context exposes the whole-batch queue; current() is the + # queue row the action pointer selects. + queue = torch.zeros((1, 2, 3), dtype=torch.float32) + queue[:, 1] = torch.from_numpy(current) + batch_state = SimpleNamespace( + current=lambda: queue[:, 1], action_queue=queue, - action_ptr=np.asarray([1], dtype=np.int64), + action_ptr=torch.asarray([1], dtype=torch.int64), ) - ctx = SimpleNamespace(actions={"joint_position": lane_state}) + ctx = SimpleNamespace(actions={"joint_position": batch_state}) term.dispatch(ctx, out, *term.args) - np.testing.assert_array_equal(out, current) + np.testing.assert_array_equal(out.numpy(), current) def test_robot_observation_cfg_uses_sim_query_layout_and_robot_key_pose() -> None: @@ -167,35 +169,33 @@ def test_robot_dof_pos_rel_obs_cfg_requires_scene_robot() -> None: def test_robot_base_velocities_are_expressed_in_the_base_local_frame() -> None: half_sqrt_two = np.float32(np.sqrt(0.5)) - base_quat = np.asarray([0.0, 0.0, half_sqrt_two, half_sqrt_two], dtype=np.float32) - linear_velocity = np.asarray([1.0, 0.0, 0.0], dtype=np.float32) - angular_velocity = np.asarray([0.0, 2.0, 0.0], dtype=np.float32) - noise = RandValue(np.asarray([1], dtype=np.uint64)) + base_quat = torch.from_numpy(np.asarray([[0.0, 0.0, half_sqrt_two, half_sqrt_two]], dtype=np.float32)) + linear_velocity = torch.from_numpy(np.asarray([[1.0, 0.0, 0.0]], dtype=np.float32)) + angular_velocity = torch.from_numpy(np.asarray([[0.0, 2.0, 0.0]], dtype=np.float32)) linear = BodyLinearVelocityObsCfg().__call__(_env()) angular = BodyAngularVelocityObsCfg().__call__(_env()) - linear_out = np.empty(3, dtype=np.float32) - angular_out = np.empty(3, dtype=np.float32) + linear_out = torch.empty((1, 3), dtype=torch.float32) + angular_out = torch.empty((1, 3), dtype=torch.float32) - # Str args name sim keys; the compiler resolves them to lane arrays, so a - # direct dispatch call substitutes them explicitly. - ctx = SimpleNamespace(rand=noise) - linear.dispatch(ctx, linear_out, base_quat, linear_velocity, *linear.args[2:]) - angular.dispatch(ctx, angular_out, base_quat, angular_velocity, *angular.args[2:]) + # Str args name sim keys; the resolver binds them to whole-batch arrays, + # so a direct dispatch call substitutes them explicitly. + linear.dispatch(None, linear_out, base_quat, linear_velocity, *linear.args[2:]) + angular.dispatch(None, angular_out, base_quat, angular_velocity, *angular.args[2:]) - np.testing.assert_allclose(linear_out, [0.0, -1.0, 0.0], atol=1e-6) - np.testing.assert_allclose(angular_out, [2.0, 0.0, 0.0], atol=1e-6) + np.testing.assert_allclose(linear_out.numpy()[0], [0.0, -1.0, 0.0], atol=1e-6) + np.testing.assert_allclose(angular_out.numpy()[0], [2.0, 0.0, 0.0], atol=1e-6) def test_robot_observation_noise_uses_one_stateful_sequence_per_environment() -> None: - dof_vel = np.asarray([1.0, 2.0], dtype=np.float32) + dof_vel = torch.from_numpy(np.tile(np.asarray([1.0, 2.0], dtype=np.float32), (2, 1))) term = BodyJointVelObsCfg(noise=UniformNoiseCfg(amplitude=0.25)).__call__(_env()) def evaluate(rng_states: np.ndarray) -> np.ndarray: - out = np.empty((2, dof_vel.shape[0]), dtype=np.float32) - for env_id in range(2): - ctx = SimpleNamespace(rand=RandValue(rng_states[env_id])) - term.dispatch(ctx, out[env_id], dof_vel, *term.args[1:]) - return out + out = torch.empty((2, dof_vel.shape[1]), dtype=torch.float32) + # The batch context shares (and advances) the per-env RNG state in place. + rand = _BatchRand(torch.from_numpy(rng_states.view(np.int64))) + term.dispatch(SimpleNamespace(rand=rand), out, dof_vel, *term.args[1:]) + return out.numpy() rng_states = initialize_rand_states(2, rand_seed=7) first = evaluate(rng_states) @@ -206,6 +206,6 @@ def evaluate(rng_states: np.ndarray) -> np.ndarray: np.testing.assert_array_equal(reproduced, first) assert np.any(continued != first) assert np.any(different_seed != first) - assert np.all(first >= dof_vel - 0.25) - assert np.all(first < dof_vel + 0.25) + assert np.all(first >= dof_vel.numpy() - 0.25) + assert np.all(first < dof_vel.numpy() + 0.25) assert not np.array_equal(first[0], first[1]) diff --git a/motrix_envs/tests/test_wbt_batch_math.py b/motrix_envs/tests/test_wbt_batch_math.py new file mode 100644 index 00000000..58290f74 --- /dev/null +++ b/motrix_envs/tests/test_wbt_batch_math.py @@ -0,0 +1,55 @@ +# Copyright Motphys Technology Co., Ltd. 2025, 2026 +# SPDX-License-Identifier: Apache-2.0 + +"""Batch quaternion math retains scalar helper semantics and broadcasting.""" + +import numpy as np +import torch + +from motrix_env_core.math import quaternion as scalar +from motrix_env_core.math import quaternion_array as batch + + +def test_quaternion_batch_matches_scalar_and_compiles(): + rng = np.random.default_rng(8) + q = rng.normal(size=(7, 4)).astype(np.float32) + other = rng.normal(size=(7, 4)).astype(np.float32) + v = rng.normal(size=(7, 3)).astype(np.float32) + + def compute(q, other, v): + return batch.inverse(q), batch.mul(q, other), batch.rotate_inverse(q, v), batch.matrix6d(q) + + expected = tuple( + np.stack([function(*args) for args in zip(*arrays, strict=True)]) + for function, arrays in ( + (scalar.inverse, (q,)), + (scalar.mul, (q, other)), + (scalar.rotate_inverse, (q, v)), + # Rows of the rotation matrix: rotate_vector(q, e_i) yields + # column i, so transpose the column stack and keep two rows. + ( + lambda quat: np.stack( + [ + np.stack( + [ + scalar.rotate_vector(unit[None, :], np.eye(3, dtype=quat.dtype)[i : i + 1])[0] + for i in range(3) + ] + ) + .T[:2] + .reshape(-1) + for unit in (quat / np.linalg.norm(quat, axis=-1, keepdims=True),) + ] + )[0], + (q,), + ), + ) + ) + for actuals in ( + compute(q, other, v), + torch.compile(compute, fullgraph=True)(*(torch.from_numpy(a) for a in (q, other, v))), + ): + for actual, reference in zip(actuals, expected, strict=True): + if isinstance(actual, torch.Tensor): + actual = actual.numpy() + np.testing.assert_allclose(actual, reference, rtol=2e-5, atol=2e-5) diff --git a/motrix_envs/tests/test_wbt_numba.py b/motrix_envs/tests/test_wbt_numba.py deleted file mode 100644 index d658e3b0..00000000 --- a/motrix_envs/tests/test_wbt_numba.py +++ /dev/null @@ -1,1054 +0,0 @@ -# Copyright Motphys Technology Co., Ltd. 2025, 2026 -# SPDX-License-Identifier: Apache-2.0 - -import pickle -from dataclasses import fields, replace -from types import SimpleNamespace - -import numba -import numpy as np -import pytest - -from motrix_env_core import registry # noqa: E402 -from motrix_env_core.manager import ( # noqa: E402 - ManagerContext, - ManagerEnv, - ManagerResetCfg, - RewardTermCfg, -) -from motrix_env_core.mdp.action import ( # noqa: E402 - JointPositionActionCfg, - JointPositionActionState, -) -from motrix_env_core.mdp.noise import add_uniform_noise_batch # noqa: E402 -from motrix_env_core.mdp.observations import ( # noqa: E402 - BodyAngularVelocityObsCfg, - UniformNoiseCfg, -) -from motrix_env_core.mdp.state import RandValue # noqa: E402 -from motrix_env_core.numba.manager.commands import ResetContext # noqa: E402 -from motrix_env_core.numba.manager.compiler.compiler import TermScalarBuffer # noqa: E402 -from motrix_env_core.numba.manager.rand import initialize_rand_states # noqa: E402 -from motrix_env_core.sim import BatchLinkPositionQuery # noqa: E402 -from motrix_envs.locomotion.wbt.cfg import ( # noqa: E402 - ActionsCfg, - RewardsCfg, - WbtEnvCfg, -) -from motrix_envs.locomotion.wbt.dex_evt import DexEvtWbtEnvCfg # noqa: E402 -from motrix_envs.locomotion.wbt.g1.common import MOTION_DIR as _G1_MOTION_DIR # noqa: E402 -from motrix_envs.locomotion.wbt.g1.common import G1WbtEnvCfg # noqa: E402 -from motrix_envs.locomotion.wbt.k1 import K1WbtEnvCfg # noqa: E402 -from motrix_envs.locomotion.wbt.mdp.command import ( # noqa: E402 - WbtMotionCommand, - WbtMotionCommandCfg, -) -from motrix_envs.locomotion.wbt.mdp.observations import ( # noqa: E402 - DofPosRelObsCfg, - DofVelObsCfg, - MotionReferenceOrientationObsCfg, -) -from motrix_envs.locomotion.wbt.mdp.reset import ( # noqa: E402 - BodyDofPosResetCfg, - BodyLinVelResetCfg, - BodyPosResetCfg, - BodyRotResetCfg, - BodyRotVelResetCfg, -) - -# Keep WBT kernels in one suite: deselecting only the first expensive tests -# shifts first-use compilation to later tests instead of removing its cost. -pytestmark = pytest.mark.numba - - -def _motion_command(env: ManagerEnv) -> WbtMotionCommand: - command = env.command_terms["motion"] - assert isinstance(command, WbtMotionCommand) - return command - - -def _deterministic_manager_cfg(*, hold_at_clip_end: bool = False) -> WbtEnvCfg: - cfg = registry.make_env_config("g1-wbt-dance", mode="play") - assert isinstance(cfg, WbtEnvCfg) - motion = cfg.commands.motion - assert isinstance(motion, WbtMotionCommandCfg) - policy = cfg.observations.policy - dof_pos = policy.dof_pos - dof_vel = policy.dof_vel - base_ang_vel = policy.base_ang_vel - motion_ref_ori = policy.motion_ref_ori_b - assert isinstance(dof_pos, DofPosRelObsCfg) - assert isinstance(dof_vel, DofVelObsCfg) - assert isinstance(base_ang_vel, BodyAngularVelocityObsCfg) - assert isinstance(motion_ref_ori, MotionReferenceOrientationObsCfg) - return replace( - cfg, - observations=replace( - cfg.observations, - policy=replace( - policy, - dof_pos=replace(dof_pos, noise=UniformNoiseCfg()), - dof_vel=replace(dof_vel, noise=UniformNoiseCfg()), - base_ang_vel=replace(base_ang_vel, noise=UniformNoiseCfg()), - motion_ref_ori_b=replace(motion_ref_ori, noise=UniformNoiseCfg()), - ), - ), - commands=replace( - cfg.commands, - motion=replace( - motion, - hold_at_clip_end=hold_at_clip_end, - # Pin the head-frame start so wrap/buffer mechanics stay - # deterministic regardless of the play start distribution. - start_at_timestep_zero_prob=1.0, - sequential_clips=False, - ), - ), - ) - - -def _assert_float_array_equal(actual: np.ndarray, expected: np.ndarray) -> None: - assert actual.dtype == expected.dtype - np.testing.assert_allclose(actual, expected, rtol=1e-5, atol=1e-6, equal_nan=True) - - -def _make_numba_env(cfg: WbtEnvCfg, *, num_envs: int, seed: int = 1) -> ManagerEnv: - return ManagerEnv(cfg, num_envs=num_envs, seed=seed) - - -def _motion_ref_ori_term(env: ManagerEnv): - return next(entry.term for entry in env.observation_groups["policy"].terms if entry.name == "motion_ref_ori_b") - - -def _policy_observation_noise(env: ManagerEnv, state, name: str) -> np.ndarray: - policy_layout = next(term for term in env.manager_layout.observations["policy"].terms if term.name == name) - value_layout = next(term for term in env.manager_layout.observations["value"].terms if term.name == name) - assert state.obs.value is not None - return state.obs.policy[:, policy_layout.output_slice] - state.obs.value[:, value_layout.output_slice] - - -def _motion_ref_ori_noise(env: ManagerEnv, state) -> np.ndarray: - return _policy_observation_noise(env, state, "motion_ref_ori_b") - - -def _motion_ref_ori_noise_amplitude(env: ManagerEnv) -> float: - term = env.cfg.observations.policy.motion_ref_ori_b - assert isinstance(term, MotionReferenceOrientationObsCfg) - return term.noise.amplitude - - -def _execute_task_kernel(env: ManagerEnv, state) -> None: - env._refresh_sim_reads() - env._execute_task_kernel(env._kernel_inputs) - - -def test_wbt_reward_terms_are_typed_manager_group() -> None: - cfg = _deterministic_manager_cfg() - rewards = cfg.rewards - assert isinstance(rewards, RewardsCfg) - assert all(isinstance(getattr(rewards, field.name), RewardTermCfg) for field in fields(rewards)) - - -def test_numba_wbt_read_plan_reuses_preallocated_arrays() -> None: - env = _make_numba_env(_deterministic_manager_cfg(), num_envs=4) - state = env.init_state() - assert env._compiled_manager_program is not None - assert isinstance(env._rand, RandValue) - motion_joint_entries = [ - entry for group in env.observation_groups.values() for entry in group.terms if entry.name == "motion_joint" - ] - assert len(motion_joint_entries) == 2 - motion = _motion_command(env) - assert all(entry.size == motion.command.shape[1] for entry in motion_joint_entries) - first = env._kernel_inputs - env._refresh_sim_reads() - second = env._kernel_inputs - - assert len(first) == len(env.manager_layout.inputs) - sources = env._compiled_manager_program.read_plan.sources - # One source per kernel-data term plus one runtime scalar buffer per float - # term argument. Scalar args must stay runtime-buffered (not baked into the - # plan/source): the deterministic config has float term args, so require at - # least one scalar-buffer source, and the manager context source either way. - scalar_buffer_sources = [source for source in sources if source.value_type is TermScalarBuffer] - assert scalar_buffer_sources - assert any(source.value_type is ManagerContext for source in sources) - context_source = next(source for source in sources if source.value_type is ManagerContext) - arrays = tuple(value for value in context_source.values if isinstance(value, np.ndarray)) - assert any(not array.flags.writeable for array in arrays) - assert motion.steps.shape == (env.num_envs, 1) - assert motion.clip.joint_pos.shape[0] == motion.clip.tracked_bodies_pos_w.shape[0] - assert env.sim_data["robot_dof_pos"].shape[0] == env.num_envs - assert first is second - for first_value, second_value in zip(first, second, strict=True): - assert first_value is second_value or np.shares_memory(first_value, second_value) or first_value.size == 0 - - action = env.action_terms["joint_position"].state - actions = np.full((env.num_envs, *env.action_space.shape), 0.25, dtype=np.float32) - env.apply_action(actions, state) - partial = env._compiled_manager_program.read_plan.read( - env, - np.asarray([1, 3], dtype=np.int64), - ) - assert partial is first - assert any(value is action.action_queue for value in first) - np.testing.assert_array_equal(action.action_queue[:, action.action_ptr[0]], actions) - - state.terminated[:] = [False, True, False, True] - env._reset_done_envs() - env._refresh_sim_reads() - assert env._kernel_inputs is first - np.testing.assert_array_equal(action.action_queue[[0, 2], action.action_ptr[0]], 0.25) - np.testing.assert_array_equal(action.action_queue[[1, 3], action.action_ptr[0]], 0.0) - - -@pytest.mark.slow -def test_numba_wbt_step_preserves_previous_actor_and_critic_observations() -> None: - env = _make_numba_env(_deterministic_manager_cfg(), num_envs=2) - initial = env.init_state() - assert initial.obs.value is not None - initial_policy = initial.obs.policy - initial_value = initial.obs.value - initial_policy_snapshot = initial_policy.copy() - initial_value_snapshot = initial_value.copy() - - task = env._task_program - assert task is not None - kernels = (task.evaluate_kernel, task.observe_kernel, task.reset_kernel) - signatures = tuple(tuple(kernel.signatures) for kernel in kernels) - assert all(len(items) == 1 for items in signatures) - env.warmup() - first = env.step(np.zeros((env.num_envs, *env.action_space.shape), dtype=np.float32)) - assert tuple(tuple(kernel.signatures) for kernel in kernels) == signatures - second_env = _make_numba_env(_deterministic_manager_cfg(), num_envs=2) - second_env.init_state() - assert second_env.manager_layout.plan_keys == env.manager_layout.plan_keys - second_task = second_env._task_program - assert second_task is not None - assert second_task.evaluate_kernel is task.evaluate_kernel - assert second_task.observe_kernel is task.observe_kernel - assert second_task.reset_kernel is task.reset_kernel - - np.testing.assert_array_equal(initial_policy, initial_policy_snapshot) - np.testing.assert_array_equal(initial_value, initial_value_snapshot) - assert not np.shares_memory(initial_policy, first.obs.policy) - assert first.obs.value is not None - assert not np.shares_memory(initial_value, first.obs.value) - - -@pytest.mark.slow -def test_numba_wbt_observation_noise_bounds_and_determinism() -> None: - """Observation noise is uniform(-amp, amp), zero for zero amplitudes, and - reproducible for a fresh environment with the same seed and episode step.""" - cfg = registry.make_env_config("g1-wbt-dance", mode="play") - assert isinstance(cfg, WbtEnvCfg) - env = _make_numba_env(cfg, num_envs=8) - state = env.init_state() - env.warmup() - term = _motion_ref_ori_term(env) - amplitude = _motion_ref_ori_noise_amplitude(env) - assert term.args == (np.float32(amplitude),) - env._refresh_sim_reads() - rng_states = env._rand.state - assert rng_states.shape == (env.num_envs, 1) - assert rng_states.dtype == np.uint64 - _execute_task_kernel(env, state) - - noise = _motion_ref_ori_noise(env, state) - assert noise.shape == (env.num_envs, 6) - assert noise.dtype == np.float32 - assert np.all(noise >= -amplitude) - assert np.all(noise <= amplitude) - assert np.any(noise != 0.0) - - # Fresh environment with the same construction sequence reproduces noise. - cfg2 = registry.make_env_config("g1-wbt-dance", mode="play") - assert isinstance(cfg2, WbtEnvCfg) - env2 = _make_numba_env(cfg2, num_envs=8) - state2 = env2.init_state() - env2.warmup() - _execute_task_kernel(env2, state2) - np.testing.assert_array_equal(_motion_ref_ori_noise(env2, state2), noise) - - -def test_numba_wbt_observation_noise_is_seeded_and_thread_count_independent() -> None: - base_cfg = registry.make_env_config("g1-wbt-dance", mode="play") - assert isinstance(base_cfg, WbtEnvCfg) - original_threads = numba.get_num_threads() - parallel_threads = min(4, original_threads) - environments = [_make_numba_env(base_cfg, num_envs=32, seed=seed) for seed in (123, 123, 124)] - states = [env.init_state() for env in environments] - for env in environments: - env.warmup() - - try: - numba.set_num_threads(1) - _execute_task_kernel(environments[0], states[0]) - - numba.set_num_threads(parallel_threads) - _execute_task_kernel(environments[1], states[1]) - _execute_task_kernel(environments[2], states[2]) - serial = _motion_ref_ori_noise(environments[0], states[0]).copy() - parallel = _motion_ref_ori_noise(environments[1], states[1]).copy() - different_seed = _motion_ref_ori_noise(environments[2], states[2]).copy() - - states[0].terminated.fill(True) - states[1].terminated.fill(True) - numba.set_num_threads(1) - environments[0]._reset_done_envs() - serial_reset = _motion_ref_ori_noise(environments[0], states[0]).copy() - numba.set_num_threads(parallel_threads) - environments[1]._reset_done_envs() - parallel_reset = _motion_ref_ori_noise(environments[1], states[1]).copy() - finally: - numba.set_num_threads(original_threads) - - np.testing.assert_array_equal(parallel, serial) - np.testing.assert_array_equal(parallel_reset, serial_reset) - assert np.any(different_seed != serial) - - -def test_numba_wbt_observation_noise_sequence_advances_during_reset() -> None: - cfg = registry.make_env_config("g1-wbt-dance", mode="play") - assert isinstance(cfg, WbtEnvCfg) - env = _make_numba_env(cfg, num_envs=4) - state = env.init_state() - first_episode = _motion_ref_ori_noise(env, state).copy() - rng_states = env._rand.state - states_before_reset = rng_states.copy() - - state.terminated.fill(True) - env._refresh_sim_reads() - env._reset_done_envs() - env.compute_observation(state) - - assert np.any(rng_states != states_before_reset) - assert np.any(_motion_ref_ori_noise(env, state) != first_episode) - - -def test_wbt_shared_noise_kernel_has_uniform_distribution() -> None: - # Exercise WBT's shared noise primitive through its batch kernel; manager - # integration, seeding and bounds are tested above with small simulations. - num_samples = 4096 - amplitude = np.float32(0.3) - noise = np.zeros((num_samples, 6), dtype=np.float32) - rng_states = initialize_rand_states(num_samples, 20260809) - add_uniform_noise_batch(noise, amplitude, rng_states, np.arange(num_samples, dtype=np.int64)) - - normalized = noise.reshape(-1) / amplitude - assert np.all(normalized >= -1.0) - assert np.all(normalized < 1.0) - assert abs(float(normalized.mean())) < 0.01 - assert float(normalized.var()) == pytest.approx(1.0 / 3.0, abs=0.01) - - -@pytest.mark.parametrize("hold_at_clip_end", [False, True]) -def test_numba_wbt_clip_end_behavior(hold_at_clip_end: bool) -> None: - env = _make_numba_env(_deterministic_manager_cfg(hold_at_clip_end=hold_at_clip_end), num_envs=2) - state = env.init_state() - motion = _motion_command(env) - motion.steps.fill(motion.clip.joint_pos.shape[0] - 1) - - env.compute_transition(state) - env.compute_observation(state) - - np.testing.assert_array_equal(motion.clip_ended, True) - expected_step = motion.clip.joint_pos.shape[0] - 1 if hold_at_clip_end else 0 - np.testing.assert_array_equal(motion.steps[:, 0], expected_step) - - -def test_numba_wbt_clip_wrap_rematerializes_sim_only() -> None: - # The single dance file keeps the wrap frame adjacent to the episode-start - # pose, so forcing a wrap there cannot diverge the physics; the corpus - # directory's final frame belongs to a different clip and would. - env = _make_numba_env(_single_file_cfg(start_at_timestep_zero_prob=1.0), num_envs=2, seed=11) - env.init_state() - motion = _motion_command(env) - action = env.action_terms["joint_position"].state - assert isinstance(action, JointPositionActionState) - num_frames = motion.clip.joint_pos.shape[0] - # 0.25 is exactly representable in float32, so equality checks stay exact. - actions = np.full((env.num_envs, *env.action_space.shape), 0.25, dtype=np.float32) - - # Prime persistent action state and the episode counter. - env.step(actions) - np.testing.assert_array_equal(action.action_queue[:, action.action_ptr[0]], 0.25) - episode_steps_before_wrap = env.state.episode_steps.copy() - - # Force every lane to wrap the clip on the next transition. - motion.steps[:, 0] = num_frames - 1 - state = env.step(actions) - - # The wrap lanes were rematerialized: frames resampled within range. - np.testing.assert_array_equal(motion.clip_ended[:, 0], True) - assert np.all(motion.steps[:, 0] <= num_frames - 2) - # Rematerialization is not an episode boundary. - np.testing.assert_array_equal(state.terminated, False) - np.testing.assert_array_equal(state.truncated, False) - np.testing.assert_array_equal(state.episode_steps, episode_steps_before_wrap + 1) - # Sim-only: the persistent action state (which an action reset would - # zero) keeps the raw policy actions of this step. - np.testing.assert_array_equal(action.action_queue[:, action.action_ptr[0]], 0.25) - np.testing.assert_array_equal( - action.action_queue[:, (action.action_ptr[0] - 1) % action.action_queue.shape[1]], 0.25 - ) - - # The request flag is cleared before the next physics step: a follow-up - # step neither rematerializes the lane nor rewinds its frame. - steps_after_wrap = motion.steps[:, 0].copy() - next_state = env.step(actions) - np.testing.assert_array_equal(env._sim_reset_requested[:, 0], False) - np.testing.assert_array_equal(motion.steps[:, 0], steps_after_wrap + 1) - np.testing.assert_array_equal(next_state.episode_steps, episode_steps_before_wrap + 2) - - -def test_numba_wbt_adaptive_sampler_state_is_manager_owned() -> None: - cfg = _deterministic_manager_cfg() - command_cfg = cfg.commands.motion - assert isinstance(command_cfg, WbtMotionCommandCfg) - command_cfg = replace( - command_cfg, - alpha=1.0, - kernel_size=1, - ) - command_cfg = replace(command_cfg, adaptive_sampling_enabled=True) - cfg = replace( - cfg, - commands=replace(cfg.commands, motion=command_cfg), - ) - env = _make_numba_env(cfg, num_envs=3) - state = env.init_state() - motion = _motion_command(env) - expected_num_bins = motion.clip.joint_pos.shape[0] // round(1.0 / cfg.ctrl_dt) + 1 - assert motion.adaptive_bin_failed_count.size == expected_num_bins - - motion.steps[:, 0] = [0, motion.clip.joint_pos.shape[0] // 2, motion.clip.joint_pos.shape[0] - 1] - motion.adaptive_bin_failed_count.fill(0.0) - motion.adaptive_current_bin_failed_count.fill(0.0) - state.terminated[:] = [False, True, True] - from motrix_env_core.numba.manager.commands import ResetContext - - motion.reset( - ResetContext( - env_ids=np.arange(env.num_envs, dtype=np.int64), - terminated=state.terminated, - metrics=state.metrics, - ) - ) - motion.on_transition() - - assert np.sum(motion.adaptive_bin_failed_count) == pytest.approx(2.0) - np.testing.assert_array_equal(motion.adaptive_current_bin_failed_count, 0.0) - assert {metric.name for metric in env.manager_layout.metrics} >= {"motion_step", "clip_ended"} - assert isinstance(state.metrics["adaptive_sampling_entropy"], float) - assert isinstance(state.metrics["adaptive_sampling_top1_prob"], float) - assert isinstance(state.metrics["adaptive_sampling_top1_bin"], float) - assert isinstance(state.metrics["adaptive_failure_mass"], float) - - -def test_numba_wbt_masked_reset_preserves_bound_buffer_identity() -> None: - env = _make_numba_env(_deterministic_manager_cfg(), num_envs=3) - state = env.init_state() - buffers = env._kernel_buffers - assert buffers is not None - env._refresh_sim_reads() - action_value = env.action_terms["joint_position"].state - motion = _motion_command(env) - identities = { - "action_queue": id(action_value.action_queue), - "action_ptr": id(action_value.action_ptr), - "reward_terms": id(buffers[0]), - "termination_masks": id(buffers[2]), - "target_body_position_relative": id(motion.target_body_position_relative), - "sim_inputs": tuple(id(env.sim_data[key]) for key in env.sim_data.keys), - } - env.apply_action(np.ones_like(action_value.action_queue[:, 0]), state) - motion.steps[:, 0] = [1, 2, 3] - robot_dof_pos = env.sim_data["robot_dof_pos"] - non_reset_dof_pos = robot_dof_pos[1].copy() - non_reset_policy = state.obs.policy[1].copy() - assert state.obs.value is not None - non_reset_value = state.obs.value[1].copy() - state.reward[:] = [1.0, 2.0, 3.0] - state.terminated[:] = [True, False, True] - reward = state.reward.copy() - terminated = state.terminated.copy() - - env._reset_done_envs() - - assert env._task_program is not None - assert env._task_program.reset_kernel.nopython_signatures - assert id(action_value.action_queue) == identities["action_queue"] - assert id(action_value.action_ptr) == identities["action_ptr"] - assert id(buffers[0]) == identities["reward_terms"] - assert id(buffers[2]) == identities["termination_masks"] - assert id(motion.target_body_position_relative) == identities["target_body_position_relative"] - assert tuple(id(env.sim_data[key]) for key in env.sim_data.keys) == identities["sim_inputs"] - ptr = int(action_value.action_ptr[0]) - previous_ptr = (ptr - 1) % action_value.action_queue.shape[1] - np.testing.assert_array_equal(action_value.action_queue[[0, 2], ptr], 0.0) - np.testing.assert_array_equal(action_value.action_queue[1, ptr], 1.0) - np.testing.assert_array_equal(action_value.action_queue[[0, 2], previous_ptr], 0.0) - np.testing.assert_array_equal(motion.steps[1, 0], 2) - np.testing.assert_allclose(robot_dof_pos[[0, 2]], motion.clip.joint_pos[motion.steps[[0, 2], 0]]) - np.testing.assert_array_equal(robot_dof_pos[1], non_reset_dof_pos) - np.testing.assert_array_equal(state.obs.policy[1], non_reset_policy) - np.testing.assert_array_equal(state.obs.value[1], non_reset_value) - np.testing.assert_array_equal(state.reward, reward) - np.testing.assert_array_equal(state.terminated, terminated) - - -def test_numba_wbt_manager_rolls_and_resets_action_buffers() -> None: - env = _make_numba_env(_deterministic_manager_cfg(), num_envs=2) - state = env.init_state() - term = env.action_terms["joint_position"] - assert term is env._action_terms["joint_position"] - value = term.state - - actions = np.ones((env.num_envs, *env.action_space.shape), dtype=np.float32) - env.apply_action(actions, state) - ptr = int(value.action_ptr[0]) - previous_ptr = (ptr - 1) % value.action_queue.shape[1] - np.testing.assert_array_equal(value.action_queue[:, ptr], 1.0) - np.testing.assert_array_equal(value.action_queue[:, previous_ptr], 0.0) - - actions.fill(2.0) - env.apply_action(actions, state) - ptr = int(value.action_ptr[0]) - previous_ptr = (ptr - 1) % value.action_queue.shape[1] - np.testing.assert_array_equal(value.action_queue[:, ptr], 2.0) - np.testing.assert_array_equal(value.action_queue[:, previous_ptr], 1.0) - - state.terminated[:] = True - ptr_before_reset = value.action_ptr.copy() - env._reset_done_envs() - np.testing.assert_array_equal(value.action_queue, 0.0) - np.testing.assert_array_equal(value.action_ptr, ptr_before_reset) - - -def test_numba_wbt_action_owns_shared_writable_model_data() -> None: - env = _make_numba_env(_deterministic_manager_cfg(), num_envs=1) - state = SimpleNamespace() - actions = np.full((1, *env.action_space.shape), 0.25, dtype=np.float32) - value = env.action_terms["joint_position"].state - - env.apply_action(actions, state) - - ptr = int(value.action_ptr[0]) - previous_ptr = (ptr - 1) % value.action_queue.shape[1] - np.testing.assert_array_equal(value.action_queue[:, ptr], actions) - np.testing.assert_array_equal(value.action_queue[:, previous_ptr], 0.0) - assert value.action_queue.flags.writeable - assert value.action_ptr.flags.writeable - assert all( - array.flags.writeable - for array in ( - value.default_angles, - value.joint_lower, - value.joint_upper, - value.action_scales, - ) - ) - np.testing.assert_allclose( - env._action_writes.buffer("joint_position"), - actions * value.action_scales + value.default_angles, - ) - - -def test_manager_rejects_overlapping_wbt_action_routes() -> None: - cfg = _deterministic_manager_cfg() - action = cfg.actions.joint_position - cfg = replace(cfg, actions={"first": action, "second": action}) - - with pytest.raises(ValueError, match="both control actuator"): - _make_numba_env(cfg, num_envs=1) - - -def test_numba_wbt_registry_uses_generic_manager_env() -> None: - manager_cfg = registry.make_env_config("g1-wbt-dance", mode="play") - env = registry.make("g1-wbt-dance", mode="play", num_envs=1) - - assert isinstance(manager_cfg, WbtEnvCfg) - motion_command_cfg = manager_cfg.commands.motion - assert isinstance(motion_command_cfg, WbtMotionCommandCfg) - action_cfg = manager_cfg.actions.joint_position - assert isinstance(action_cfg, JointPositionActionCfg) - tracked_body_pos = env.sim_data.query("tracked_body_pos") - assert isinstance(tracked_body_pos, BatchLinkPositionQuery) - assert motion_command_cfg.tracked_body_names == tracked_body_pos.links - assert motion_command_cfg.reference_body_name in motion_command_cfg.tracked_body_names - assert isinstance(manager_cfg.actions, ActionsCfg) - assert isinstance(manager_cfg.actions.joint_position, JointPositionActionCfg) - assert isinstance(manager_cfg.sim_reset, ManagerResetCfg) - assert isinstance(manager_cfg.sim_reset.body_pos, BodyPosResetCfg) - assert isinstance(manager_cfg.sim_reset.body_rot, BodyRotResetCfg) - assert isinstance(manager_cfg.sim_reset.body_lin_vel, BodyLinVelResetCfg) - assert isinstance(manager_cfg.sim_reset.body_rot_vel, BodyRotVelResetCfg) - assert isinstance(manager_cfg.sim_reset.body_dof_pos, BodyDofPosResetCfg) - assert isinstance(env, ManagerEnv) - assert isinstance(_motion_command(env), WbtMotionCommand) - - -def test_wbt_manager_play_disables_each_reset_term_noise() -> None: - cfg = registry.make_env_config("g1-wbt-dance", mode="play") - - assert all(term.noise_scale == 0.0 for term in cfg.sim_reset.to_dict().values()) - - -def test_wbt_manager_play_enables_sequential_playback() -> None: - """Play walks the motion source as one ordered sequence: start at the - corpus head, cross clip boundaries in order, and at the corpus end loop - back to the head — hold_at_clip_end is cleared so playback loops the - corpus instead of freezing on the final frame.""" - for name in ("g1-29dof-wbt-largebox", "g1-wbt-dance"): - cfg = registry.make_env_config(name, mode="play") - assert isinstance(cfg, WbtEnvCfg) - assert cfg.commands.motion.sequential_clips - assert not cfg.commands.motion.hold_at_clip_end - assert not cfg.commands.motion.adaptive_sampling_enabled - - -def test_wbt_robot_config_subclasses_isolate_nested_overrides() -> None: - g1 = registry.make_env_config("g1-wbt-dance") - another_g1 = registry.make_env_config("g1-wbt-dance") - k1 = registry.make_env_config("k1-wbt-freekick") - dex_evt = registry.make_env_config("dex-evt-wbt-dance") - - assert isinstance(g1, G1WbtEnvCfg) - assert isinstance(another_g1, G1WbtEnvCfg) - assert isinstance(k1, K1WbtEnvCfg) - assert isinstance(dex_evt, DexEvtWbtEnvCfg) - assert g1.commands.motion.joint_names - assert g1.commands.motion.tracked_body_names - assert g1.commands.motion.reference_body_name == "torso_link" - assert k1.commands.motion.reference_body_name == "Trunk" - assert dex_evt.commands.motion.reference_body_name == "waist_pitch_link" - assert g1.queries.model["actuator_kp"] == k1.queries.model["actuator_kp"] - assert g1.queries.model["actuator_kp"] == dex_evt.queries.model["actuator_kp"] - assert g1.queries.model["robot_joint_position_limits"].body == "pelvis" - assert k1.queries.model["robot_joint_position_limits"].body == "Trunk" - assert dex_evt.queries.model["robot_joint_position_limits"].body == "pelvis" - assert g1.queries is not another_g1.queries - assert g1.queries.data is not another_g1.queries.data - assert g1.queries.model is not another_g1.queries.model - assert g1.rewards is not another_g1.rewards - assert g1.commands.motion is not another_g1.commands.motion - - -def test_wbt_queries_use_authoritative_motion_orders() -> None: - cfg = registry.make_env_config("g1-wbt-dance") - assert isinstance(cfg, G1WbtEnvCfg) - joint_names = cfg.commands.motion.joint_names - tracked_body_names = cfg.commands.motion.tracked_body_names - assert cfg.queries.model["robot_joint_position_limits"].body == cfg.scene.objs.robot.resolved_base_link_name - assert cfg.queries.data["undesired_contact_forces"].body == cfg.scene.objs.robot.resolved_base_link_name - assert cfg.queries.data["robot_dof_pos"].joints == joint_names - assert cfg.queries.data["robot_dof_vel"].joints == joint_names - assert cfg.queries.data["tracked_body_pos"].links == tracked_body_names - assert cfg.queries.data["tracked_body_quat"].links == tracked_body_names - assert cfg.queries.data["tracked_body_linear_velocity"].links == tracked_body_names - assert cfg.queries.data["tracked_body_angular_velocity"].links == tracked_body_names - - -def test_wbt_build_spec_is_spawn_pickle_safe() -> None: - spec = registry.resolve("g1-wbt-dance") - - pickle.dumps(spec) - assert spec.env_cfg.queries.data["robot_dof_pos"].joints == spec.env_cfg.commands.motion.joint_names - assert spec.env_cfg.queries.data["tracked_body_pos"].links == spec.env_cfg.commands.motion.tracked_body_names - - -@pytest.mark.parametrize( - "env_name", - ["dex-evt-wbt-dance", "g1-29dof-wbt-largebox", "g1-wbt-dance", "k1-wbt-freekick"], -) -def test_numba_wbt_manager_builds_for_all_wbt_presets(env_name: str) -> None: - env = registry.make(env_name, mode="play", num_envs=2) - assert isinstance(env, ManagerEnv) - assert env.cfg.queries.data["robot_dof_pos"].joints == env.cfg.commands.motion.joint_names - assert env.sim_data.query("robot_dof_pos").joints == env.cfg.commands.motion.joint_names - env.init_state() - action = env.action_terms["joint_position"].state - assert isinstance(action, JointPositionActionState) - joint_lower, joint_upper = env.model.others["robot_joint_position_limits"] - np.testing.assert_array_equal(action.joint_lower, joint_lower) - np.testing.assert_array_equal(action.joint_upper, joint_upper) - assert action.joint_lower.shape == env.sim_data["robot_dof_pos"].shape[1:] - - policy_terms = env.manager_layout.observations["policy"].terms - value_terms = env.manager_layout.observations["value"].terms - assert [term.name for term in policy_terms] == [ - "motion_joint", - "motion_ref_ori_b", - "base_ang_vel", - "dof_pos", - "dof_vel", - "actions", - ] - assert [term.name for term in value_terms] == [ - "motion_joint", - "motion_ref_pos_b", - "motion_ref_ori_b", - "robot_body_pos_b", - "robot_body_ori_b", - "base_lin_vel", - "base_ang_vel", - "dof_pos", - "dof_vel", - "actions", - ] - observation_space = env.observation_space - assert observation_space.value is not None - for terms, width in ( - (policy_terms, observation_space.policy.shape[0]), - (value_terms, observation_space.value.shape[0]), - ): - assert terms[0].output_slice.start == 0 - assert all(left.output_slice.stop == right.output_slice.start for left, right in zip(terms, terms[1:])) - assert terms[-1].output_slice.stop == width - - assert [term.name for term in env.manager_layout.rewards] == [field.name for field in fields(env.cfg.rewards)] - assert [term.name for term in env.manager_layout.terminations] == [ - "bad_ref_z", - "bad_ref_ori", - "bad_body_z", - "bad_dof_pos", - "bad_dof_vel", - ] - - -def test_g1_wbt_dance_mgr_uses_manager_environment() -> None: - spec = registry.resolve("g1-wbt-dance") - - assert spec.env_cls is ManagerEnv - assert isinstance(spec.env_cfg, WbtEnvCfg) - - -# --------------------------------------------------------------------------- -# Multi-clip corpus (MotionLibrary) behavior -# --------------------------------------------------------------------------- - -_DANCE_MOTION = _G1_MOTION_DIR / "dance" / "dance1_subject2.npz" -_DANCE_SPLIT = 500 - - -def _split_dance_corpus(tmp_path, *, split=_DANCE_SPLIT, extension_channels=()): - """Split the bundled G1 dance motion into two schema v1 files. - - Slicing preserves every per-frame field, so the concatenated corpus - reproduces the original clip bit-for-bit on the global frame axis. - """ - with np.load(_DANCE_MOTION, allow_pickle=False) as data: - fields = {key: data[key] for key in data.files} - total = int(np.asarray(fields["num_frames"]).reshape(-1)[0]) - rng = np.random.default_rng(7) - paths = [] - bounds = (0, split, total) - for index, (start, stop) in enumerate(zip(bounds[:-1], bounds[1:])): - part = {"num_frames": np.int32(stop - start)} - for name, value in fields.items(): - if name == "num_frames": - continue - # Per-frame arrays (including ext_ channels) carry the frame axis; - # names and scalars are corpus-wide metadata. - if isinstance(value, np.ndarray) and value.shape[:1] == (total,): - part[name] = value[start:stop] - else: - part[name] = value - for channel in extension_channels: - part[f"ext_{channel}"] = rng.standard_normal((stop - start, 3)).astype(np.float32) - path = tmp_path / f"part{index}.npz" - np.savez(path, **part) - paths.append(str(path)) - return tuple(paths) - - -def _multi_clip_cfg( - tmp_path, - *, - hold_at_clip_end: bool = False, - sequential_clips: bool = False, - start_at_timestep_zero_prob: float | None = None, - adaptive_sampling_enabled: bool | None = None, - alpha: float | None = None, - split: int = _DANCE_SPLIT, - extension_channels: tuple[str, ...] = (), - motion_files: tuple[str, ...] | None = None, -) -> WbtEnvCfg: - cfg = registry.make_env_config("g1-wbt-dance", mode="play") - assert isinstance(cfg, WbtEnvCfg) - motion = cfg.commands.motion - assert isinstance(motion, WbtMotionCommandCfg) - replacements = { - "motion_files": motion_files - if motion_files is not None - else _split_dance_corpus(tmp_path, split=split, extension_channels=extension_channels), - "extension_channels": tuple(extension_channels), - "hold_at_clip_end": hold_at_clip_end, - "sequential_clips": sequential_clips, - } - if start_at_timestep_zero_prob is not None: - replacements["start_at_timestep_zero_prob"] = start_at_timestep_zero_prob - if adaptive_sampling_enabled is not None: - replacements["adaptive_sampling_enabled"] = adaptive_sampling_enabled - if alpha is not None: - replacements["alpha"] = alpha - return replace(cfg, commands=replace(cfg.commands, motion=replace(motion, **replacements))) - - -def _single_file_cfg(*, start_at_timestep_zero_prob: float | None = None) -> WbtEnvCfg: - """A play config pinned to the bundled dance file. - - The g1-wbt-dance corpus directory holds however many clips are currently - dropped into it, so tests that need one exact clip load the file directly - instead of going through the registry preset. - """ - cfg = registry.make_env_config("g1-wbt-dance", mode="play") - assert isinstance(cfg, WbtEnvCfg) - motion = cfg.commands.motion - assert isinstance(motion, WbtMotionCommandCfg) - replacements: dict[str, object] = { - "motion_files": (str(_DANCE_MOTION),), - "hold_at_clip_end": False, - "sequential_clips": False, - } - if start_at_timestep_zero_prob is not None: - replacements["start_at_timestep_zero_prob"] = start_at_timestep_zero_prob - return replace(cfg, commands=replace(cfg.commands, motion=replace(motion, **replacements))) - - -def test_numba_wbt_multi_clip_wrap_rematerializes_sim_only(tmp_path) -> None: - # Split after two frames: the first clip's final frame is adjacent to the - # episode-start pose, so forcing a wrap there cannot diverge the physics. - # The head-frame start keeps the priming step on that same pose. - cfg = _multi_clip_cfg(tmp_path, split=2, start_at_timestep_zero_prob=1.0) - env = _make_numba_env(cfg, num_envs=2, seed=11) - env.init_state() - motion = _motion_command(env) - action = env.action_terms["joint_position"].state - assert isinstance(action, JointPositionActionState) - total = motion.clip.joint_pos.shape[0] - actions = np.full((env.num_envs, *env.action_space.shape), 0.25, dtype=np.float32) - - env.step(actions) - episode_steps_before_wrap = env.state.episode_steps.copy() - - # A lane inside the second clip advances normally: no wrap, no rematerialize. - motion.steps[:, 0] = 2 - state = env.step(actions) - np.testing.assert_array_equal(motion.clip_ended[:, 0], False) - np.testing.assert_array_equal(motion.steps[:, 0], 3) - np.testing.assert_array_equal(state.terminated, False) - np.testing.assert_array_equal(env._sim_reset_requested[:, 0], False) - - # A lane on the first clip's final frame wraps at the clip boundary, not at - # the corpus end: the frame is resampled over the whole corpus and the reset - # pipeline rematerializes the lane without ending the episode. - motion.steps[:, 0] = 1 - state = env.step(actions) - np.testing.assert_array_equal(motion.clip_ended[:, 0], True) - assert np.all(motion.steps[:, 0] <= total - 2) - np.testing.assert_array_equal(state.terminated, False) - np.testing.assert_array_equal(state.truncated, False) - np.testing.assert_array_equal(state.episode_steps, episode_steps_before_wrap + 2) - np.testing.assert_array_equal(action.action_queue[:, action.action_ptr[0]], 0.25) - np.testing.assert_array_equal( - action.action_queue[:, (action.action_ptr[0] - 1) % action.action_queue.shape[1]], 0.25 - ) - - steps_after_wrap = motion.steps[:, 0].copy() - env.step(actions) - np.testing.assert_array_equal(env._sim_reset_requested[:, 0], False) - np.testing.assert_array_equal(motion.steps[:, 0], steps_after_wrap + 1) - - -def test_numba_wbt_multi_clip_hold_clamps_to_own_clip_end(tmp_path) -> None: - env = _make_numba_env(_multi_clip_cfg(tmp_path, hold_at_clip_end=True), num_envs=2) - state = env.init_state() - motion = _motion_command(env) - total = motion.clip.joint_pos.shape[0] - - # Hold on the first clip keeps the lane on that clip's final frame, not on - # the corpus final frame. - motion.steps.fill(_DANCE_SPLIT - 1) - env.compute_transition(state) - np.testing.assert_array_equal(motion.clip_ended[:, 0], True) - np.testing.assert_array_equal(motion.steps[:, 0], _DANCE_SPLIT - 1) - - # The corpus-final clip still holds on the global final frame. - motion.steps.fill(total - 1) - env.compute_transition(state) - np.testing.assert_array_equal(motion.clip_ended[:, 0], True) - np.testing.assert_array_equal(motion.steps[:, 0], total - 1) - - -def test_numba_wbt_multi_clip_sequential_crosses_boundaries_and_loops(tmp_path) -> None: - # Split after two frames so boundary crossings stay adjacent to the - # episode-start pose and cannot diverge the physics (same trick as the - # wrap test above). - cfg = _multi_clip_cfg(tmp_path, sequential_clips=True, split=2) - env = _make_numba_env(cfg, num_envs=2, seed=11) - env.init_state() - motion = _motion_command(env) - action = env.action_terms["joint_position"].state - assert isinstance(action, JointPositionActionState) - total = motion.clip.joint_pos.shape[0] - actions = np.full((env.num_envs, *env.action_space.shape), 0.25, dtype=np.float32) - - # The initial reset keeps the initialized corpus-head frame. - np.testing.assert_array_equal(motion.steps[:, 0], 0) - - env.step(actions) - episode_steps_before_boundary = env.state.episode_steps.copy() - - # A lane inside the second clip advances normally: no boundary crossed. - motion.steps[:, 0] = 2 - state = env.step(actions) - np.testing.assert_array_equal(motion.steps[:, 0], 3) - np.testing.assert_array_equal(motion.clip_ended[:, 0], False) - np.testing.assert_array_equal(env._sim_reset_requested[:, 0], False) - - # The first clip's final frame crosses into the second clip's head frame - # instead of resampling, rematerializing the lane there. - motion.steps[:, 0] = 1 - state = env.step(actions) - np.testing.assert_array_equal(motion.clip_ended[:, 0], True) - np.testing.assert_array_equal(motion.steps[:, 0], 2) - np.testing.assert_array_equal(state.terminated, False) - np.testing.assert_array_equal(state.truncated, False) - np.testing.assert_array_equal(state.episode_steps, episode_steps_before_boundary + 2) - # Rematerialization is sim-only: the persistent action state keeps the - # raw policy actions of this step. - np.testing.assert_array_equal(action.action_queue[:, action.action_ptr[0]], 0.25) - np.testing.assert_array_equal( - action.action_queue[:, (action.action_ptr[0] - 1) % action.action_queue.shape[1]], 0.25 - ) - - # The corpus-final frame loops back to the corpus head frame. - motion.steps[:, 0] = total - 1 - state = env.step(actions) - np.testing.assert_array_equal(motion.clip_ended[:, 0], True) - np.testing.assert_array_equal(motion.steps[:, 0], 0) - np.testing.assert_array_equal(state.terminated, False) - - -def test_numba_wbt_multi_clip_sequential_reset_keeps_frame(tmp_path) -> None: - env = _make_numba_env(_multi_clip_cfg(tmp_path, sequential_clips=True), num_envs=3) - state = env.init_state() - motion = _motion_command(env) - frames = [_DANCE_SPLIT - 1, _DANCE_SPLIT, motion.clip.joint_pos.shape[0] - 1] - - # Episode resets keep the lane on its current frame: the robot is - # rematerialized onto that reference frame and the sequence continues. - motion.steps[:, 0] = frames - env._refresh_sim_reads() - state.terminated[:] = True - env._reset_done_envs() - - np.testing.assert_array_equal(motion.steps[:, 0], frames) - np.testing.assert_allclose(env.sim_data["robot_dof_pos"], motion.clip.joint_pos[motion.steps[:, 0]]) - - -def test_numba_wbt_multi_clip_sequential_hold_stops_at_corpus_end(tmp_path) -> None: - env = _make_numba_env(_multi_clip_cfg(tmp_path, sequential_clips=True, hold_at_clip_end=True, split=2), num_envs=2) - state = env.init_state() - motion = _motion_command(env) - total = motion.clip.joint_pos.shape[0] - - # Intermediate boundaries still cross into the next clip's head frame. - motion.steps[:, 0] = 1 - env.compute_transition(state) - np.testing.assert_array_equal(motion.clip_ended[:, 0], True) - np.testing.assert_array_equal(motion.steps[:, 0], 2) - - # The corpus end holds on the final frame instead of looping back. - env._sim_reset_requested.fill(False) - motion.steps[:, 0] = total - 1 - env.compute_transition(state) - np.testing.assert_array_equal(motion.clip_ended[:, 0], True) - np.testing.assert_array_equal(motion.steps[:, 0], total - 1) - np.testing.assert_array_equal(env._sim_reset_requested[:, 0], False) - - -def test_numba_wbt_multi_clip_sampling_sequence_matches_single_clip(tmp_path) -> None: - """Same-seed start-frame draws over the split corpus match the original file.""" - - def build(cfg: WbtEnvCfg) -> list[np.ndarray]: - env = _make_numba_env(cfg, num_envs=4, seed=123) - motion = _motion_command(env) - state = env.init_state() - env.warmup() - sequence = [motion.steps[:, 0].copy()] - for _ in range(3): - state.terminated.fill(True) - env._refresh_sim_reads() - env._reset_done_envs() - sequence.append(motion.steps[:, 0].copy()) - return sequence - - single_cfg = _single_file_cfg(start_at_timestep_zero_prob=0.0) - multi_cfg = _multi_clip_cfg(tmp_path, start_at_timestep_zero_prob=0.0) - - for multi_draws, single_draws in zip( - build(multi_cfg), - build(single_cfg), - strict=True, - ): - np.testing.assert_array_equal(multi_draws, single_draws) - - -def test_numba_wbt_multi_clip_extension_channels_reach_kernel_inputs(tmp_path) -> None: - cfg = _multi_clip_cfg(tmp_path, extension_channels=("aux_reference",)) - env = _make_numba_env(cfg, num_envs=2) - env.init_state() - motion = _motion_command(env) - channel = motion.clip.extensions["aux_reference"].data - total = motion.clip.joint_pos.shape[0] - assert channel.shape == (total, 3) - assert channel.dtype == np.float32 - - env._refresh_sim_reads() - arrays = (value for value in env._kernel_inputs if isinstance(value, np.ndarray)) - assert any(np.shares_memory(value, channel) for value in arrays) - - -def test_numba_wbt_multi_clip_adaptive_bins_span_whole_corpus(tmp_path) -> None: - # alpha=1.0 folds only this update's failure counts into the histogram. - cfg = _multi_clip_cfg(tmp_path, adaptive_sampling_enabled=True, alpha=1.0) - env = _make_numba_env(cfg, num_envs=3) - state = env.init_state() - motion = _motion_command(env) - total = motion.clip.joint_pos.shape[0] - env_fps = round(1.0 / cfg.ctrl_dt) - expected_num_bins = total // env_fps + 1 - assert motion.adaptive_bin_failed_count.size == expected_num_bins - - # Failures on frames of both clips fold into global bins without going out - # of range, and the rebuilt CDF stays normalized. - motion.steps[:, 0] = [_DANCE_SPLIT - 1, _DANCE_SPLIT, total - 1] - motion.adaptive_bin_failed_count.fill(0.0) - motion.adaptive_current_bin_failed_count.fill(0.0) - state.terminated[:] = [True, True, True] - motion.reset( - ResetContext( - env_ids=np.arange(env.num_envs, dtype=np.int64), - terminated=state.terminated, - metrics=state.metrics, - ) - ) - motion.on_transition() - assert np.sum(motion.adaptive_bin_failed_count) == pytest.approx(3.0) - assert np.all(motion.sampling_cdf[:-1] <= motion.sampling_cdf[1:] + 1e-6) - assert motion.sampling_cdf[-1] == pytest.approx(1.0) - - -def test_numba_wbt_multi_clip_cfg_survives_pickle_round_trip(tmp_path) -> None: - """Async collectors receive the env spec via pickle; the assembled corpus - must survive the round-trip unchanged.""" - cfg = _multi_clip_cfg(tmp_path) - restored = pickle.loads(pickle.dumps(cfg)) - assert isinstance(restored, WbtEnvCfg) - env = _make_numba_env(restored, num_envs=2) - motion_term = _motion_command(env) - assert motion_term.clip.frame_clip_end[_DANCE_SPLIT - 1] == _DANCE_SPLIT - 1 - assert motion_term.clip.frame_clip_end[-1] == motion_term.clip.joint_pos.shape[0] - 1 diff --git a/pyproject.toml b/pyproject.toml index eb80c97f..2896af22 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -39,12 +39,6 @@ docs = [ # GPU profiles are owned by the installable RL package. Keep the workspace # installation entry points without duplicating versions or index routing. cuda = ["motrix-rl[cuda]==0.4.0"] -# TBB threading layer for numba parallel manager kernels on many-core -# machines (see motrix-env-core's tbb extra); optional, combines with any -# GPU profile. -tbb = [ - "motrix-env-core[tbb]==0.4.0", -] rocm = ["motrix-rl[rocm]==0.4.0"] [dependency-groups] @@ -136,14 +130,12 @@ version_files = [ [tool.pytest.ini_options] markers = [ "slow: tests with substantial compilation or runtime cost", - "numba: dedicated Numba kernel regression tests", "integration: full-environment or subprocess integration tests", ] [tool.mypy] files = [ "motrix_env_core/tests/typecheck/configclass.py", - "motrix_env_core/tests/typecheck/numba_env.py", ] mypy_path = "motrix_env_core/src" python_version = "3.10" diff --git a/uv.lock b/uv.lock index 6c4740e8..3938214d 100644 --- a/uv.lock +++ b/uv.lock @@ -1183,34 +1183,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/d9/3a/d68cb2b334d53fd30fac81d3a489ce4ba0d9506f4df43fcf676b68352b19/librt-0.13.0-cp313-cp313-win_arm64.whl", hash = "sha256:d4cb6fbfdf874340ab5e51450753c0f817b6958a3621125ee695bbc3de866566", size = 112127, upload-time = "2026-07-08T12:25:35.981Z" }, ] -[[package]] -name = "llvmlite" -version = "0.44.0" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/89/6a/95a3d3610d5c75293d5dbbb2a76480d5d4eeba641557b69fe90af6c5b84e/llvmlite-0.44.0.tar.gz", hash = "sha256:07667d66a5d150abed9157ab6c0b9393c9356f229784a4385c02f99e94fc94d4", size = 171880, upload-time = "2025-01-20T11:14:41.342Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/41/75/d4863ddfd8ab5f6e70f4504cf8cc37f4e986ec6910f4ef8502bb7d3c1c71/llvmlite-0.44.0-cp310-cp310-macosx_10_14_x86_64.whl", hash = "sha256:9fbadbfba8422123bab5535b293da1cf72f9f478a65645ecd73e781f962ca614", size = 28132306, upload-time = "2025-01-20T11:12:18.634Z" }, - { url = "https://files.pythonhosted.org/packages/37/d9/6e8943e1515d2f1003e8278819ec03e4e653e2eeb71e4d00de6cfe59424e/llvmlite-0.44.0-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:cccf8eb28f24840f2689fb1a45f9c0f7e582dd24e088dcf96e424834af11f791", size = 26201096, upload-time = "2025-01-20T11:12:24.544Z" }, - { url = "https://files.pythonhosted.org/packages/aa/46/8ffbc114def88cc698906bf5acab54ca9fdf9214fe04aed0e71731fb3688/llvmlite-0.44.0-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:7202b678cdf904823c764ee0fe2dfe38a76981f4c1e51715b4cb5abb6cf1d9e8", size = 42361859, upload-time = "2025-01-20T11:12:31.839Z" }, - { url = "https://files.pythonhosted.org/packages/30/1c/9366b29ab050a726af13ebaae8d0dff00c3c58562261c79c635ad4f5eb71/llvmlite-0.44.0-cp310-cp310-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:40526fb5e313d7b96bda4cbb2c85cd5374e04d80732dd36a282d72a560bb6408", size = 41184199, upload-time = "2025-01-20T11:12:40.049Z" }, - { url = "https://files.pythonhosted.org/packages/69/07/35e7c594b021ecb1938540f5bce543ddd8713cff97f71d81f021221edc1b/llvmlite-0.44.0-cp310-cp310-win_amd64.whl", hash = "sha256:41e3839150db4330e1b2716c0be3b5c4672525b4c9005e17c7597f835f351ce2", size = 30332381, upload-time = "2025-01-20T11:12:47.054Z" }, - { url = "https://files.pythonhosted.org/packages/b5/e2/86b245397052386595ad726f9742e5223d7aea999b18c518a50e96c3aca4/llvmlite-0.44.0-cp311-cp311-macosx_10_14_x86_64.whl", hash = "sha256:eed7d5f29136bda63b6d7804c279e2b72e08c952b7c5df61f45db408e0ee52f3", size = 28132305, upload-time = "2025-01-20T11:12:53.936Z" }, - { url = "https://files.pythonhosted.org/packages/ff/ec/506902dc6870249fbe2466d9cf66d531265d0f3a1157213c8f986250c033/llvmlite-0.44.0-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:ace564d9fa44bb91eb6e6d8e7754977783c68e90a471ea7ce913bff30bd62427", size = 26201090, upload-time = "2025-01-20T11:12:59.847Z" }, - { url = "https://files.pythonhosted.org/packages/99/fe/d030f1849ebb1f394bb3f7adad5e729b634fb100515594aca25c354ffc62/llvmlite-0.44.0-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:c5d22c3bfc842668168a786af4205ec8e3ad29fb1bc03fd11fd48460d0df64c1", size = 42361858, upload-time = "2025-01-20T11:13:07.623Z" }, - { url = "https://files.pythonhosted.org/packages/d7/7a/ce6174664b9077fc673d172e4c888cb0b128e707e306bc33fff8c2035f0d/llvmlite-0.44.0-cp311-cp311-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:f01a394e9c9b7b1d4e63c327b096d10f6f0ed149ef53d38a09b3749dcf8c9610", size = 41184200, upload-time = "2025-01-20T11:13:20.058Z" }, - { url = "https://files.pythonhosted.org/packages/5f/c6/258801143975a6d09a373f2641237992496e15567b907a4d401839d671b8/llvmlite-0.44.0-cp311-cp311-win_amd64.whl", hash = "sha256:d8489634d43c20cd0ad71330dde1d5bc7b9966937a263ff1ec1cebb90dc50955", size = 30331193, upload-time = "2025-01-20T11:13:26.976Z" }, - { url = "https://files.pythonhosted.org/packages/15/86/e3c3195b92e6e492458f16d233e58a1a812aa2bfbef9bdd0fbafcec85c60/llvmlite-0.44.0-cp312-cp312-macosx_10_14_x86_64.whl", hash = "sha256:1d671a56acf725bf1b531d5ef76b86660a5ab8ef19bb6a46064a705c6ca80aad", size = 28132297, upload-time = "2025-01-20T11:13:32.57Z" }, - { url = "https://files.pythonhosted.org/packages/d6/53/373b6b8be67b9221d12b24125fd0ec56b1078b660eeae266ec388a6ac9a0/llvmlite-0.44.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:5f79a728e0435493611c9f405168682bb75ffd1fbe6fc360733b850c80a026db", size = 26201105, upload-time = "2025-01-20T11:13:38.744Z" }, - { url = "https://files.pythonhosted.org/packages/cb/da/8341fd3056419441286c8e26bf436923021005ece0bff5f41906476ae514/llvmlite-0.44.0-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:c0143a5ef336da14deaa8ec26c5449ad5b6a2b564df82fcef4be040b9cacfea9", size = 42361901, upload-time = "2025-01-20T11:13:46.711Z" }, - { url = "https://files.pythonhosted.org/packages/53/ad/d79349dc07b8a395a99153d7ce8b01d6fcdc9f8231355a5df55ded649b61/llvmlite-0.44.0-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:d752f89e31b66db6f8da06df8b39f9b91e78c5feea1bf9e8c1fba1d1c24c065d", size = 41184247, upload-time = "2025-01-20T11:13:56.159Z" }, - { url = "https://files.pythonhosted.org/packages/e2/3b/a9a17366af80127bd09decbe2a54d8974b6d8b274b39bf47fbaedeec6307/llvmlite-0.44.0-cp312-cp312-win_amd64.whl", hash = "sha256:eae7e2d4ca8f88f89d315b48c6b741dcb925d6a1042da694aa16ab3dd4cbd3a1", size = 30332380, upload-time = "2025-01-20T11:14:02.442Z" }, - { url = "https://files.pythonhosted.org/packages/89/24/4c0ca705a717514c2092b18476e7a12c74d34d875e05e4d742618ebbf449/llvmlite-0.44.0-cp313-cp313-macosx_10_14_x86_64.whl", hash = "sha256:319bddd44e5f71ae2689859b7203080716448a3cd1128fb144fe5c055219d516", size = 28132306, upload-time = "2025-01-20T11:14:09.035Z" }, - { url = "https://files.pythonhosted.org/packages/01/cf/1dd5a60ba6aee7122ab9243fd614abcf22f36b0437cbbe1ccf1e3391461c/llvmlite-0.44.0-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:9c58867118bad04a0bb22a2e0068c693719658105e40009ffe95c7000fcde88e", size = 26201090, upload-time = "2025-01-20T11:14:15.401Z" }, - { url = "https://files.pythonhosted.org/packages/d2/1b/656f5a357de7135a3777bd735cc7c9b8f23b4d37465505bd0eaf4be9befe/llvmlite-0.44.0-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:46224058b13c96af1365290bdfebe9a6264ae62fb79b2b55693deed11657a8bf", size = 42361904, upload-time = "2025-01-20T11:14:22.949Z" }, - { url = "https://files.pythonhosted.org/packages/d8/e1/12c5f20cb9168fb3464a34310411d5ad86e4163c8ff2d14a2b57e5cc6bac/llvmlite-0.44.0-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:aa0097052c32bf721a4efc03bd109d335dfa57d9bffb3d4c24cc680711b8b4fc", size = 41184245, upload-time = "2025-01-20T11:14:31.731Z" }, - { url = "https://files.pythonhosted.org/packages/d0/81/e66fc86539293282fd9cb7c9417438e897f369e79ffb62e1ae5e5154d4dd/llvmlite-0.44.0-cp313-cp313-win_amd64.whl", hash = "sha256:2fb7c4f2fb86cbae6dca3db9ab203eeea0e22d73b99bc2341cdf9de93612e930", size = 30331193, upload-time = "2025-01-20T11:14:38.578Z" }, -] - [[package]] name = "markdown" version = "3.10" @@ -1473,15 +1445,6 @@ dependencies = [ all = [ { name = "imageio" }, { name = "imageio-ffmpeg" }, - { name = "numba" }, - { name = "scipy" }, -] -numba = [ - { name = "numba" }, - { name = "scipy" }, -] -tbb = [ - { name = "tbb" }, ] video = [ { name = "imageio" }, @@ -1494,15 +1457,12 @@ requires-dist = [ { name = "gymnasium", specifier = "==1.1.1" }, { name = "imageio", marker = "extra == 'video'", specifier = ">=2.30" }, { name = "imageio-ffmpeg", marker = "extra == 'video'", specifier = ">=0.5" }, - { name = "motrix-env-core", extras = ["numba", "video"], marker = "extra == 'all'", editable = "motrix_env_core" }, - { name = "numba", marker = "extra == 'numba'", specifier = "==0.61.2" }, + { name = "motrix-env-core", extras = ["video"], marker = "extra == 'all'", editable = "motrix_env_core" }, { name = "numpy", specifier = ">=1.26" }, { name = "omegaconf", specifier = ">=2.3,<2.4" }, - { name = "scipy", marker = "extra == 'numba'", specifier = "==1.15.3" }, - { name = "tbb", marker = "extra == 'tbb'", specifier = "==2021.13.0" }, { name = "typing-extensions", specifier = ">=4.1" }, ] -provides-extras = ["numba", "video", "all", "tbb"] +provides-extras = ["video", "all"] [[package]] name = "motrix-env-motrixsim" @@ -1544,7 +1504,7 @@ version = "0.4.0" source = { editable = "motrix_envs" } dependencies = [ { name = "gymnasium" }, - { name = "motrix-env-core", extra = ["numba", "video"] }, + { name = "motrix-env-core", extra = ["video"] }, { name = "motrix-env-motrixsim" }, { name = "motrix-robots" }, { name = "numpy" }, @@ -1613,9 +1573,6 @@ rslrl = [ skrl-jax = [ { name = "motrix-rl-skrl", extra = ["jax"] }, ] -tbb = [ - { name = "motrix-env-core", extra = ["tbb"] }, -] unitree = [ { name = "motrix-deploy-unitree", extra = ["native"] }, ] @@ -1646,7 +1603,6 @@ requires-dist = [ { name = "motrix-deploy-mujoco", editable = "motrix_deploy_mujoco" }, { name = "motrix-deploy-tasks", editable = "motrix_deploy_tasks" }, { name = "motrix-deploy-unitree", extras = ["native"], marker = "extra == 'unitree'", editable = "motrix_deploy_unitree" }, - { name = "motrix-env-core", extras = ["tbb"], marker = "extra == 'tbb'", editable = "motrix_env_core" }, { name = "motrix-env-motrixsim", editable = "motrix_env_motrixsim" }, { name = "motrix-env-mujoco", editable = "motrix_env_mujoco" }, { name = "motrix-envs", editable = "motrix_envs" }, @@ -1668,7 +1624,7 @@ requires-dist = [ { name = "sphinxcontrib-video", marker = "extra == 'docs'" }, { name = "tensorboard", marker = "extra == 'docs'", specifier = "==2.20.0" }, ] -provides-extras = ["rslrl", "skrl-jax", "unitree", "docs", "cuda", "tbb", "rocm"] +provides-extras = ["rslrl", "skrl-jax", "unitree", "docs", "cuda", "rocm"] [package.metadata.requires-dev] dev = [ @@ -1696,7 +1652,7 @@ dependencies = [ { name = "gymnasium" }, { name = "hydra-core" }, { name = "motrix-deploy" }, - { name = "motrix-env-core", extra = ["numba", "video"] }, + { name = "motrix-env-core", extra = ["video"] }, { name = "numpy" }, { name = "nvidia-ml-py" }, { name = "omegaconf" }, @@ -2072,38 +2028,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/b9/54/dd730b32ea14ea797530a4479b2ed46a6fb250f682a9cfb997e968bf0261/networkx-3.4.2-py3-none-any.whl", hash = "sha256:df5d4365b724cf81b8c6a7312509d0c22386097011ad1abe274afd5e9d3bbc5f", size = 1723263, upload-time = "2024-10-21T12:39:36.247Z" }, ] -[[package]] -name = "numba" -version = "0.61.2" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "llvmlite" }, - { name = "numpy" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/1c/a0/e21f57604304aa03ebb8e098429222722ad99176a4f979d34af1d1ee80da/numba-0.61.2.tar.gz", hash = "sha256:8750ee147940a6637b80ecf7f95062185ad8726c8c28a2295b8ec1160a196f7d", size = 2820615, upload-time = "2025-04-09T02:58:07.659Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/eb/ca/f470be59552ccbf9531d2d383b67ae0b9b524d435fb4a0d229fef135116e/numba-0.61.2-cp310-cp310-macosx_10_14_x86_64.whl", hash = "sha256:cf9f9fc00d6eca0c23fc840817ce9f439b9f03c8f03d6246c0e7f0cb15b7162a", size = 2775663, upload-time = "2025-04-09T02:57:34.143Z" }, - { url = "https://files.pythonhosted.org/packages/f5/13/3bdf52609c80d460a3b4acfb9fdb3817e392875c0d6270cf3fd9546f138b/numba-0.61.2-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:ea0247617edcb5dd61f6106a56255baab031acc4257bddaeddb3a1003b4ca3fd", size = 2778344, upload-time = "2025-04-09T02:57:36.609Z" }, - { url = "https://files.pythonhosted.org/packages/e2/7d/bfb2805bcfbd479f04f835241ecf28519f6e3609912e3a985aed45e21370/numba-0.61.2-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:ae8c7a522c26215d5f62ebec436e3d341f7f590079245a2f1008dfd498cc1642", size = 3824054, upload-time = "2025-04-09T02:57:38.162Z" }, - { url = "https://files.pythonhosted.org/packages/e3/27/797b2004745c92955470c73c82f0e300cf033c791f45bdecb4b33b12bdea/numba-0.61.2-cp310-cp310-manylinux_2_28_aarch64.whl", hash = "sha256:bd1e74609855aa43661edffca37346e4e8462f6903889917e9f41db40907daa2", size = 3518531, upload-time = "2025-04-09T02:57:39.709Z" }, - { url = "https://files.pythonhosted.org/packages/b1/c6/c2fb11e50482cb310afae87a997707f6c7d8a48967b9696271347441f650/numba-0.61.2-cp310-cp310-win_amd64.whl", hash = "sha256:ae45830b129c6137294093b269ef0a22998ccc27bf7cf096ab8dcf7bca8946f9", size = 2831612, upload-time = "2025-04-09T02:57:41.559Z" }, - { url = "https://files.pythonhosted.org/packages/3f/97/c99d1056aed767503c228f7099dc11c402906b42a4757fec2819329abb98/numba-0.61.2-cp311-cp311-macosx_10_14_x86_64.whl", hash = "sha256:efd3db391df53aaa5cfbee189b6c910a5b471488749fd6606c3f33fc984c2ae2", size = 2775825, upload-time = "2025-04-09T02:57:43.442Z" }, - { url = "https://files.pythonhosted.org/packages/95/9e/63c549f37136e892f006260c3e2613d09d5120672378191f2dc387ba65a2/numba-0.61.2-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:49c980e4171948ffebf6b9a2520ea81feed113c1f4890747ba7f59e74be84b1b", size = 2778695, upload-time = "2025-04-09T02:57:44.968Z" }, - { url = "https://files.pythonhosted.org/packages/97/c8/8740616c8436c86c1b9a62e72cb891177d2c34c2d24ddcde4c390371bf4c/numba-0.61.2-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:3945615cd73c2c7eba2a85ccc9c1730c21cd3958bfcf5a44302abae0fb07bb60", size = 3829227, upload-time = "2025-04-09T02:57:46.63Z" }, - { url = "https://files.pythonhosted.org/packages/fc/06/66e99ae06507c31d15ff3ecd1f108f2f59e18b6e08662cd5f8a5853fbd18/numba-0.61.2-cp311-cp311-manylinux_2_28_aarch64.whl", hash = "sha256:bbfdf4eca202cebade0b7d43896978e146f39398909a42941c9303f82f403a18", size = 3523422, upload-time = "2025-04-09T02:57:48.222Z" }, - { url = "https://files.pythonhosted.org/packages/0f/a4/2b309a6a9f6d4d8cfba583401c7c2f9ff887adb5d54d8e2e130274c0973f/numba-0.61.2-cp311-cp311-win_amd64.whl", hash = "sha256:76bcec9f46259cedf888041b9886e257ae101c6268261b19fda8cfbc52bec9d1", size = 2831505, upload-time = "2025-04-09T02:57:50.108Z" }, - { url = "https://files.pythonhosted.org/packages/b4/a0/c6b7b9c615cfa3b98c4c63f4316e3f6b3bbe2387740277006551784218cd/numba-0.61.2-cp312-cp312-macosx_10_14_x86_64.whl", hash = "sha256:34fba9406078bac7ab052efbf0d13939426c753ad72946baaa5bf9ae0ebb8dd2", size = 2776626, upload-time = "2025-04-09T02:57:51.857Z" }, - { url = "https://files.pythonhosted.org/packages/92/4a/fe4e3c2ecad72d88f5f8cd04e7f7cff49e718398a2fac02d2947480a00ca/numba-0.61.2-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:4ddce10009bc097b080fc96876d14c051cc0c7679e99de3e0af59014dab7dfe8", size = 2779287, upload-time = "2025-04-09T02:57:53.658Z" }, - { url = "https://files.pythonhosted.org/packages/9a/2d/e518df036feab381c23a624dac47f8445ac55686ec7f11083655eb707da3/numba-0.61.2-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:5b1bb509d01f23d70325d3a5a0e237cbc9544dd50e50588bc581ba860c213546", size = 3885928, upload-time = "2025-04-09T02:57:55.206Z" }, - { url = "https://files.pythonhosted.org/packages/10/0f/23cced68ead67b75d77cfcca3df4991d1855c897ee0ff3fe25a56ed82108/numba-0.61.2-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:48a53a3de8f8793526cbe330f2a39fe9a6638efcbf11bd63f3d2f9757ae345cd", size = 3577115, upload-time = "2025-04-09T02:57:56.818Z" }, - { url = "https://files.pythonhosted.org/packages/68/1d/ddb3e704c5a8fb90142bf9dc195c27db02a08a99f037395503bfbc1d14b3/numba-0.61.2-cp312-cp312-win_amd64.whl", hash = "sha256:97cf4f12c728cf77c9c1d7c23707e4d8fb4632b46275f8f3397de33e5877af18", size = 2831929, upload-time = "2025-04-09T02:57:58.45Z" }, - { url = "https://files.pythonhosted.org/packages/0b/f3/0fe4c1b1f2569e8a18ad90c159298d862f96c3964392a20d74fc628aee44/numba-0.61.2-cp313-cp313-macosx_10_14_x86_64.whl", hash = "sha256:3a10a8fc9afac40b1eac55717cece1b8b1ac0b946f5065c89e00bde646b5b154", size = 2771785, upload-time = "2025-04-09T02:57:59.96Z" }, - { url = "https://files.pythonhosted.org/packages/e9/71/91b277d712e46bd5059f8a5866862ed1116091a7cb03bd2704ba8ebe015f/numba-0.61.2-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:7d3bcada3c9afba3bed413fba45845f2fb9cd0d2b27dd58a1be90257e293d140", size = 2773289, upload-time = "2025-04-09T02:58:01.435Z" }, - { url = "https://files.pythonhosted.org/packages/0d/e0/5ea04e7ad2c39288c0f0f9e8d47638ad70f28e275d092733b5817cf243c9/numba-0.61.2-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:bdbca73ad81fa196bd53dc12e3aaf1564ae036e0c125f237c7644fe64a4928ab", size = 3893918, upload-time = "2025-04-09T02:58:02.933Z" }, - { url = "https://files.pythonhosted.org/packages/17/58/064f4dcb7d7e9412f16ecf80ed753f92297e39f399c905389688cf950b81/numba-0.61.2-cp313-cp313-manylinux_2_28_aarch64.whl", hash = "sha256:5f154aaea625fb32cfbe3b80c5456d514d416fcdf79733dd69c0df3a11348e9e", size = 3584056, upload-time = "2025-04-09T02:58:04.538Z" }, - { url = "https://files.pythonhosted.org/packages/af/a4/6d3a0f2d3989e62a18749e1e9913d5fa4910bbb3e3311a035baea6caf26d/numba-0.61.2-cp313-cp313-win_amd64.whl", hash = "sha256:59321215e2e0ac5fa928a8020ab00b8e57cda8a97384963ac0dfa4d4e6aa54e7", size = 2831846, upload-time = "2025-04-09T02:58:06.125Z" }, -] - [[package]] name = "numpy" version = "2.2.6" @@ -3268,7 +3192,7 @@ name = "scipy" version = "1.15.3" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "numpy" }, + { name = "numpy", marker = "sys_platform != 'darwin' or (extra == 'extra-10-motrix-lab-cuda' and extra == 'extra-10-motrix-lab-rocm') or (extra == 'extra-9-motrix-rl-cuda' and extra == 'extra-9-motrix-rl-rocm')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/0f/37/6964b830433e654ec7485e45a00fc9a27cf868d622838f6b6d9c5ec0d532/scipy-1.15.3.tar.gz", hash = "sha256:eae3cf522bc7df64b42cad3925c876e1b0b6c35c1337c93e12c0f366f55b0eaf", size = 59419214, upload-time = "2025-05-08T16:13:05.955Z" } wheels = [ @@ -3597,17 +3521,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/a2/09/77d55d46fd61b4a135c444fc97158ef34a095e5681d0a6c10b75bf356191/sympy-1.14.0-py3-none-any.whl", hash = "sha256:e091cc3e99d2141a0ba2847328f5479b05d94a6635cb96148ccb3f34671bd8f5", size = 6299353, upload-time = "2025-04-27T18:04:59.103Z" }, ] -[[package]] -name = "tbb" -version = "2021.13.0" -source = { registry = "https://pypi.org/simple" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/c9/32/9c88009a9feb81f0525bf1c4729a7b39a4523d626a40682c792c6a53e323/tbb-2021.13.0-py2.py3-none-manylinux1_i686.whl", hash = "sha256:a2567725329639519d46d92a2634cf61e76601dac2f777a05686fea546c4fe4f", size = 5205096, upload-time = "2024-06-20T15:34:29.818Z" }, - { url = "https://files.pythonhosted.org/packages/5d/b8/994132d9f9493d7b220637f614f3f87727e405b851d96aaffc7d43c2b298/tbb-2021.13.0-py2.py3-none-manylinux1_x86_64.whl", hash = "sha256:aaf667e92849adb012b8874d6393282afc318aca4407fc62f912ee30a22da46a", size = 5386610, upload-time = "2024-06-20T15:35:11.596Z" }, - { url = "https://files.pythonhosted.org/packages/1f/ce/1dd6f2988d06d850374be2c56350d17399d4faaed53b489048f12a3d3474/tbb-2021.13.0-py3-none-win32.whl", hash = "sha256:6669d26703e9943f6164c6407bd4a237a45007e79b8d3832fe6999576eaaa9ef", size = 248996, upload-time = "2024-06-20T15:38:31.467Z" }, - { url = "https://files.pythonhosted.org/packages/f1/24/500811330b3b070e5995c3275181dbcd00c06cef26c6ebfe6ee1ca9b6223/tbb-2021.13.0-py3-none-win_amd64.whl", hash = "sha256:3528a53e4bbe64b07a6112b4c5a00ff3c61924ee46c9c68e004a1ac7ad1f09c3", size = 286910, upload-time = "2024-06-20T15:38:18.376Z" }, -] - [[package]] name = "tensorboard" version = "2.20.0" diff --git a/wiki/design/index.md b/wiki/design/index.md index 511bc16c..5b75c57b 100644 --- a/wiki/design/index.md +++ b/wiki/design/index.md @@ -29,8 +29,6 @@ 定义标准 dataclass 兼容的配置装饰器语义:直接声明可变容器与嵌套配置默认值、逐实例深拷贝隔离、显式 `field()` 优先,以及类型检查和 Hydra/OmegaConf 兼容契约。 - [Scene / Robot 配置设计](./scene-model-config.md) 使用字段式 `SceneAssetsCfg` 与 `SceneObjsCfg` 声明具名 asset 和有序 scene object,通过 `SceneVisualCfg` 配置全局视觉环境,并由可继承的 `StandardSceneCfg` 提供标准 skybox、地面、材质、贴图、haze 和方向光。 -- [G1 WBT 环境设计](./g1-wbt-env.md) - MotrixLab 中 G1 whole-body tracking 环境(direct `DirectEnv`,注册 `g1-29dof-wbt-largebox` / `g1-wbt-dance`)的设计:`WbtMotion` 数据加载与 name-based 状态映射、reset/step 时间推进、effort-scaled PD 控制、actor/critic observation、tracking reward、bad-tracking termination、adaptive timestep sampler、`motrix.fastsac` 执行拓扑切换与 play 变体。 - [MotrixLab Motion NPZ Schema 设计](./motrixlab-motion-npz-schema.md) MotrixLab 自己的 WBT `.npz` 动作文件格式 v1:在 BeyondMimic 原生 schema 上加 `schema_version / joint_names / body_names / num_frames` 与 `ext_*` 扩展槽位,四元数约定切到 xyzw 与 MotrixSim 内部 API 对齐。提供通用 `MotrixMotion` loader、Holosoma/BeyondMimic → MotrixLab 一次性 converter、`scripts/motion/replay.py` 通用化。 - [多 Motion Clip 训练基础设施设计](./motion-multi-clip-infra.md) diff --git a/wiki/design/manager/index.md b/wiki/design/manager/index.md index 05c298af..6b7f1847 100644 --- a/wiki/design/manager/index.md +++ b/wiki/design/manager/index.md @@ -9,11 +9,11 @@ ### 核心设计 - [Manager Runtime 设计](./runtime.md) - ManagerEnv 的运行时边界、配置到 term 的构造链路、step/reset 生命周期、term 协议和 Numba 性能约束。 + ManagerEnv 的运行时边界、配置到 term 的构造链路、step/reset 生命周期、term 协议和执行/性能约束。 - [Manager Task API](./task-authoring.md) 面向任务开发者的 ManagerBasedEnvCfg、manager groups、term 创建、query 声明、action、observation、reward、termination 和 reset 指南。 - [Manager Kernel ABI](./kernel-abi.md) - 定义 ManagerContext、KernelData tree、leaf scope、term ownership、metric backing 和 fused kernel fingerprint。 + 定义 ManagerContext、KernelData tree、leaf scope、term ownership、metric backing 和 resolved layout。 ### Simulator 集成 @@ -27,9 +27,11 @@ ### 专题设计 - [Term 自声明 Sim Query 设计](./term-required-queries.md) - 定义 observation term 的 `required_queries()`、任务声明与 term 贡献的合并和冲突规则。 + 定义 term 通过 `SimDataQuery` 参数自声明 sim 数据需求、与任务声明合并的 key 单一属主规则,以及 resolver 到批量 sim 视图的映射。 - [IsaacLab ManagerCfg Adapter 设计](./isaaclab-manager-cfg-adapter.md) 定义 IsaacLab-like ManagerCfg 到 MotrixLab native term/query 配置的适配边界;本文是目标设计,尚未表示全部 API 已实现。 +- [Torch Manager 执行器设计](./torch-executor.md) + 定义 Torch 执行器(唯一数值执行器)的调度语义、批量 term ctx 约定与 resolver 分工。 ## 使用说明 diff --git a/wiki/design/manager/isaaclab-manager-cfg-adapter.md b/wiki/design/manager/isaaclab-manager-cfg-adapter.md index e8262062..d2325c15 100644 --- a/wiki/design/manager/isaaclab-manager-cfg-adapter.md +++ b/wiki/design/manager/isaaclab-manager-cfg-adapter.md @@ -2,7 +2,7 @@ ## 摘要 -本设计定义一个面向 IsaacLab manager-based 任务配置的薄适配层,目标是尽量保留用户熟悉的 ManagerCfg 分组、`func + params`、term 命名、嵌套覆写和 selector 表达。Numba-compatible 自定义 MDP 函数在程序构建时被特化为私有 concrete `@kernel_data` Term,selector 被归一化为显式 `SimDataQuery`。适配层不直接执行任意 IsaacLab Torch callback,也不改变 [Manager Task API](./task-authoring.md) 和 [Manager Runtime](./runtime.md) 中的 Numba 运行时边界。 +本设计定义一个面向 IsaacLab manager-based 任务配置的薄适配层,目标是尽量保留用户熟悉的 ManagerCfg 分组、`func + params`、term 命名、嵌套覆写和 selector 表达。Torch 批量自定义 MDP 函数在程序构建时被特化为私有 concrete `@kernel_data` Term,selector 被归一化为显式 `SimDataQuery`。适配层不直接执行任意 IsaacLab Torch callback,也不改变 [Manager Task API](./task-authoring.md) 和 [Manager Runtime](./runtime.md) 中的 Torch 执行器边界。 本文是目标设计,不表示下述 API 均已实现。兼容基线参考本地 IsaacLab 快照 `21bcb476b2`;该快照检查时落后 `origin/main` 202 个 commit,因此不声称覆盖所有上游版本。 @@ -16,7 +16,7 @@ IsaacLab-like ManagerCfg -> native TermCfg.__call__() / functional term generation -> concrete @kernel_data Term -> ManagerContext / SimDataQuery / ReadPlan - -> fused Numba kernel + -> Torch executor(torch.compile kernel,整批张量) ``` 这意味着: @@ -31,15 +31,15 @@ IsaacLab-like ManagerCfg - 不直接执行 IsaacLab 中以 `env` 为参数的通用 Python/Torch `func`。 - 不允许 term 隐式读取整个 environment 或在 kernel 外隐藏 simulator 依赖。 - 不静默忽略 noise、clip、history 等未支持选项。 -- 不让 callable、`dict[str, Any]` 或完整 environment 进入 KernelData/Numba ABI。 -- 不为表面兼容牺牲 typed cfg、compiler-visible dependency 或 fused Numba kernel。 +- 不让 callable、`dict[str, Any]` 或完整 environment 进入 kernel_data ABI 或编译图。 +- 不为表面兼容牺牲 typed cfg、compiler-visible dependency 或批量 Torch kernel。 ## 2. 当前主要差异 | 配置维度 | IsaacLab 常见形态 | MotrixLab 当前形态 | Adapter 责任 | | --- | --- | --- | --- | | Manager group | `@configclass` 容器,主要依赖字段结构 | `dict` 或指定的 `Manager*Cfg` nominal type | 按结构枚举配置字段,不要求必须继承特定 group base | -| Term 定义 | `func` + `params` + 通用 term cfg | concrete typed `*TermCfg.__call__()` | Numba-compatible 函数自动特化为 concrete Term;foreign callback 通过注册 adapter 映射 | +| Term 定义 | `func` + `params` + 通用 term cfg | concrete typed `*TermCfg.__call__()` | Torch 批量函数自动特化为 concrete Term;foreign callback 通过注册 adapter 映射 | | 禁用 term | 继承配置中常设为 `None` | term group 目前期望有效 cfg 实例 | normalization 时过滤 `None` | | Reward 权重 | 常通过 `.weight` 嵌套覆写,0 通常表示关闭 | typed reward cfg,运行时统一加权 | 保留 `.weight` 覆写;`weight == 0` 在 program 构建前裁剪 | | Scene selector | `SceneEntityCfg` 将 asset/body/joint/sensor 选择放在 term params 内 | `queries.data` 中集中声明 `SimDataQuery` | 解析 backend-neutral selector,产生并去重 `SimDataQuery` | @@ -78,7 +78,7 @@ track_lin_vel_xy_exp = RewTerm( ) ``` -后一种形态应成为 adapter 的一等输入,而不是只能通过逐函数 factory 注册使用的例外路径。用户经常需要自定义 MDP 函数,因此普通 Numba-compatible 函数不应额外要求 `CommandRef`、参数语义 decorator 或 adapter registry。 +后一种形态应成为 adapter 的一等输入,而不是只能通过逐函数 factory 注册使用的例外路径。用户经常需要自定义 MDP 函数,因此普通 Torch 批量函数不应额外要求 `CommandRef`、参数语义 decorator 或 adapter registry。 #### IsaacLab 的实际解析语义 @@ -135,7 +135,7 @@ class _GeneratedTrackLinVelXYExp: #### 自动生成与 foreign adapter 边界 -Numba-compatible 纯函数使用上述自动生成路径,无需逐函数注册。原始 IsaacLab callback 若依赖 `ManagerBasedRLEnv`、Torch tensor 或其 scene/manager object model,不能原样进入 Numba kernel;这类 foreign callable 需要注册到等价的 Numba-native function 或 typed `*TermCfg` factory。 +Torch 批量纯函数使用上述自动生成路径,无需逐函数注册。原始 IsaacLab callback 若依赖 `ManagerBasedRLEnv` 的 scene/manager object model,不能原样进入 Torch 执行器的编译图;这类 foreign callable 需要注册到等价的批量 Torch function 或 typed `*TermCfg` factory。 自动生成首先覆盖无状态 Reward 和 Termination。Observation 还需要固定输出 `size`,可由 config 或函数 metadata 提供;Action、Command 和其他拥有 reset/update 生命周期的 stateful term 继续使用显式 typed Term class。 @@ -167,7 +167,7 @@ self.rewards.action_rate_l2.weight = -0.005 self.rewards.undesired_contacts = None ``` -第一行修改 typed cfg 中的纯数值字段;第二行在 normalization 时删除 term。对 reward,`weight == 0` 也应在 `NumbaTaskProgram` 构建之前裁剪,避免为明确关闭的 term 创建 runtime object、query 和 kernel slot。 +第一行修改 typed cfg 中的纯数值字段;第二行在 normalization 时删除 term。对 reward,`weight == 0` 也应在 manager program 构建之前裁剪,避免为明确关闭的 term 创建 runtime object、query 和 kernel slot。 ## 4. Observation group 兼容范围 @@ -188,8 +188,8 @@ self.rewards.undesired_contacts = None - term 字段支持 `None`,reward `weight == 0` 在程序构建前裁剪。 - 提供 backend-neutral `SceneEntityCfg`,归一化为显式、可去重的 `SimDataQuery`。 - 保留 `.weight = ...` 与 `term = None` 等嵌套覆写方式。 -- 对无状态 Reward/Termination 接受 Numba-compatible `func + params`,自动生成 concrete Term,不要求用户注册 factory。 -- `params` 字典只存在于 host config,不进入 KernelData/Numba ABI。 +- 对无状态 Reward/Termination 接受 Torch 批量 `func + params`,自动生成 concrete Term,不要求用户注册 factory。 +- `params` 字典只存在于 host config,不进入 kernel_data ABI 或编译图。 ### P1:Observation 常用选项 @@ -201,7 +201,7 @@ self.rewards.undesired_contacts = None ### P2:Foreign callback 与扩展 term -- 为依赖 IsaacLab env/Torch object model 的 foreign callable 提供显式 adapter registry,映射到 Numba-native function 或 native `*TermCfg` factory。 +- 为依赖 IsaacLab env/Torch object model 的 foreign callable 提供显式 adapter registry,映射到批量 Torch function 或 native `*TermCfg` factory。 - 对 functional Observation 补充固定 `size` 契约,再按同一生成模型支持。 - Action、Command 与 stateful term 只在能保留其生命周期契约时扩展,不将其伪装为无状态 callback。 - 未注册 foreign callable、不兼容 params 和未支持 option 在 config normalization 阶段立即报错。 @@ -212,9 +212,9 @@ self.rewards.undesired_contacts = None 适配层完成时至少应验证: 1. 一个典型 IsaacLab locomotion ManagerCfg 可保留 `func + params`、原有分组、term 名称与大部分派生配置覆写。 -2. 用户自定义 Numba-compatible Reward/Termination 函数无需创建 Term class、注册 factory 或声明 `CommandRef`。 +2. 用户自定义 Torch 批量 Reward/Termination 函数无需创建 Term class、注册 factory 或声明 `CommandRef`。 3. `term = None`、`weight == 0`、`critic -> value` 和 selector 去重有独立配置层测试。 4. 生成的 KernelData/ABI 中不包含 callable、`params` 字典或 environment reference;compiler plan 记录原函数 fingerprint 与生成 term schema。 5. 数值参数、string literal、enum、ndarray 和 selector 有独立 lowering 测试,不支持的值会报出包含 group/term/param 路径的错误。 -6. `ctx.commands[command_name]` 能在生成 wrapper 传入 string literal 时编译为 nopython;无效 key 在 warm-up 阶段报错。 +6. `ctx.commands[command_name]` 能在生成 wrapper 传入 string literal 时进入 torch.compile 编译图;无效 key 在 warm-up 阶段报错。 7. 归一化前后 term 顺序、observation layout、reward 权重和 termination 语义有 parity 测试。 diff --git a/wiki/design/manager/kernel-abi.md b/wiki/design/manager/kernel-abi.md index e5c7404a..cba36dca 100644 --- a/wiki/design/manager/kernel-abi.md +++ b/wiki/design/manager/kernel-abi.md @@ -2,33 +2,33 @@ ## 摘要 -本文定义 Manager fused Numba kernel 的数据 ABI:`@kernel_data` logical tree 如何 flatten,`ManagerContext` 如何按 lane +本文定义 Manager term 的数据 ABI:`@kernel_data` logical tree 如何 flatten,`ManagerContext` 如何以整批张量 提供固定 namespace,以及 `PER_ENV`、`SHARED` 和 metric backing 的 ownership 语义。任务开发者只需遵守 -[Manager Task API](./task-authoring.md) 中的 term 协议;backend 读侧和 simulator 边界见 [SimBackend](./sim-backend.md)。 +[Manager Task API](./task-authoring.md) 中的 term 协议;数值执行器见 [Torch 执行器设计](./torch-executor.md), +backend 读侧和 simulator 边界见 [SimBackend](./sim-backend.md)。 ## 1. ManagerContext -`ManagerContext` 是 compiled term 的唯一 manager/runtime 输入入口: +`ManagerContext` 是 term dispatch 的唯一 manager/runtime 输入入口: ```python @kernel_data class ManagerContext: env_id: np.ndarray actions: Map - observations: Map - rewards: Map - terminations: Map commands: Map rand: RandValue sim: Map dt: np.float32 + sim_reset_requested: np.ndarray ``` Context 不包含 config、provider、backend handle、dynamic Python container 或 generic `states` / `resources` store。 Persistent runtime data 直接归属 concrete term;framework-owned random state 通过 `rand` 暴露,`dt` 提供当前 -`ctrl_dt` 的固定 float32 标量。 +`ctrl_dt` 的固定 float32 标量。执行器把该结构解析为固定 layout,运行时以 namespace 视图(`ctx.sim[key]`、 +`ctx.commands[name].field` 等)暴露整批张量。 -## 2. Lane-local 语义 +## 2. 批量语义 Host 侧 `env_id` 是: @@ -36,28 +36,30 @@ Host 侧 `env_id` 是: np.arange(num_envs, dtype=np.int64) ``` -Generated kernel 在 `prange` loop 中按当前 environment id 重建 context。普通 step 和 partial reset 都使用原始 environment id, -因此 `ctx.env_id`、action/command state 和 simulator snapshot 始终属于同一 environment,不会把 reset batch row 当作 environment id。 +所有 `PER_ENV` leaves 是 `(N, ...)` 整批数组;term 一次调用处理整批。partial reset 时执行器按原始 +environment id gather `PER_ENV` leaves(第 0 维 = len(env_ids)),term 写 gathered 行后 scatter 回常驻缓冲, +不会把 reset batch row 当作 environment id。 -所有 Map 的 key 顺序、concrete value schema、field tree、leaf dtype、scope 和 input slot 在编译前固定。Map lookup 必须使用 literal key: +所有 Map 的 key 顺序、concrete value schema、field tree、leaf dtype、scope 和 input slot 在 resolve 前固定。 +Map lookup 使用 literal key: ```python motion = ctx.commands["motion"] dof_pos = ctx.sim["robot_dof_pos"] ``` -动态字符串 key、运行期插入/删除和 fallback lookup 不属于 ABI;错误 key 应在 typing/lowering 或 warm-up 阶段失败。 +动态字符串 key、运行期插入/删除和 fallback lookup 不属于 ABI;错误 key 在 resolve 或首次执行阶段失败。 ## 3. Leaf scope -| 类型 | 默认 scope | lane 访问语义 | +| 类型 | 默认 scope | term 访问语义 | | --- | --- | --- | -| `np.ndarray` | `PER_ENV` | 取当前 environment 的 row/view | -| `SharedArray` | `SHARED` | 所有 lane 访问完整可写 ndarray | -| Python scalar | `SHARED` | 所有 lane 使用同一个标量 | -| 固定 dtype NumPy scalar | `SHARED` | 所有 lane 使用同一个定 dtype 标量 | +| `np.ndarray` | `PER_ENV` | 整批 `(N, ...)` 数组/张量;reset 场景为 gathered 行 | +| `SharedArray` | `SHARED` | 完整可写 ndarray,所有环境共享同一份 | +| Python scalar | `SHARED` | 所有环境使用同一个标量 | +| 固定 dtype NumPy scalar | `SHARED` | 所有环境使用同一个定 dtype 标量 | -`SharedArray` 只描述 lowering scope,不代表 immutable。数组是否可以更新由 concrete term 的 owner/lifecycle 契约决定。 +`SharedArray` 只描述 scope,不代表 immutable。数组是否可以更新由 concrete term 的 owner/lifecycle 契约决定。 Term receiver 的参数通常按 shared lowering 处理,以免把静态参数误切片。 ## 4. Logical tree 与 flat ABI @@ -71,40 +73,41 @@ rebuilt = unflatten_kernel_data(tree_def, leaves) Record 可以嵌套另一个 KernelData record;不支持 multiple inheritance、覆盖 inherited field、普通 dataclass、递归结构、dict、list、字符串、object array 或 backend handle。 -Logical tree 与 physical lowering 分离:同一 concrete value 在 host term 和 context 中保留相同的 tree schema,但 lowering 根据使用位置决定 scope、slot 和 lane proxy。 -Lowered ABI 只允许 ndarray/scalar leaves 跨越 `prange` boundary;nested proxy 在当前 lane 重建,不应产生稳态临时 ndarray 分配。 +Logical tree 与 execution lowering 分离:同一 concrete value 在 host term 和 context 中保留相同的 tree schema, +但 resolver 根据使用位置决定 scope、slot 与运行时 namespace 视图。flat ABI 只允许 ndarray/scalar leaves 跨越 +执行边界;term 内不应产生稳态临时大数组分配。 ## 5. Ownership 与 metric - config 只创建 term,不持有 environment-local mutable arrays; -- 每个 concrete term 只 canonicalize 一次,host 和 compiled path 共享 backing identity; +- 每个 concrete term 只 canonicalize 一次,host 和执行路径共享 backing identity; - action/command term 拥有自己的 persistent buffers; -- observation output 由调用方提供,generated kernel 在当前 invocation 写入; -- simulator snapshot 由 `PhysicsReadProgram` 拥有并写入,compiled consumer 只读; -- `metric()` 字段由 owning term 写入自己的当前 lane,host collector 负责聚合 scalar metric。 +- observation output 由调用方提供,执行器在当前 invocation 写入; +- simulator snapshot 由 `PhysicsReadProgram` 拥有并写入,term 只读; +- `metric()` 字段由 owning term 写入自己的当前行,host collector 负责聚合 scalar metric。 Term 可以读取其他 term 的数据,但不应调用其他 term 的行为方法。Context store 表达数据依赖,不承担 manager scheduling。 -## 6. 编译与 fingerprint +## 6. Resolver 职责 -Compiler 负责: +`ManagerProgramResolver`(backend-neutral)负责: - flatten concrete runtime terms 和 `ManagerContext`; - 收集 read program 的 sim input map; -- lowering context、term receiver 和 literal Map lookup; -- 校验方法签名、返回类型、dtype 和 layout; -- 生成 observation、reward、termination、command evaluation/step 和 partial-reset(含 command reset)kernels; -- 生成 layout/fingerprint,覆盖 concrete type、Map key、field tree、scope、sim binding 和 invocation schema。 +- 解析 context、term receiver 和 literal Map lookup 的 slot/表达式; +- 校验 `@dispatch` 标记、方法签名、返回注解和 layout; +- 产出 `ManagerLayout`、`ReadPlan`、`KernelInputSource` 与 `PreparedInvocation`,供 Torch 执行器调度。 -构造期 warm-up 必须确认 generated dispatcher 的 nopython signature。schema rename、term type 变化、scope 变化和 query layout 变化都必须产生新的 fingerprint,避免错误复用 compiled kernel 或旧 Numba cache。 +构造期 warm-up 以克隆输入执行一次 evaluate/observe,确认输入/输出契约(shape、dtype、batch 维、可写性); +schema rename、term type 变化、scope 变化和 query layout 变化都会改变 resolved layout。 ## 7. 核心不变量 - Context schema 在一个 environment 生命周期内固定; - 每个 key 对应唯一 concrete type 和 stable slot; -- 所有 `PER_ENV` leaves 使用同一个 current lane id; -- `ctx.env_id` 在 step/reset 中与其他 lane-local data 对齐; -- simulator snapshots 只来自 compiled read program; +- 所有 `PER_ENV` leaves 共享同一个第 0 维(整批为 N,reset 为 len(env_ids)); +- `ctx.env_id` 在 step/reset 中与其他 per-env data 对齐; +- simulator snapshots 只来自编译后的 read program; - 不建立 generic `states`、`resources` 或 `ManagerValue` registry; -- generated kernel 不接收完整 environment、config、info 或 simulator object; -- nopython、buffer identity 和 deterministic parity 属于 CI 验收,绝对性能和 scaling 属于独立 benchmark。 +- term dispatch 不接收完整 environment、config、info 或 simulator object; +- 数值行为、buffer identity 和 deterministic parity 属于 CI 验收,绝对性能和 scaling 属于独立 benchmark。 diff --git a/wiki/design/manager/runtime.md b/wiki/design/manager/runtime.md index 02171510..15378050 100644 --- a/wiki/design/manager/runtime.md +++ b/wiki/design/manager/runtime.md @@ -3,10 +3,10 @@ ## 摘要 `ManagerEnv` 是基于 `ArrayEnv` 的 NumPy 前端,使用配置声明 action、command、observation、reward、termination 和 reset, -并将可编译的逐环境数值计算融合为 Numba kernel。它保持现有环境公共契约:action 和 observation 使用 NumPy,状态使用 -`ArrayEnvState`,仿真通过 backend-neutral 的 `SimBackend` 访问。 +数值执行由通用 Torch 执行器按整批张量调度(见 [Torch 执行器设计](./torch-executor.md))。它保持现有环境公共契约:action 和 +observation 使用 NumPy,状态使用 `ArrayEnvState`,仿真通过 backend-neutral 的 `SimBackend` 访问。 -本文集中说明 Manager 的运行时边界、构造与生命周期、term ownership,以及编译和性能约束。任务配置接口见 +本文集中说明 Manager 的运行时边界、构造与生命周期、term ownership,以及执行和性能约束。任务配置接口见 [Manager Task API](./task-authoring.md),kernel 数据 ABI 见 [Manager Kernel ABI](./kernel-abi.md),仿真后端边界见 [SimBackend](./sim-backend.md)。 @@ -16,7 +16,7 @@ ArrayEnv lifecycle ├── action terms(host) ├── SimBackend.step / read program(backend) - ├── fused task kernel(Numba) + ├── task program(Torch executor,整批张量调度) ├── observation / reward / termination outputs └── partial reset(host + backend + reset kernel) ``` @@ -27,7 +27,7 @@ ArrayEnv lifecycle - `DirectEnv` 和 `ManagerEnv` 都通过 registry 的 `sim` 参数选择 `SimBackend`;未指定时使用注册的默认 backend。 - `motrix_env_core` 的 Manager 代码不得 import 具体 simulator 类型; - MotrixSim 对象、backend handle、Python 容器和 `info` 字典停留在 host 边界; -- Numba kernel 只接收固定的 `ManagerContext` 和 flat ndarray/scalar leaves。 +- term dispatch 只接收固定的 `ManagerContext` 视图和 flat ndarray/scalar leaves(整批张量)。 Manager 不是完整的 `step()` 替代品,也不编译物理推进。物理时间只由 control/training loop 推进一次,reset 不推进物理。 @@ -59,10 +59,11 @@ Manager group 的顺序是行为契约的一部分:action slice、observation 2. 解析并编译 model、read、write 和 reset declarations; 3. 创建 action、command、observation、reward 和 termination terms; 4. 为每个 term 执行一次 `canonicalize_kernel_data`; -5. 构造 `ManagerContext` layout 和 fused kernels; -6. 编译 fused kernels:term dispatch 以签名校验(`inspect` + 类型标注)后内联进 kernel, - nopython signature、dtype 与返回契约由 fused kernel 编译一次性确认;term 不再单独编译 - dispatcher(避免逐 term 编译开销,issue #56)。 +5. resolver(`ManagerProgramResolver`)构造 `ManagerContext` layout、read plan 和 + prepared invocations:term dispatch 以签名校验(`inspect` + 类型标注),不做任何 + 代码生成或 dispatcher 编译; +6. Torch 执行器按元数据调度;首次真实调用完成惰性 specialization,`warmup()` 用克隆 + 输入执行一次 evaluate/observe 预热。 ### Step @@ -71,7 +72,7 @@ Manager group 的顺序是行为契约的一部分:action slice、observation 3. backend 推进 `sim_substeps`; 4. command term 在 host 侧更新 persistent state; 5. read program 刷新 simulator snapshot; -6. fused kernel 执行 command evaluation、observation、reward 和 termination; +6. Torch 执行器执行 command evaluation、observation、reward 和 termination(一次整批调用/term); 7. host 折叠 command 统计(`on_transition()`),`ArrayEnv` 处理 truncation 和 done reset;`state.metrics` 是常驻活视图(kernel 原地写),跨步持有走 `process_metrics()` 归约快照。 ### Partial reset @@ -81,7 +82,7 @@ reset 只处理指定的原始 `env_ids`: - action/command/reset terms 只修改目标 rows; - backend reset program 只写目标 rows; - read program 只刷新目标 rows; -- reset kernel 使用原始 environment id 重建 context,不把 reset batch row 当作环境 id; +- reset 路径按原始 environment id gather `PER_ENV` leaves,不把 reset batch row 当作环境 id; - `ArrayEnv` 继续负责 auto-reset、episode steps 和 truncation 语义。 ## 4. Term 协议 @@ -120,33 +121,35 @@ class ObservationTerm: class RewardTerm: - dispatch: Callable[..., float] + dispatch: Callable[..., torch.Tensor] args: tuple[Any, ...] - # dispatch(ctx, *args) -> float + # dispatch(ctx, *args) -> (N,) float32 tensor class TerminationTerm: - dispatch: Callable[..., bool] + dispatch: Callable[..., torch.Tensor] args: tuple[Any, ...] metrics: dict[str, np.ndarray] - # dispatch(ctx, metrics, *args) -> bool when metrics are present - # dispatch(ctx, *args) -> bool otherwise + # dispatch(ctx, *args) -> (N,) bool tensor; metrics written via ctx.metrics class CommandTerm: def reset(ctx: ResetContext) -> None: ... # host: prepare shared reset data for episode resets def on_transition() -> None: ... # host: fold per-step statistics - def reset_env(ctx) -> None: ... # @dispatch: reset one lane in the reset kernel - def advance(ctx) -> None: ... # @dispatch: advance one lane in the transition kernel - def update(ctx) -> None: ... # @dispatch: update derived per-lane command data + @staticmethod + def reset_env(receiver, ctx) -> None: ... # @dispatch: batched, lane-gathered reset rows + @staticmethod + def advance(receiver, ctx) -> None: ... # @dispatch: batched, whole batch + @staticmethod + def update(receiver, ctx) -> None: ... # @dispatch: batched, whole batch ``` -Command term 的生命周期拆成 fused-kernel hooks 与 host 侧方法:`advance(ctx)` 在 transition kernel 末尾推进 lane 状态, -`reset_env(ctx)` 在 reset kernel 开头、sim reset terms 之前重采样 lane 状态。lane 在 `advance` 中置 `ctx.sim_reset_requested`(Manager 常驻 +Command term 的生命周期拆成 batched hooks 与 host 侧方法:`advance(receiver, ctx)` 在 evaluate 末尾推进整批状态, +`reset_env(receiver, ctx)` 在 reset 路径开头、sim reset terms 之前重采样(lane-gathered)状态。lane 在 `advance` 中置 `ctx.sim_reset_requested`(Manager 常驻 buffer,每次 physics step 前清零)即可请求 sim-only reset:该 lane 与当步的 episode reset -合并进**同一次** reset kernel 运行(`reset_env` 重采样 → sim reset terms 重写仿真状态 → 按行重读 kernel inputs),但不会 +合并进**同一次** reset 运行(`reset_env` 重采样 → sim reset terms 重写仿真状态 → 按行重读 kernel inputs),但不会 触发 episode bookkeeping(`episode_steps`、truncation)与 command/action 的 host reset。WBT 在 motion clip wrap 时使用 该机制。host `reset(ResetContext)` 只在 episode reset 时准备跨 lane 的共享数据(如采样分布);`on_transition()` 每步折叠 统计。 @@ -163,17 +166,16 @@ route-local controls,不重复推进 history,也不直接写 simulator state Manager 根据 actuator route 合并 controls 后写入 backend。Term 不应调用其他 term 的行为方法,跨 term 依赖通过 `ManagerContext` 的数据 store 表达。 -Observation terms use the host-side ``ObservationTerm(size, dispatch, *args)`` form. ``size`` is the fixed output width, and additional arguments are passed positionally; simple scalar/array arguments do not require a separate Args class. Reward terms use the analogous ``RewardTerm(dispatch, *args)`` form and return one numeric scalar per environment. Termination terms use ``TerminationTerm(dispatch, *args, metrics=...)``; metrics are optional per-environment outputs exposed to the dispatch as a static ``Map[np.ndarray]``. The compiler validates and lowers these values into the static kernel ABI, then emits ``dispatch(ctx, out, *args)`` for observations, ``dispatch(ctx, *args)`` for rewards, and ``dispatch(ctx, metrics, *args)`` when termination metrics are present. +Observation terms use the host-side ``ObservationTerm(size, dispatch, *args)`` form. ``size`` is the fixed output width, and additional arguments are passed positionally; simple scalar/array arguments do not require a separate Args class. Reward terms use the analogous ``RewardTerm(dispatch, *args)`` form and return one float32 tensor per environment batch. Termination terms use ``TerminationTerm(dispatch, *args, metrics=...)``; metrics are optional per-environment outputs exposed to the dispatch through ``ctx.metrics``. The resolver validates these values into the static execution metadata, and the Torch executor calls ``dispatch(ctx, out, *args)`` for observations and ``dispatch(ctx, *args)`` for rewards and terminations. -## 5. 编译与性能边界 +## 5. 执行与性能边界 -- 只有纯数组计算进入 Numba nopython kernel;不得传入 `SceneModel`、`SceneData`、`ArrayEnvState`、`NpObs`、config 或 `info`; -- 不提供 object-mode fallback,不自动转换任意 Python/NumPy 环境方法; -- kernel 通过 `prange` 沿 environment batch 维度并行,使用预分配 output 和 term-owned buffers; -- 不默认启用 `fastmath`; -- 单个环境实例不修改进程级 Numba 线程数,线程策略由训练 runtime 管理; -- 编译时间、warm-up 和稳态 step latency 分开统计;绝对性能门槛属于 benchmark,不属于环境行为测试; -- CI 验证 nopython、数值 parity、reset/step 生命周期和 buffer identity。 +- 只有纯数组计算进入 term dispatch;不得传入 `SceneModel`、`SceneData`、`ArrayEnvState`、`NpObs`、config 或 `info`; +- term 一次调用处理整批张量,写入预分配 output 和 term-owned buffers; +- Torch 执行器默认以 `torch.compile` 编译整段 term 调度(可回退 eager),不自动转换任意 Python/NumPy 环境方法; +- 单个环境实例不修改进程级 Torch 线程数,线程策略由训练 runtime 管理; +- 首次调用(惰性 specialization)、warm-up 和稳态 step latency 分开统计;绝对性能门槛属于 benchmark,不属于环境行为测试; +- CI 验证数值行为、reset/step 生命周期和 buffer identity。 ## 6. 不变量 @@ -182,4 +184,4 @@ Observation terms use the host-side ``ObservationTerm(size, dispatch, *args)`` f - action policy layout 与 actuator route 分离,route 非空且不重叠; - sim snapshot 只来自编译后的 read program; - `ctx.env_id` 和所有 `PER_ENV` leaves 使用同一原始 environment id; -- Manager 不引入 `NumbaEnvState`、`NumbaObs`、generic `ManagerValue` registry 或独立 Hydra backend 维度。 +- Manager 不引入 generic `ManagerValue` registry 或独立 Hydra backend 维度。 diff --git a/wiki/design/manager/sim-read-program.md b/wiki/design/manager/sim-read-program.md index d5c18db2..60874f96 100644 --- a/wiki/design/manager/sim-read-program.md +++ b/wiki/design/manager/sim-read-program.md @@ -6,18 +6,18 @@ 框架负责 query 的配置级解析和 key 校验;backend 负责 simulator handle 解析、physical memory planning、相等 query 的去重、 logical view 和 full/partial execute。 -Manager task 的 query 来源和 `required_queries()` 规则见 [Manager Task API](./task-authoring.md);整体 simulator 边界见 +Manager task 的 query 来源(任务 `queries.data` 与 term 的 `SimDataQuery` 参数)见 +[Manager Task API](./task-authoring.md);整体 simulator 边界见 [Manager SimBackend](./sim-backend.md)。 ## 1. 数据流与责任 ```text -cfg.queries.data + observation.required_queries() - -> ManagerBasedEnvCfg.sim_query_cfgs() - -> resolve_sim_queries(key 校验 + EnvCfg resolver) +cfg.queries.data + term args 中的 SimDataQuery + -> ManagerEnv 构造期合并(_merge_query_declarations) -> SimBackend.compile_reads -> PhysicsReadProgram - -> ManagerContext.sim + -> ManagerContext.sim / term 参数视图 ``` | 组件 | 职责 | @@ -71,12 +71,13 @@ program 和 kernel 不保留 callable。 ## 4. Term required queries -声明来源有两处:任务 `queries.data` / `queries.model`,以及 observation term 的 `required_queries()`。Manager 在 backend 编译前合并它们: +声明来源有两处:任务 `queries.data` / `queries.model`,以及 reward/termination/observation term 参数中的 +`SimDataQuery` 对象。Manager 在 backend 编译前合并它们: - 任务显式声明的 key 由任务拥有; -- 任务未声明的 key 由 term contribution 补齐; -- 多个 term 对同一 key 贡献不等 query 时构造期报错; -- 相等 term contributions 折叠为一条声明; +- term 参数中的 query 按相等去重,登记为合成 key; +- 同一 key 下不相等声明构造期报错; +- 相等声明折叠为一条声明、一次 physical read; - data query 进入 read program;model query 在 term construction 时解析为静态参数。 具体协议和限制见 [Term 自声明 Sim Query](./term-required-queries.md)。 diff --git a/wiki/design/manager/task-authoring.md b/wiki/design/manager/task-authoring.md index 28e4ea5c..b79181f2 100644 --- a/wiki/design/manager/task-authoring.md +++ b/wiki/design/manager/task-authoring.md @@ -3,7 +3,7 @@ ## 摘要 本文面向任务开发者,说明如何使用 `ManagerBasedEnvCfg` 和 Manager term 创建一个 config-driven `ManagerEnv`。 -内容覆盖任务文件组织、manager groups、term 创建、`required_queries()`、`@kernel_data` 字段和常用 term 接口。 +内容覆盖任务文件组织、manager groups、term 创建、`SimDataQuery` 参数、`@kernel_data` 字段和常用 term 接口。 运行时边界见 [Manager Runtime](./runtime.md),kernel ABI 见 [Manager Kernel ABI](./kernel-abi.md)。 ## 1. 任务文件组织 @@ -101,26 +101,29 @@ class TrackingErrorTermination: ## 4. Simulator query 声明 -任务可以在 `cfg.queries.data` 和 `cfg.queries.model` 中声明 query。Observation term 也可以通过 `required_queries()` 声明自己需要的 query: +任务可以在 `cfg.queries.data` 和 `cfg.queries.model` 中声明 query;此外,reward、termination 和 +observation term 可以把 `SimDataQuery` 对象直接作为 term 参数传入,声明自己需要的 sim 数据: ```python @configclass(kw_only=True) class RobotPositionObsCfg(ObservationTermCfg): - def required_queries(self) -> SimQueriesCfg: - return SimQueriesCfg( - data={"robot_pos": LinkPositionQuery(link=lambda cfg: cfg.scene.objs.robot.resolved_base_link_name)} + def __call__(self, env): + return ObservationTerm( + size=3, + dispatch=robot_position_obs, + args=(LinkPositionQuery(link=env.cfg.scene.objs.robot.resolved_base_link_name), 0.5), ) ``` -合并规则: +dispatch 收到的对应参数是 read program 提供的整批 sim 视图(`(N, ...)`)。合并规则: -1. 任务显式声明拥有该 key; -2. 任务未声明时,由 term 提供默认 query; -3. 多个 term 对同一 key 提供不相等 query 时,构造期报错; -4. 相等 query 由 backend read program 折叠为一次 physical read,并保留各 key 的 logical view。 +1. Manager 构造期扫描 term 参数中的 `SimDataQuery`,按相等去重登记为合成 key,与任务声明统一合并; +2. 同一 key 下不相等的声明构造期报错; +3. 相等 query 由 backend read program 折叠为一次 physical read,并保留各 key 的 logical view。 -Query key 是 compiled interface;重命名 key 需要同步修改所有 `ctx.sim["key"]` 消费者。 -Term 不持有 backend handle,也不在运行期创建或执行 simulator query。 +完整机制见 [Term 自声明 Sim Query](./term-required-queries.md)。Query key 是 compiled interface;重命名 +key 需要同步修改所有 `ctx.sim["key"]` 消费者。Term 不持有 backend handle,也不在运行期创建或执行 +simulator query。 ## 5. Action diff --git a/wiki/design/manager/term-required-queries.md b/wiki/design/manager/term-required-queries.md index 5de49032..bc60317d 100644 --- a/wiki/design/manager/term-required-queries.md +++ b/wiki/design/manager/term-required-queries.md @@ -1,127 +1,90 @@ -# Term 自声明 Sim Query 设计(required_queries 协议) +# Term 自声明 Sim Query 设计(SimDataQuery 参数) ## 摘要 -本设计解决通用 MDP term(框架库中的 `RobotDofPosRelObsCfg` 等)与任务侧 sim query 声明之间的跨层 key 冲突:框架 term 的 kernel 以编译期字面量 `ctx.sim["robot_dof_pos"]` 消费 sim 数据,而该 key 目前的唯一声明点在任务的 `queries.data` 里,契约完全隐式。方案是给 term cfg 增加自声明协议 `required_queries()`——term 自己声明需要的 data/model query,构造期与任务声明按"key 单一属主、任务优先、term 贡献兜底、冲突报错"合并;相等 query 由 backend 既有的 exact-equal 折叠自动 alias,不产生重复物理读。kernel 访问方式(literal `ctx.sim[key]`)、numba lowering、codegen 全部不变。 +本设计定义 manager term 如何声明自己的 simulator 数据需求:reward、termination 和 observation term 把 +`SimDataQuery` 对象直接作为 term 参数传入(`RewardTerm(dispatch, *args)` 中的 `args` 元素)。`ManagerEnv` +构造期扫描全部 term 参数,把每个 distinct query 登记为合成 key `terms.args[i]`,与任务侧 +`cfg.queries.data` 声明统一合并:同一 key 下不相等的声明构造期报错,相等声明折叠为一次物理读。resolver +(`ManagerProgramResolver`)随后把 query 参数映射为 read program 的批量 sim 视图,Torch 执行器在参数原位置 +把整批张量传给 dispatch。本文覆盖声明机制、合并规则与执行映射;任务侧集中声明方式见 +[Manager Task API](./task-authoring.md),读侧编译契约见 [读侧编译](./sim-read-program.md)。 -本文描述已落地的设计;`required_queries()` 协议和 data/model query 合并已实现。manager kernel 只为环境自身配置构建;同一 runtime 不支持重建另一份 manager 配置。 +## 1. 动机:term 声明 vs 任务声明 -## 1. 问题 +任务通过 `SimQueriesCfg` 字段名集中声明 data query(key 即字段名,如 +`queries.data["robot_dof_pos"]`)。集中声明适合有自身 authority 的任务(如 WBT 的 motion joint 顺序), +但对通用 term 有两个问题: -现状数据流(三个事实): +1. **契约隐式**。term 的 dispatch 需要哪些 sim 输入,任务作者只能阅读 term 源码发现;唯一保障是运行期 + key 访问失败。 +2. **声明权与消费权分离**。通用 term(框架 MDP 库中的 robot observation/reward 等)需要的 query 由任务 + 代为声明,框架每新增一个 term 都可能要求所有使用它的任务补声明。 -1. **key 的声明权在任务层**。`SimDataQueriesCfg` 字段名即 query key,查询参数(关节集合与顺序、link 名)来自任务 authority,如 `WbtSimQueriesCfg.robot_dof_pos = JointPositionQuery(joints=resolve_wbt_joint_names)`。 -2. **key 的消费权在框架层,且被编译期钉死**。`motrix_env_core.mdp.observations` 的 kernel 全部使用 `ctx.sim["robot_dof_pos"]` 字面量;`Map.__getitem__` 的 numba lowering(`prefer_literal` overload)要求 key 是 `StringLiteral`,即 key 必须出现在 kernel 源码里,无法参数化。 -3. **契约隐式**。框架没有任何声明指出"我需要哪些 key",任务作者只能阅读框架 kernel 源码发现约定,唯一保障是构造期 query 访问失败。 +机制目标是让 term 自己携带输入声明:把 `SimDataQuery` 作为 term 参数,声明与消费落在同一个对象上; +同时保留任务侧集中声明,二者在构造期合并,且保持"一个 key 只对应一个 logical query"的单一属主规则。 -冲突本质:同一个字符串 key 承担两个角色——读程序中的布线身份(理应任务所有)与通用量的语义标识(框架所有)——而名字的归属被劈成两层。更深一层的结构性缺陷是:契约面随框架 term 数量线性增长,任何"任务侧集中声明框架所需 key"的方案(如框架发布 canonical query group 基类让任务继承)都会把框架每次新增 term 的成本转嫁给所有任务。 +## 2. 声明机制:SimDataQuery 作为 term 参数 -## 2. 目标与非目标 - -目标: - -- 框架通用 term 自带输入声明;框架新增 term 不要求任务侧新增任何声明字段。 -- 有自身 authority 的任务(如 WBT 的 motion joint 顺序)保持现有声明方式与语义,零迁移。 -- 相同 query 在不同 key 下自动共享物理读(alias),无重复读取。 -- data query 与 model query 统一支持。 -- kernel 访问路径、numba ABI、codegen、编译 fingerprint 机制全部不变。 - -非目标: - -- 不引入 key 参数化/注入路径(不改 [sim-read-program](./sim-read-program.md) §7 "literal `ctx.sim[key]`,不提供参数注入路径"的决策)。 -- 不支持单环境多机器人的 per-entity key 命名空间;canonical key 是单数角色("the robot")。 -- 不改变 "key 等于声明字段名、不维护第二份 registry" 的既有契约,只扩展声明的来源。 - -## 3. 协议设计 - -`ObservationTermCfg` 新增一个默认为空的协议方法: - -```python -@configclass(kw_only=True) -class ObservationTermCfg(abc.ABC): - def required_queries(self) -> SimQueriesCfg: - """本 term 需要的 sim 输入声明。 - - data query 的 key 必须与 kernel 中的 ``ctx.sim`` 字面量一致; - model query 的 key 必须与 ``__call__`` 中 ``env.model`` 的读取 key 一致。 - 返回值必须是 cfg 级可计算的(允许 ``lambda cfg:`` 晚绑定字段),不得依赖已构造的 env。 - """ - return SimQueriesCfg() -``` - -框架通用机器人 term 各自实现(示例): +reward、termination 和 observation term 的 `args` 元组中可以直接包含 `SimDataQuery` 实例: ```python -def _default_robot_dof_pos_query() -> SimDataQuery: - return JointPositionQuery(joints=lambda cfg: resolve_robot_joint_names(cfg)) # 模块级共享 +joint_pos_query = JointPositionQuery(joints=resolve_robot_joint_names) - -class RobotJointPosObsCfg(ObservationTermCfg): - def required_queries(self) -> SimQueriesCfg: - return SimQueriesCfg(data={"robot_dof_pos": _default_robot_dof_pos_query()}) +RewardTerm( + dispatch=robot_joint_pos_reward, + args=(joint_pos_query, -0.5), # dispatch(ctx, joint_pos, penalty_scale) +) ``` -约束: - -- **同一 key 的所有贡献者必须共享同一个模块级默认构造函数**,保证贡献的 query 在解析后精确相等(见 §5)。 -- 默认 query 从 `scene.objs.robot` 推导(全关节 / `resolved_base_link_name`),经 `lambda cfg:` 晚绑定求值;机器人 DOF 默认使用 `BodyJointPositionQuery` / `BodyJointVelocityQuery`,由 `resolved_base_link_name` 确定 body,不需要在 `RobotCfg` 中重复维护关节序。 -- 协议本期只加在 `ObservationTermCfg`(当前唯一需求方)。收集循环按 observation groups 实现;未来核心层新增通用 reward/termination 时将同一方法加到对应 cfg base,收集点按需扩展,不预留抽象。 - -## 4. 合并规则:key 单一属主,任务优先 +`ManagerEnv` 构造时(`motrix_env_core.manager.env`)在 read program 编译之前扫描 reward、termination 和 +全部 observation term 的 `term.args`: -唯一收集点是 `ManagerBasedEnvCfg` 的两个 cfg 级方法(均在 backend 编译之前执行;方法内部先使用现有 resolver 得到 concrete query): +- 每个 distinct query(按 dataclass 相等去重)登记进 `_term_query_keys`,合成 key 为 + `terms.args[i]`(`i` 按收集顺序递增);多个 term 传入相等 query 时共享同一 key。 +- 合成 key 与任务声明一起进入 `_merge_query_declarations` 合并,随后一次性交给 + `SimBackend.compile_reads` 编译为 `PhysicsReadProgram`。 -- data:`sim_query_cfgs()` = `queries.data_dict()` ∪ Σ observation term 的 `required_queries().data_dict()`; -- model:`model_query_cfgs()` = `queries.model` ∪ Σ observation term 的 `required_queries().model`。 +## 3. 合并规则:key 单一属主,不等即报错 -两侧都先经现有 resolver 把 `lambda cfg:` 解析为 concrete query,再按键合并: +`_merge_query_declarations` 对任务声明与 term 贡献统一处理: | 情形 | 结果 | | --- | --- | -| 任务已声明该 key | 与 term 声明重复,构造期失败;term 的特殊语义校验(例如相对关节位置的 key-pose 对齐)继续在 `__call__` 中执行 | -| 无人声明 | term 贡献自己的默认 query | -| 多个来源对同一 key 贡献不等 query | 构造期 `ValueError`,报错信息含 key、两侧 query repr 与贡献者路径(如 `observations.policy.dof_pos`) | - -一个 key 仍然只对应一个 logical query,"rename key 是 compiled interface 变更"的既有契约不变;本设计只增加声明的第二来源,并让任务声明在语义上保持权威。 - -## 5. Alias 与去重语义 +| 仅任务声明该 key | 使用任务声明,语义权威保持在任务侧 | +| 仅 term 贡献该 key | term 的 query 进入 read program | +| 多个来源对同一 key 声明相等 query | 折叠为一个 logical query、一次物理读(read program 的既有 exact-equal 去重不变量,见 [读侧编译](./sim-read-program.md)) | +| 同一 key 下声明不相等 | 构造期 `ValueError`,报错含 key、两侧 query repr 与来源 | -"相同 query 不同 key 自动共享"是 `PhysicsReadProgram` 的既有文档化不变量,本设计直接依赖、不新增机制: +term-arg 的合成 key(`terms.args[i]`)与任务声明的字段名 key 处于同一命名空间:任务对同一字符串 key +的声明与 term 贡献不等时同样构造期失败。"rename key 是 compiled interface 变更"的既有契约不变。 -- `compile_reads` 收到完整声明集(含重复),exact-equal 声明折叠到同一物理区域,每个 key 仍是可服务的别名,`view()` 返回同一视图对象(见 [读侧编译](./sim-read-program.md) §3/§5)。 -- 典型场景:任务把某 query 声明在自定义名字下(如 `motion_dof_pos`),term 的 kernel 字面量是 `robot_dof_pos` 且无人声明该 key → term 贡献默认 query → 与任务声明相等 → 一次物理读、两个别名 key、零重复。 -- 相等是 **lambda 解析后的 dataclass 精确相等**。`joints` 等序敏感字段顺序不同即不相等——那是不同数据(顺序即语义),本来就该各自读取,不算重复。 -- model query 没有折叠契约(按 name 逐个编译),但它是构造期一次性 host 读,重复成本可忽略;且 §4 的单属主规则使同一 key 不会出现两条声明。 +## 4. 执行映射:resolver 到批量 sim 视图 -## 6. Model query 的绑定方式 +`ManagerProgramResolver` 解析 term 参数时(`_resolve_term_argument` 路径)识别 `SimDataQuery` 值: -model query 是静态数据,不经 `ctx.sim`,不参与每步刷新:term 在 `__call__(env)` 时从 `env.model` 读出结果并烘焙进 term 的 `@kernel_data` params(与 `RelativePositionParams.reference` 烘焙 key pose、WBT action 读取 `env.model.others["robot_joint_position_limits"]` 同构)。term 贡献的 model query key 与它烘焙时读取的 key 是同一个类级常量,天然一致。 +- 以 query 反查 `_term_query_keys` 得到 key,再取其在 read program 中的 `ManagerSimInput`(logical + view + query 绑定); +- dispatch 调用时,该参数位置收到的是 context 的**整批** sim 视图张量(`(N, ...)`),与 + `ctx.sim[key]` 的批量语义一致(见 [Torch 执行器设计](./torch-executor.md)); +- view 的 shape/stride 进入 resolver 的 layout fingerprint,保证输入布局变化触发重新解析。 -## 7. Kernel 字面量的所有权翻转 +term 因此不需要在 dispatch 内部用字符串 key 访问 `ctx.sim`:query 声明即参数,参数即批量视图。 -实现本协议后: +## 5. 边界与已知限制 -- `ctx.sim["robot_dof_pos"]` 字面量不变,但该 key 由框架 term 的 `required_queries()` 声明进 read program——**字面量回到了它的声明者手里**,跨层依赖方向被理顺; -- 任务侧 kernel(WBT rewards/terminations/actions 读取 `robot_dof_pos` 等)行为不变:key 仍存在,数据仍来自同一声明; -- 单一 literal 约束(每 term 类一个 key)保持现状表达力:同一 term 类的两个实例若需要不同 query,仍不支持,与现状一致。 - -## 8. 与现有任务和文档的关系 - -- **WBT:零迁移。** `WbtSimQueriesCfg` 现有声明按 §4 任务优先规则继续胜出;可选择删除与框架默认相等的声明以精简,但非必需。`WbtManagerEnvCfg.__post_init__` 对 `queries.data` 具体类型的校验不受影响。 -- **[IsaacLab ManagerCfg Adapter](./isaaclab-manager-cfg-adapter.md):** 适配层把 `SceneEntityCfg` 式 selector 归一化为"去重后的显式 SimDataQuery",本协议是该归一化在原生侧的落点——适配器生成的 concrete term cfg 携带自己的 `required_queries()`,无需适配器另行维护 key 分配。 -- **[Manager Task API](./task-authoring.md) 与 [读侧编译](./sim-read-program.md):** query 声明来源和合并规则分别在任务开发指南与本篇定义。 -- **诊断:** `SimInputLayout.consumers` 可从常量 `"ManagerContext.sim"` 扩展为贡献者 cfg 路径(可选,不阻塞)。 - -## 9. 边界与已知限制 - -- **`_scene_robot_cfg` 的场景扫描**(`RobotDofPosRelObsCfg` 为 `BodyJointPositionQuery` 分支寻找 key-pose robot)是同类的隐式"the robot"角色,另行处理,不阻塞本设计。 -- **多机器人**:单环境多机器人需要 per-entity key 命名空间或注入式绑定,届时另立设计;本设计明确不支持。 +- 收集范围是 reward、termination 和 observation term 的 `args`;command/action term 的 kernel_data + 走各自的 canonicalize 路径,不参与本机制。 +- 合成 key `terms.args[i]` 依赖收集顺序(配置映射顺序),属于内部执行元数据,不是面向任务的稳定命名 + 契约;任务需要稳定 key 时应在 `queries.data` 中显式声明。 +- 不支持单环境多机器人的 per-entity key 命名空间;canonical key 是单数角色("the robot")。 +- model query 是静态数据,不经 read program:term 在 `__call__(env)` 时从 `env.model` 读出并烘焙进 + term 的 `@kernel_data` params,与本机制无关。 ## 验收标准 -- 使用框架通用机器人 term 的任务,若其默认 query 适用,`queries.data` / `queries.model` 中无需任何 robot 相关声明,环境可构造并训练。 -- 任务对某 key 的显式声明总是胜出;term 行为与该声明一致;缺失 key 或 backend 无法解析 query 时在构造期失败。 -- 相同 query 的不同 key 声明在 read program 中折叠为一次物理读,别名 `view()` 为同一对象(既有 backend 测试覆盖 + 新增协议路径用例)。 -- 两个 term 对同一 key 贡献不等 query 时构造期报错,报错含 key 与双方来源。 -- WBT 全部现有测试在零配置改动下通过。 -- kernel 源码、numba lowering、codegen、`plan_key` 指纹机制无结构性变化。 +- 传入 `SimDataQuery` 参数的 term 无需任务在 `queries.data` 中重复声明,环境可构造并训练;dispatch 在 + 参数位置收到正确 shape 的批量张量。 +- 相等 query(无论来自任务声明还是多个 term 参数)在 read program 中折叠为一次物理读,view 为同一对象。 +- 同一 key 下不等声明在构造期报错,报错含 key 与两侧来源。 +- 使用任务侧集中声明的现有任务(如 WBT)零迁移。 diff --git a/wiki/design/manager/torch-executor.md b/wiki/design/manager/torch-executor.md new file mode 100644 index 00000000..422f7f76 --- /dev/null +++ b/wiki/design/manager/torch-executor.md @@ -0,0 +1,107 @@ +# Torch Manager 执行器设计 + +## 摘要 + +本文定义 `ManagerEnv` 的数值执行器:通用 Torch 批量执行器 +(`motrix_env_core.torch_manager.runtime.TorchManagerRuntime`)。resolver +(`motrix_env_core.manager.resolver.ManagerProgramResolver`)把 manager 配置 +解析为 read plan / layout / prepared invocations 等执行元数据,Torch 执行器 +以整批张量按固定语句顺序调度 term;host 生命周期与 term 配置 API 不变。 +Torch 是唯一执行器——历史上并存过的 Numba fused-kernel 路径与 +`executor="numba"|"torch"` 双执行器切换已经移除。相关背景见 +[Manager Runtime 设计](./runtime.md)。 + +执行方式默认为 `torch.compile`:每个 kernel(evaluate/observe/reset)把整段 +term 调度编译为 Inductor 图(observe 限制 `cpp.threads=min(4, nthreads)`, +可用 `MOTRIX_TORCH_COMPILE_THREADS` 覆盖;reset 以 env_ids 为动态维度参数), +context 命名空间、invocation 参数与张量视图按 inputs 身份一次性缓存。 +RNG draw 顺序与 eager 完全一致;Inductor 重结合仅带来 ≤1e-8 量级的 reward +差异。调试可用 `TorchManagerRuntime(env, program, compile=False)` 或 +`runtime.compile_enabled = False` 回到逐 term eager 执行。 + +## 解析与执行的分工 + +- `ManagerProgramResolver`(backend-neutral)负责输入准备:flatten + `@kernel_data` term 与 `ManagerContext`,构建 `ManagerLayout`、 + `KernelInputSource`/`ReadPlan`、`PreparedInvocation`(每个 term 的 + dispatcher、参数表达式/值/prepared index),并校验 `@dispatch` 标记与 + 签名契约。它不做任何代码生成、缓存或 dispatcher 编译。 +- `TorchManagerRuntime` 消费这些元数据并执行 term 调度;所有写入落回 + host 持有的 numpy 缓冲(reward terms/masks/obs/command state/metrics), + `ManagerEnv` 的公共 API 与状态面保持不变。 + +## 执行器行为 + +- evaluate 按固定语句顺序调度:command `update` → rewards + (raw、`weighted = raw * weight * dt`、Σ)→ terminations(masks、OR)→ + command `advance`(rewards/terminations 读 pre-advance command state); + observe 按 group/term 顺序写 obs 切片;reset = command `reset_env` + (lane-gathered)→ sim-reset terms 写 compact 行。 +- RNG:per-lane SplitMix64(`torch_manager.rand.draw_uniform`),批量语义为 + term 一次 draw `count=k` 等价于 per-lane 连续调用 k 次。 +- 输入元组持久复用:simulator-backed 数组由 read program 原地重写, + `_TensorCache` 在同一输入元组上缓存 tensor 视图。 +- reset 路径对 `PER_ENV` slot 做 lane gather,term 写 gathered 行后 scatter + 回常驻缓冲;compact sim-reset write 缓冲布局与 `SimResetRuntime.buffers` + 一致。 + +## 批量 term 约定 + +ctx(`_BatchContext`)字段全部是**整批**张量: + +- `ctx.sim[key]` → `(N, ...)` 张量(SimDataQuery 参数同样是整批张量) +- `ctx.commands[name].field` → per-env 字段 `(N, ...)`;嵌套 kernel_data(如 + `motion.clip.joint_pos`)为 shared `(F, J)`,term 自行按 `steps[:, 0]` gather +- `ctx.actions[name].current()/previous()` → `(N, A)`;`action_queue/action_ptr` + 同理 +- `ctx.metrics[name]` → `(N, 1)` 可写张量 +- `ctx.rand.next_uniform(count=k)` → `(N, k)`,`[-1,1)`,原地推进状态; + `ctx.rand.uniform_range(count, low, high)` 同理 +- `ctx.dt` float、`ctx.env_id` `(N,)`、`ctx.sim_reset_requested` `(N,1)` bool + +函数签名约定(第一个参数注解为 `ManagerContext`,observation 的 `out` +注解为 `torch.Tensor`(`(N, size)` 切片),reset 的 `sim_writes` 注解为 +`Map[torch.Tensor]`——resolver 校验这些契约注解;标量参数到达时为 Python +float,数组参数为批量张量): + +- reward:返回 `(N,)` float32 张量 +- termination:返回 `(N,)` bool 张量;metric 写 `ctx.metrics[name][:, 0] = value` +- observation:`out` 为 `(N, size)` 张量切片,直接写入;noise 用 + `ctx.rand.next_uniform(size) * amplitude`(count == out.shape[1] 保持 RNG 顺序) +- reset:`(ctx, sim_writes, *args)`,ctx 为 lane-gathered(第 0 维 = + len(env_ids)),`sim_writes[name]` 为 compact `(M, *write_shape)` 张量 + +Command hooks(`update`/`advance`/`reset_env`)是 CommandTerm 类上的 +`@staticmethod` + `@dispatch` 函数,签名为 `(receiver, ctx)`;`receiver` 是该 +command 的 kernel-data namespace 视图(batch/reset 场景下分别为整批/gathered +行)。host 侧 `reset(ResetContext)` 与 `on_transition()` 保持 numpy host 实现, +ActionTerm 的 `process/_reset` 也保持 numpy host 路径不变。 + +标量参数(float)以 Python float 传入;字符串参数直接传入。 + +## 参考 math 模块 + +- `motrix_envs/.../wbt/mdp/batch_math.py`(张量四元数)、 + `batch_rewards/batch_terminations/batch_observations`(WBT 项的批量数学) + 与 `walk_manager_mdp/terrain.heightfield_lookup_batch`(批量地形插值)。 + +## 性能特性 + +以下为 G1 locomotion(2048 envs、zero actions、100 steps)迁移到单 Torch 执行器时的基准记录, +用于解释当前实现中的若干设计选择: + +- **机器对比**。EPYC(192 logical CPUs,双方 16 线程,PASSIVE):Numba 基线 median 12.99ms + (157,641 steps/s),Torch 迁移后 median 14.0–14.5ms;breakdown 差距集中在 observe + (3.31 → 1.85ms)与 reset kernel(0.80 → 0.72ms,Numba 0.19ms),evaluate 反超 + (1.94 → 1.66 vs Numba 1.78)。本机(32 核桌面)已打平:16.99 vs 16.98ms + (evaluate 0.83 反超 Numba 0.92)。 +- **`min_chunk_size` 发现**。图本身已是单融合 kernel;真正瓶颈是 Inductor `cpp.min_chunk_size` + 默认 4096——manager 批量循环(2048×3…2048×84)全部低于阈值,被生成为 `#pragma omp single` + 单线程段(observe 图约 19 段)。修复:`min_chunk_size=256`(env `MOTRIX_TORCH_MIN_CHUNK`)、 + 默认图内线程 4→8、observe 改 per-term 连续缓冲 + 单 kernel 组装、输出张量改为图参数使双缓冲槽 + 共享编译图、observe 按 group 拆分独立编译图。图内线程扫描( + `MOTRIX_TORCH_COMPILE_THREADS=1/2/3/4/8/16`)4 最优;numactl 在单 NUMA node 上无效。 +- **冷惩罚测量**。物理后首个编译 kernel 存在冷缓存惩罚:同段计算重复执行 0.29ms vs 物理后执行 + 0.59ms(Numba 因单小并行区天然免疫)。 +- **已知残余**。EPYC 上相对 Numba 约 1ms 残余 = 冷缓存惩罚 + reset 静态全批布尔掩码设计 + (RNG 状态 `where(mask, ...)` 门控)带来的约 3x lane 计算结构性成本;本机持平 16.8–17.0ms。 From f45a4d2ce7bf91ecf34b7d20d50cde4d4b6b863c Mon Sep 17 00:00:00 2001 From: zilch <147668916@qq.com> Date: Thu, 5 Nov 2026 21:30:45 +0800 Subject: [PATCH 02/41] chore(bench): drop the concluded batch-mdp microbenchmark 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. --- bench/bench_batch_mdp.py | 103 --------------------------------------- 1 file changed, 103 deletions(-) delete mode 100644 bench/bench_batch_mdp.py diff --git a/bench/bench_batch_mdp.py b/bench/bench_batch_mdp.py deleted file mode 100644 index c8c640be..00000000 --- a/bench/bench_batch_mdp.py +++ /dev/null @@ -1,103 +0,0 @@ -# Copyright Motphys Technology Co., Ltd. 2025, 2026 -# SPDX-License-Identifier: Apache-2.0 - -"""One batch MDP source executed by NumPy and Torch; no simulator required.""" - -import argparse -import json -from time import perf_counter - -import numpy as np -import torch - -from motrix_env_core.array import mx -from motrix_env_core.manager.dispatch import dispatch - - -@dispatch -def batch_rewards(position, default, weights, action, previous, velocity, command): - delta = position - default - pose = mx.sum(delta * delta * weights, axis=-1) - change = action - previous - rate = mx.sum(change * change, axis=-1) - error = velocity - command - tracking = mx.exp(-mx.sum(error * error, axis=-1)) - return pose, rate, tracking - - -def lane_rewards(position, default, weights, action, previous, velocity, command): - """Independent scalar oracle and lane-fused performance reference.""" - count, joints = position.shape - pose = np.empty(count, dtype=position.dtype) - rate = np.empty(count, dtype=position.dtype) - tracking = np.empty(count, dtype=position.dtype) - for lane in range(count): - p = np.float32(0) - r = np.float32(0) - e = np.float32(0) - for joint in range(joints): - delta = position[lane, joint] - default[joint] - p += delta * delta * weights[joint] - change = action[lane, joint] - previous[lane, joint] - r += change * change - for component in range(velocity.shape[1]): - error = velocity[lane, component] - command[lane, component] - e += error * error - pose[lane] = p - rate[lane] = r - tracking[lane] = np.exp(-e) - return pose, rate, tracking - - -def inputs(count, joints, seed=7): - rng = np.random.default_rng(seed) - shapes = ((count, joints), (joints,), (joints,), (count, joints), (count, joints), (count, 3), (count, 3)) - return tuple(rng.normal(size=shape).astype(np.float32) for shape in shapes) - - -def measure(function, args, repeats): - start = perf_counter() - result = function(*args) - first_ms = (perf_counter() - start) * 1000 - for _ in range(10): - function(*args) - samples = [] - for _ in range(7): - start = perf_counter() - for _ in range(repeats): - function(*args) - samples.append((perf_counter() - start) * 1e6 / repeats) - return result, {"first_call_ms": first_ms, "median_us": float(np.median(samples))} - - -def main(): - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--num-envs", type=int, default=256) - parser.add_argument("--joints", type=int, default=29) - parser.add_argument("--threads", type=int, default=4) - parser.add_argument("--repeats", type=int, default=200) - options = parser.parse_args() - torch.set_num_threads(options.threads) - arrays = inputs(options.num_envs, options.joints) - tensors = tuple(torch.from_numpy(value) for value in arrays) - expected = lane_rewards(*arrays) - functions = { - "numpy_batch": (batch_rewards, arrays), - "numpy_lane": (lane_rewards, arrays), - "torch_eager_cpu": (batch_rewards, tensors), - "torch_inductor_cpu": (torch.compile(batch_rewards, fullgraph=True), tensors), - } - report = {"num_envs": options.num_envs, "joints": options.joints, "threads": options.threads, "results": {}} - for name, (function, args) in functions.items(): - result, timing = measure(function, args, options.repeats) - for actual, reference in zip(result, expected, strict=True): - if isinstance(actual, torch.Tensor): - actual = actual.detach().numpy() - np.testing.assert_allclose(actual, reference, rtol=3e-5, atol=3e-5) - assert actual.dtype == reference.dtype - report["results"][name] = timing - print(json.dumps(report, indent=2)) - - -if __name__ == "__main__": - main() From 97327f0745103e6e61bc25946f1863c94b33fe18 Mon Sep 17 00:00:00 2001 From: zilch <147668916@qq.com> Date: Thu, 5 Nov 2026 23:17:48 +0800 Subject: [PATCH 03/41] fix(ci): drop the removed tbb/numba extras and markers from CI 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. --- .github/workflows/README.md | 17 ++++++------- .github/workflows/tests.yml | 25 +++---------------- .../motrix_env_core/math/quaternion_array.py | 1 - 3 files changed, 11 insertions(+), 32 deletions(-) diff --git a/.github/workflows/README.md b/.github/workflows/README.md index dd6f9c24..abf11c73 100644 --- a/.github/workflows/README.md +++ b/.github/workflows/README.md @@ -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 @@ -46,21 +45,19 @@ 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, diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index 79cd2f8a..f47388af 100644 --- a/.github/workflows/tests.yml +++ b/.github/workflows/tests.yml @@ -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: @@ -64,7 +47,7 @@ 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 }}"' @@ -72,8 +55,8 @@ jobs: 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 diff --git a/motrix_env_core/src/motrix_env_core/math/quaternion_array.py b/motrix_env_core/src/motrix_env_core/math/quaternion_array.py index 9f7a222a..e347e896 100644 --- a/motrix_env_core/src/motrix_env_core/math/quaternion_array.py +++ b/motrix_env_core/src/motrix_env_core/math/quaternion_array.py @@ -113,7 +113,6 @@ def rotate_inverse_components(quaternion, vector): ``quaternion`` is ``(..., 4)`` and ``vector`` is ``(..., 3)``; the rotated components are returned as three tensors of the leading shape. """ - xp = array_namespace(quaternion, vector) x, y, z, w = (quaternion[..., i] for i in range(4)) vx, vy, vz = (vector[..., i] for i in range(3)) cx = y * vz - z * vy - w * vx From a41325f5f4ed35f86ccbe1db70e44100c5058044 Mon Sep 17 00:00:00 2001 From: zilch <147668916@qq.com> Date: Thu, 5 Nov 2026 23:23:08 +0800 Subject: [PATCH 04/41] refactor(core)!: remove the unused shared-primitive namespace 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. --- .github/workflows/README.md | 34 +---- .../src/motrix_env_core/__init__.py | 2 - .../src/motrix_env_core/array/__init__.py | 5 - .../src/motrix_env_core/array/mx.py | 14 -- .../src/motrix_env_core/array/primitives.py | 74 ---------- .../src/motrix_env_core/array_api.py | 136 ------------------ .../src/motrix_env_core/manager/__init__.py | 3 - motrix_env_core/tests/test_array_api.py | 74 ---------- motrix_env_core/tests/test_array_wrapper.py | 25 ---- motrix_env_core/tests/test_batch_mdp.py | 47 ------ 10 files changed, 1 insertion(+), 413 deletions(-) delete mode 100644 motrix_env_core/src/motrix_env_core/array/mx.py delete mode 100644 motrix_env_core/src/motrix_env_core/array/primitives.py delete mode 100644 motrix_env_core/src/motrix_env_core/array_api.py delete mode 100644 motrix_env_core/tests/test_array_api.py delete mode 100644 motrix_env_core/tests/test_array_wrapper.py delete mode 100644 motrix_env_core/tests/test_batch_mdp.py diff --git a/.github/workflows/README.md b/.github/workflows/README.md index abf11c73..0205c4d6 100644 --- a/.github/workflows/README.md +++ b/.github/workflows/README.md @@ -39,7 +39,7 @@ 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): @@ -63,38 +63,6 @@ regressions. The existing nightly smoke test automatically covers all registered 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` diff --git a/motrix_env_core/src/motrix_env_core/__init__.py b/motrix_env_core/src/motrix_env_core/__init__.py index 787b7b37..a2fd0180 100644 --- a/motrix_env_core/src/motrix_env_core/__init__.py +++ b/motrix_env_core/src/motrix_env_core/__init__.py @@ -12,7 +12,6 @@ """ from motrix_env_core import registry -from motrix_env_core.array import Array from motrix_env_core.base import ABEnv, EnvCfg, ObsSpace, SimCfg from motrix_env_core.config import configclass from motrix_env_core.config.scene import SceneCfg @@ -31,7 +30,6 @@ __all__ = [ "ABEnv", - "Array", "BoundedGamePadPlanarVelocityBinding", "CommandBinding", "ConstantPlanarVelocityBinding", diff --git a/motrix_env_core/src/motrix_env_core/array/__init__.py b/motrix_env_core/src/motrix_env_core/array/__init__.py index 3949e63f..a655a514 100644 --- a/motrix_env_core/src/motrix_env_core/array/__init__.py +++ b/motrix_env_core/src/motrix_env_core/array/__init__.py @@ -1,7 +1,2 @@ # Copyright Motphys Technology Co., Ltd. 2025, 2026 # SPDX-License-Identifier: Apache-2.0 - -from motrix_env_core.array import mx -from motrix_env_core.array_api import Array - -__all__ = ["Array", "mx"] diff --git a/motrix_env_core/src/motrix_env_core/array/mx.py b/motrix_env_core/src/motrix_env_core/array/mx.py deleted file mode 100644 index 3e507941..00000000 --- a/motrix_env_core/src/motrix_env_core/array/mx.py +++ /dev/null @@ -1,14 +0,0 @@ -# Copyright Motphys Technology Co., Ltd. 2025, 2026 -# SPDX-License-Identifier: Apache-2.0 - -"""Shared MDP primitive namespace. - -The functions in this module are the only array operations that shared MDP -code should call when it must run eagerly on NumPy and Torch values. Compiler -lowerings may replace these calls in compiled regions; this module itself is -kept small and transparent for host-side tests. -""" - -from motrix_env_core.array.primitives import clip, exp, sum, where, zeros_like - -__all__ = ["clip", "exp", "sum", "where", "zeros_like"] diff --git a/motrix_env_core/src/motrix_env_core/array/primitives.py b/motrix_env_core/src/motrix_env_core/array/primitives.py deleted file mode 100644 index f27f0ad6..00000000 --- a/motrix_env_core/src/motrix_env_core/array/primitives.py +++ /dev/null @@ -1,74 +0,0 @@ -# Copyright Motphys Technology Co., Ltd. 2025, 2026 -# SPDX-License-Identifier: Apache-2.0 - -"""Backend-neutral primitive namespace for shared MDP functions. - -Public functions dispatch eagerly to native NumPy or Torch arrays. -""" - -from __future__ import annotations - -import numpy as np -from array_api_compat import is_torch_array - -from motrix_env_core.array_api import Array, NativeArray, unwrap_array - - -def sum(x: Array | NativeArray, axis: int | tuple[int, ...] | None = None, *, keepdims: bool = False) -> NativeArray: - x = unwrap_array(x) - if isinstance(x, np.ndarray): - return np.sum(x, axis=axis, keepdims=keepdims) - if is_torch_array(x): - import torch - - return torch.sum(x, dim=axis, keepdim=keepdims) - raise TypeError(f"Unsupported array type {type(x).__name__}.") - - -def exp(x: Array | NativeArray) -> NativeArray: - x = unwrap_array(x) - if isinstance(x, np.ndarray): - return np.exp(x) - if is_torch_array(x): - import torch - - return torch.exp(x) - raise TypeError(f"Unsupported array type {type(x).__name__}.") - - -def where(condition: Array | NativeArray, x: Array | NativeArray, y: Array | NativeArray) -> NativeArray: - condition = unwrap_array(condition) - x = unwrap_array(x) - y = unwrap_array(y) - if isinstance(condition, np.ndarray): - return np.where(condition, x, y) - if is_torch_array(condition): - import torch - - return torch.where(condition, x, y) - raise TypeError(f"Unsupported array type {type(condition).__name__}.") - - -def clip(x: Array | NativeArray, low: float | int, high: float | int) -> NativeArray: - x = unwrap_array(x) - if isinstance(x, np.ndarray): - return np.clip(x, low, high) - if is_torch_array(x): - import torch - - return torch.clamp(x, min=low, max=high) - raise TypeError(f"Unsupported array type {type(x).__name__}.") - - -def zeros_like(x: Array | NativeArray) -> NativeArray: - x = unwrap_array(x) - if isinstance(x, np.ndarray): - return np.zeros_like(x) - if is_torch_array(x): - import torch - - return torch.zeros_like(x) - raise TypeError(f"Unsupported array type {type(x).__name__}.") - - -__all__ = ["clip", "exp", "sum", "where", "zeros_like"] diff --git a/motrix_env_core/src/motrix_env_core/array_api.py b/motrix_env_core/src/motrix_env_core/array_api.py deleted file mode 100644 index 9d15bdd8..00000000 --- a/motrix_env_core/src/motrix_env_core/array_api.py +++ /dev/null @@ -1,136 +0,0 @@ -# Copyright Motphys Technology Co., Ltd. 2025, 2026 -# SPDX-License-Identifier: Apache-2.0 - -"""Backend-neutral array values and namespace helpers. - -``Array`` is a host-side value object. Compiled Manager regions must receive its -``native`` storage instead of the wrapper. -""" - -from __future__ import annotations - -from types import ModuleType -from typing import Any, TypeGuard, Union, overload - -import numpy as np -import numpy.typing as npt -from array_api_compat import array_namespace as _array_namespace -from array_api_compat import is_numpy_array, is_torch_array - -NativeArray = Union[npt.NDArray[Any], "torch.Tensor"] # noqa: F821 - - -class Array: - """A small host-side wrapper around a NumPy array or Torch tensor.""" - - __slots__ = ("_value",) - - def __init__(self, value: NativeArray | Array) -> None: - if isinstance(value, Array): - value = value._value - if not is_array(value): - raise TypeError(f"Array only supports NumPy arrays and Torch tensors, got {type(value).__name__}.") - self._value = value - - @property - def native(self) -> NativeArray: - """Return the backing native array without copying.""" - return self._value - - @property - def shape(self) -> tuple[int, ...]: - return tuple(self._value.shape) - - @property - def dtype(self) -> Any: - return self._value.dtype - - @property - def device(self) -> Any: - return getattr(self._value, "device", None) - - def numpy(self) -> np.ndarray: - if isinstance(self._value, np.ndarray): - return self._value - return self._value.detach().cpu().numpy() - - def torch(self, *, device: Any = None): - import torch - - if isinstance(self._value, torch.Tensor): - return self._value if device is None else self._value.to(device) - return torch.as_tensor(self._value, device=device) - - def __array_namespace__(self, *, api_version: str | None = None) -> ModuleType: - del api_version - return array_namespace(self._value) - - def __getitem__(self, index: Any) -> Array: - return Array(self._value[index]) - - def _binary(self, other: Any, operation) -> Array: - other = unwrap_array(other) - return Array(operation(self._value, other)) - - def __add__(self, other: Any) -> Array: - return self._binary(other, lambda left, right: left + right) - - def __radd__(self, other: Any) -> Array: - return self._binary(other, lambda left, right: right + left) - - def __sub__(self, other: Any) -> Array: - return self._binary(other, lambda left, right: left - right) - - def __rsub__(self, other: Any) -> Array: - return self._binary(other, lambda left, right: right - left) - - def __mul__(self, other: Any) -> Array: - return self._binary(other, lambda left, right: left * right) - - def __rmul__(self, other: Any) -> Array: - return self._binary(other, lambda left, right: right * left) - - def __truediv__(self, other: Any) -> Array: - return self._binary(other, lambda left, right: left / right) - - def __rtruediv__(self, other: Any) -> Array: - return self._binary(other, lambda left, right: right / left) - - def __neg__(self) -> Array: - return Array(-self._value) - - def __repr__(self) -> str: - return f"Array({self._value!r})" - - -@overload -def unwrap_array(value: Array) -> NativeArray: ... - - -@overload -def unwrap_array(value: NativeArray) -> NativeArray: ... - - -def unwrap_array(value: Array | NativeArray) -> NativeArray: - return value.native if isinstance(value, Array) else value - - -def array_namespace(*arrays: Array | NativeArray) -> ModuleType: - """Return the NumPy or Torch Array API namespace shared by ``arrays``.""" - native = tuple(unwrap_array(array) for array in arrays) - for array in native: - if not is_array(array): - raise TypeError( - f"array_namespace only supports NumPy arrays and Torch tensors, got {type(array).__name__}." - ) - return _array_namespace(*native) - - -def is_array(value: object) -> TypeGuard[NativeArray | Array]: - """Return whether ``value`` is a supported array or host wrapper.""" - if isinstance(value, Array): - return True - return (is_numpy_array(value) and isinstance(value, np.ndarray)) or is_torch_array(value) - - -__all__ = ["Array", "NativeArray", "array_namespace", "is_array", "unwrap_array"] diff --git a/motrix_env_core/src/motrix_env_core/manager/__init__.py b/motrix_env_core/src/motrix_env_core/manager/__init__.py index 2a4c4f3e..5f3c9752 100644 --- a/motrix_env_core/src/motrix_env_core/manager/__init__.py +++ b/motrix_env_core/src/motrix_env_core/manager/__init__.py @@ -3,7 +3,6 @@ """Public API for manager-based environments and terms.""" -from motrix_env_core.array import Array, mx from motrix_env_core.config.sim_reset import ManagerResetCfg, ResetTermCfg from motrix_env_core.manager.actions import ( ActionCfg, @@ -38,9 +37,7 @@ __all__ = [ "ActionCfg", "ActionTerm", - "Array", "ActionState", - "mx", "CommandCfg", "CommandTerm", "ContextMap", diff --git a/motrix_env_core/tests/test_array_api.py b/motrix_env_core/tests/test_array_api.py deleted file mode 100644 index 7ce1653c..00000000 --- a/motrix_env_core/tests/test_array_api.py +++ /dev/null @@ -1,74 +0,0 @@ -# Copyright Motphys Technology Co., Ltd. 2025, 2026 -# SPDX-License-Identifier: Apache-2.0 - -from typing import Any - -import numpy as np -import numpy.typing as npt -import pytest -import torch - -from motrix_env_core.array_api import Array, array_namespace, is_array - - -class _OtherArrayApiObject: - def __array_namespace__(self, *, api_version: str | None = None) -> object: - del api_version - return np - - -def _statistics(x: Array, y: Array) -> Array: - xp = array_namespace(x, y) - return xp.mean(x, axis=0) + 2 * xp.std(y, axis=0) - - -def _to_numpy(value: Array) -> npt.NDArray[Any]: - if isinstance(value, torch.Tensor): - return value.detach().cpu().numpy() - return value - - -def test_is_array_recognizes_supported_runtime_arrays() -> None: - assert is_array(np.zeros(1, dtype=np.float32)) - assert is_array(torch.zeros(1)) - assert not is_array([0.0]) - assert not is_array(np.float32(0.0)) - assert not is_array(_OtherArrayApiObject()) - - -def test_array_namespace_rejects_other_array_api_backends() -> None: - with pytest.raises(TypeError, match="only supports NumPy arrays and Torch tensors"): - array_namespace(_OtherArrayApiObject()) - - -def test_array_namespace_rejects_mixed_supported_backends() -> None: - with pytest.raises(TypeError, match="Multiple namespaces"): - array_namespace(np.zeros(1, dtype=np.float32), torch.zeros(1)) - - -@pytest.mark.parametrize( - ("x", "y", "expected_type"), - [ - pytest.param( - np.asarray([[1.0, 2.0], [3.0, 4.0]], dtype=np.float32), - np.asarray([[1.0, 3.0], [5.0, 7.0]], dtype=np.float32), - np.ndarray, - id="numpy", - ), - pytest.param( - torch.tensor([[1.0, 2.0], [3.0, 4.0]], dtype=torch.float32), - torch.tensor([[1.0, 3.0], [5.0, 7.0]], dtype=torch.float32), - torch.Tensor, - id="torch", - ), - ], -) -def test_array_namespace_runs_backend_neutral_statistics( - x: Array, - y: Array, - expected_type: type[object], -) -> None: - result = _statistics(x, y) - - assert isinstance(result, expected_type) - np.testing.assert_allclose(_to_numpy(result), np.asarray([6.0, 7.0], dtype=np.float32)) diff --git a/motrix_env_core/tests/test_array_wrapper.py b/motrix_env_core/tests/test_array_wrapper.py deleted file mode 100644 index 1dbd20bf..00000000 --- a/motrix_env_core/tests/test_array_wrapper.py +++ /dev/null @@ -1,25 +0,0 @@ -# Copyright Motphys Technology Co., Ltd. 2025, 2026 -# SPDX-License-Identifier: Apache-2.0 - -import numpy as np -import torch - -from motrix_env_core.array_api import Array - - -def test_array_wrapper_preserves_numpy_storage_and_operators() -> None: - source = np.asarray([1.0, 2.0], dtype=np.float32) - value = Array(source) - result = (value + 1.0) * 2.0 - - assert value.native is source - np.testing.assert_allclose(result.native, [4.0, 6.0]) - - -def test_array_wrapper_preserves_torch_storage_and_operators() -> None: - source = torch.tensor([1.0, 2.0]) - value = Array(source) - result = (value + 1.0) * 2.0 - - assert value.native is source - torch.testing.assert_close(result.native, torch.tensor([4.0, 6.0])) diff --git a/motrix_env_core/tests/test_batch_mdp.py b/motrix_env_core/tests/test_batch_mdp.py deleted file mode 100644 index 39ce683c..00000000 --- a/motrix_env_core/tests/test_batch_mdp.py +++ /dev/null @@ -1,47 +0,0 @@ -# Copyright Motphys Technology Co., Ltd. 2025, 2026 -# SPDX-License-Identifier: Apache-2.0 - -"""Shared batch reductions preserve environment independence and native dtype.""" - -import numpy as np -import pytest -import torch - -from motrix_env_core.array import mx -from motrix_env_core.manager.dispatch import dispatch - - -@dispatch -def _batch_pose(position, default, weights): - delta = position - default - return mx.sum(delta * delta * weights, axis=-1) - - -@pytest.mark.parametrize("shape", [(1, 1), (7, 29), (4, 0)]) -def test_shared_batch_pose_preserves_lane_independence(shape): - rng = np.random.default_rng(5) - position = rng.normal(size=shape).astype(np.float32) - default = rng.normal(size=shape[1]).astype(np.float32) - weights = rng.uniform(size=shape[1]).astype(np.float32) - expected = np.sum((position - default) ** 2 * weights, axis=-1) - for native_args in ( - (position, default, weights), - tuple(torch.from_numpy(value) for value in (position, default, weights)), - ): - result = _batch_pose(*native_args) - assert np.asarray(result).shape == (shape[0],) - assert np.asarray(result).dtype == np.float32 - np.testing.assert_allclose(np.asarray(result), expected, rtol=2e-6) - - -def test_sum_supports_runtime_axes_and_strided_input(): - value = np.arange(48, dtype=np.float32).reshape(4, 12)[:, ::2] - for axis in (-2, -1, 0, 1): - np.testing.assert_allclose(mx.sum(value, axis=axis), np.sum(value, axis=axis)) - for axis in (-3, 2): - with pytest.raises(ValueError, match="bounds"): - mx.sum(np.asarray(value), axis=axis) - - -def test_sum_without_axis(): - assert mx.sum(np.ones((2, 3), dtype=np.float32)) == 6 From 68f97b3c08f167d2d86a35b7a46af3f208bdbcc1 Mon Sep 17 00:00:00 2001 From: zilch <147668916@qq.com> Date: Thu, 5 Nov 2026 23:33:57 +0800 Subject: [PATCH 05/41] refactor(manager)!: drop the @dispatch requirement from manager term 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). --- AGENTS.md | 10 +-- .../src/motrix_env_core/manager/resolver.py | 15 ++--- .../src/motrix_env_core/manager/terms.py | 2 - .../src/motrix_env_core/mdp/observations.py | 8 --- .../src/motrix_env_core/mdp/rewards.py | 5 -- .../src/motrix_env_core/mdp/terminations.py | 3 - motrix_env_core/tests/test_manager_env.py | 63 ------------------- .../tests/test_manager_sim_backend.py | 4 -- .../ball_balance/mdp/observations.py | 5 -- .../locomotion/ball_balance/mdp/reset.py | 7 --- .../locomotion/ball_balance/mdp/rewards.py | 5 -- .../ball_balance/mdp/terminations.py | 4 -- .../humanoid/walk_manager_mdp/observations.py | 2 - .../humanoid/walk_manager_mdp/reset.py | 3 - .../humanoid/walk_manager_mdp/rewards.py | 8 --- .../locomotion/wbt/mdp/observations.py | 8 --- .../motrix_envs/locomotion/wbt/mdp/reset.py | 6 -- .../motrix_envs/locomotion/wbt/mdp/rewards.py | 10 --- .../locomotion/wbt/mdp/terminations.py | 6 -- wiki/design/manager/kernel-abi.md | 2 +- wiki/design/manager/torch-executor.md | 2 +- 21 files changed, 12 insertions(+), 166 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index 30057716..b7d5f45b 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -78,10 +78,12 @@ README 或教程;插件与任务 package 的 README 不重复这些内容, ### Manager term 入口方法 -Manager 数值 term 入口方法必须用 `@dispatch` 装饰。`@dispatch` 是透明的标记装饰器:resolver 只做签名/标注 -校验,数值执行由 Torch 执行器按整批张量调度,不存在 Numba 编译路径。Command hooks(`update`/`advance`/ -`reset_env`)是类上的 `@staticmethod` + `@dispatch` 函数,签名为 `(receiver, ctx)`。哪些方法是 term 入口、 -哪些是普通 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`。 ### 环境配置与注册规范 diff --git a/motrix_env_core/src/motrix_env_core/manager/resolver.py b/motrix_env_core/src/motrix_env_core/manager/resolver.py index ece04dd4..ccbbba23 100644 --- a/motrix_env_core/src/motrix_env_core/manager/resolver.py +++ b/motrix_env_core/src/motrix_env_core/manager/resolver.py @@ -26,7 +26,6 @@ import torch from motrix_env_core.manager.context import ManagerContext -from motrix_env_core.manager.dispatch import is_dispatched from motrix_env_core.manager.env import ( KernelInputSource, ManagerEnv, @@ -388,7 +387,7 @@ def _resolve_reset( f"Manager simulator reset term {name!r} dispatch must annotate its second parameter as " f"Map[torch.Tensor] (conventionally named sim_writes)." ) - dispatcher = self._term_entry(function) + dispatcher = function expressions = [] prepared_indices = [] for index, value in enumerate(term.args): @@ -602,7 +601,7 @@ def _resolve_observation_dispatch( expressions.append(expression) warmup_values.append(warmup) expressions = tuple(expressions) - dispatcher = self._term_entry(function) + dispatcher = function return PreparedInvocation( f"{group}.{name}", "observation", @@ -703,7 +702,7 @@ def _resolve_dispatch_term( prepared_indices.append(prepared_index) expressions.append(expression) warmup_values.append(warmup) - dispatcher = self._term_entry(function) + dispatcher = function return PreparedInvocation( f"{group}.{name}", group, @@ -762,7 +761,7 @@ def _resolve_receiver_term( ) kind = f"command_{function_name}" if group == "command" else group context = f"{group}.{name} ({function.__module__}.{function.__qualname__})" - dispatcher = self._term_entry(function) + dispatcher = function return PreparedInvocation(context, kind, dispatcher, receiver_key, output_size=output_size) # -- input preparation ---------------------------------------------------- @@ -863,9 +862,3 @@ def _validate_observation_spaces(self, groups: dict[str, ObservationGroupEntry]) raise TypeError("Manager value observation space must use float32 dtype.") elif self._env.has_value_observation: raise ValueError("Environment declares a value observation space but manager config has no 'value' group.") - - def _term_entry(self, function: Any) -> Any: - """Validate one dispatch entry and keep it as the Python callable.""" - if not is_dispatched(function): - raise TypeError(f"Manager kernel entry {function.__qualname__!r} must be decorated with @dispatch.") - return function diff --git a/motrix_env_core/src/motrix_env_core/manager/terms.py b/motrix_env_core/src/motrix_env_core/manager/terms.py index a4fe864a..88110c43 100644 --- a/motrix_env_core/src/motrix_env_core/manager/terms.py +++ b/motrix_env_core/src/motrix_env_core/manager/terms.py @@ -63,8 +63,6 @@ def __init__(self, dispatch: Callable[..., Any], *args: Any) -> None: def __post_init__(self) -> None: if not inspect.isfunction(self.dispatch): raise TypeError(f"Term dispatch must be a Python function, got {type(self.dispatch).__name__}.") - if not getattr(self.dispatch, "__motrix_manager_dispatch__", False): - raise TypeError(f"Term dispatch {self.dispatch.__qualname__!r} must be decorated with @dispatch.") if not isinstance(self.args, tuple): raise TypeError("Term invocation args must be a tuple.") diff --git a/motrix_env_core/src/motrix_env_core/mdp/observations.py b/motrix_env_core/src/motrix_env_core/mdp/observations.py index c7b41f5a..b9e43787 100644 --- a/motrix_env_core/src/motrix_env_core/mdp/observations.py +++ b/motrix_env_core/src/motrix_env_core/mdp/observations.py @@ -13,7 +13,6 @@ from motrix_env_core.config import configclass from motrix_env_core.config.scene import RobotCfg from motrix_env_core.manager.context import BuildContext, ManagerContext -from motrix_env_core.manager.dispatch import dispatch from motrix_env_core.manager.observations import ObservationTermCfg, ObsTerm from motrix_env_core.sim import ( BodyJointPositionQuery, @@ -61,7 +60,6 @@ def _rotate_inverse(quaternion: torch.Tensor, vector: torch.Tensor, out: torch.T out[:, 2] = vz + 2.0 * (x * cy - y * cx) -@dispatch def body_joint_vel_obs( ctx: ManagerContext, out: torch.Tensor, dof_vel: torch.Tensor, scale: float, noise_amplitude: float ) -> None: @@ -70,7 +68,6 @@ def body_joint_vel_obs( add_uniform_noise(ctx, out, noise_amplitude) -@dispatch def actions_obs(ctx: ManagerContext, out: torch.Tensor, action_name: str) -> None: out.copy_(ctx.actions[action_name].current()) @@ -86,7 +83,6 @@ def __call__(self, ctx: BuildContext) -> ObsTerm: return ObsTerm(action.action_queue.shape[2], actions_obs, self.action_name) -@dispatch def command_obs(ctx: ManagerContext, out: torch.Tensor, command_name: str) -> None: out.copy_(ctx.commands[command_name].command) @@ -127,7 +123,6 @@ def __call__(self, ctx: BuildContext) -> ObsTerm: ) -@dispatch def body_linear_velocity_obs( ctx: ManagerContext, out: torch.Tensor, @@ -161,7 +156,6 @@ def __call__(self, ctx: BuildContext) -> ObsTerm: ) -@dispatch def body_angular_velocity_obs( ctx: ManagerContext, out: torch.Tensor, @@ -195,7 +189,6 @@ def __call__(self, ctx: BuildContext) -> ObsTerm: ) -@dispatch def body_projected_gravity_obs( ctx: ManagerContext, out: torch.Tensor, base_quat: torch.Tensor, noise_amplitude: float ) -> None: @@ -222,7 +215,6 @@ def __call__(self, ctx: BuildContext) -> ObsTerm: ) -@dispatch def body_joint_pos_rel_obs( ctx: ManagerContext, out: torch.Tensor, diff --git a/motrix_env_core/src/motrix_env_core/mdp/rewards.py b/motrix_env_core/src/motrix_env_core/mdp/rewards.py index b633a062..631401fe 100644 --- a/motrix_env_core/src/motrix_env_core/mdp/rewards.py +++ b/motrix_env_core/src/motrix_env_core/mdp/rewards.py @@ -9,7 +9,6 @@ from motrix_env_core.config import configclass from motrix_env_core.manager import RewardTerm, RewardTermCfg from motrix_env_core.manager.context import BuildContext, ManagerContext -from motrix_env_core.manager.dispatch import dispatch from motrix_env_core.math.quaternion_array import rotate_inverse_components from motrix_env_core.sim import ( LinkAngularVelocityQuery, @@ -18,7 +17,6 @@ ) -@dispatch def alive_reward(ctx: ManagerContext) -> torch.Tensor: return torch.ones_like(ctx.env_id, dtype=torch.float32) @@ -31,7 +29,6 @@ def __call__(self, ctx) -> RewardTerm: return RewardTerm(alive_reward) -@dispatch def action_rate_reward(ctx: ManagerContext, action_name: str) -> torch.Tensor: action = ctx.actions[action_name] delta = action.current() - action.previous() @@ -48,7 +45,6 @@ def __call__(self, ctx) -> RewardTerm: return RewardTerm(action_rate_reward, self.action_name) -@dispatch def tracking_lin_vel_xy_reward( ctx: ManagerContext, sigma: float, @@ -81,7 +77,6 @@ def __call__(self, ctx: BuildContext) -> RewardTerm: ) -@dispatch def tracking_ang_vel_z_reward( ctx: ManagerContext, sigma: float, diff --git a/motrix_env_core/src/motrix_env_core/mdp/terminations.py b/motrix_env_core/src/motrix_env_core/mdp/terminations.py index 365ddd70..519da342 100644 --- a/motrix_env_core/src/motrix_env_core/mdp/terminations.py +++ b/motrix_env_core/src/motrix_env_core/mdp/terminations.py @@ -9,16 +9,13 @@ from motrix_env_core.config import configclass from motrix_env_core.manager import ManagerContext, TerminationTerm, TerminationTermCfg from motrix_env_core.manager.context import BuildContext -from motrix_env_core.manager.dispatch import dispatch from motrix_env_core.sim import BodyJointVelocityQuery, GeomPairCollidingQuery -@dispatch def colliding_termination(ctx: ManagerContext, colliding: torch.Tensor) -> torch.Tensor: return colliding.any(dim=1) -@dispatch def bad_dof_velocity_termination(ctx: ManagerContext, dof_vel: torch.Tensor, threshold: float) -> torch.Tensor: # Non-finite components count as +inf, so the max over the masked # absolute velocities terminates lanes with non-finite joint speeds diff --git a/motrix_env_core/tests/test_manager_env.py b/motrix_env_core/tests/test_manager_env.py index 4f46aab0..ea20051b 100644 --- a/motrix_env_core/tests/test_manager_env.py +++ b/motrix_env_core/tests/test_manager_env.py @@ -47,8 +47,6 @@ ) from motrix_env_core.manager.dispatch import dispatch # noqa: E402 from motrix_env_core.manager.observations import create_observation_groups -from motrix_env_core.manager.rewards import create_reward_terms -from motrix_env_core.manager.terminations import TerminationManager @kernel_data @@ -104,7 +102,6 @@ class _ObservationParams: scale: np.float32 -@dispatch def _injected_observation(ctx: ManagerContext, out: torch.Tensor, params: _ObservationParams) -> None: action = ctx.actions["test"] out[:, 0] = action.source[:, 0] * params.scale + params.offset[0] @@ -124,7 +121,6 @@ def __call__(self, env: ManagerEnv) -> ObsTerm: ) -@dispatch def _derived_observation(ctx: ManagerContext, out: torch.Tensor) -> None: out[:, 0] = ctx.commands["counter"].double[:, 0] @@ -140,7 +136,6 @@ def __call__(self, env: ManagerEnv) -> ObsTerm: return ObsTerm(1, _derived_observation) -@dispatch def _lane_observation(ctx: ManagerContext, out: torch.Tensor) -> None: out[:, 0] = ctx.env_id @@ -166,7 +161,6 @@ class _LaneObservationsCfg(ManagerObservationsCfg): policy: _LanePolicyObsCfg = _LanePolicyObsCfg() -@dispatch def _injected_reward(ctx: ManagerContext): return ctx.actions["test"].source[:, 0] @@ -178,7 +172,6 @@ def __call__(self, env: ManagerEnv) -> RewardTerm: return RewardTerm(_injected_reward) -@dispatch def _injected_termination(ctx: ManagerContext, threshold) -> bool: value = ctx.actions["test"].source[:, 0] ctx.metrics["source_at_termination"][:, 0] = value @@ -798,34 +791,6 @@ class _StaticObservationsCfg(ManagerObservationsCfg): assert static_groups["policy"].terms[0].term.dispatch is _static_observation -def test_reward_term_requires_dispatch_descriptor() -> None: - class _InvalidReward: - pass - - @configclass(kw_only=True) - class _InvalidRewardCfg(RewardTermCfg): - def __call__(self, env: ManagerEnv) -> RewardTerm: - del env - return _InvalidReward() - - env = _ManagerEnv(num_envs=1) - with pytest.raises(TypeError, match="must return RewardTerm"): - create_reward_terms({"invalid": _InvalidRewardCfg(weight=1.0)}, env) - - def _undecorated_reward(ctx: ManagerContext) -> float: - del ctx - return 1.0 - - @configclass(kw_only=True) - class _UndecoratedRewardCfg(RewardTermCfg): - def __call__(self, env: ManagerEnv) -> RewardTerm: - del env - return RewardTerm(_undecorated_reward) - - with pytest.raises(TypeError, match="must be decorated with @dispatch"): - create_reward_terms({"undecorated": _UndecoratedRewardCfg(weight=1.0)}, env) - - def test_dispatch_term_parameter_names_are_not_constrained() -> None: @dispatch def _renamed_reward(context: ManagerContext): @@ -848,31 +813,3 @@ class _RenamedRewardsCfg(ManagerRewardsCfg): env.compute_transition(state) np.testing.assert_allclose(state.reward, 2.0 * 0.5 * env.cfg.ctrl_dt) - - -def test_termination_term_requires_dispatch_descriptor() -> None: - class _InvalidTermination: - pass - - @configclass(kw_only=True) - class _InvalidTerminationCfg(TerminationTermCfg): - def __call__(self, env: ManagerEnv) -> TerminationTerm: - del env - return _InvalidTermination() - - env = _ManagerEnv(num_envs=1) - with pytest.raises(TypeError, match="must return TerminationTerm"): - TerminationManager({"invalid": _InvalidTerminationCfg()}, env) - - def _undecorated_termination(ctx: ManagerContext) -> bool: - del ctx - return False - - @configclass(kw_only=True) - class _UndecoratedTerminationCfg(TerminationTermCfg): - def __call__(self, env: ManagerEnv) -> TerminationTerm: - del env - return TerminationTerm(_undecorated_termination) - - with pytest.raises(TypeError, match="must be decorated with @dispatch"): - TerminationManager({"undecorated": _UndecoratedTerminationCfg()}, env) diff --git a/motrix_env_core/tests/test_manager_sim_backend.py b/motrix_env_core/tests/test_manager_sim_backend.py index 05152dd7..1980b5e7 100644 --- a/motrix_env_core/tests/test_manager_sim_backend.py +++ b/motrix_env_core/tests/test_manager_sim_backend.py @@ -15,7 +15,6 @@ from motrix_env_core.config.scene import SceneCfg from motrix_env_core.manager.actions import ActionCfg, ActionState, ActionTerm, ManagerActionsCfg from motrix_env_core.manager.context import ManagerContext -from motrix_env_core.manager.dispatch import dispatch from motrix_env_core.manager.env import ManagerBasedEnvCfg, ManagerEnv from motrix_env_core.manager.kernel_data import kernel_data from motrix_env_core.manager.observations import ( @@ -237,7 +236,6 @@ def __call__(self, env, actuator_indices): ) -@dispatch def _sim_reading_observation(ctx: ManagerContext, out: torch.Tensor) -> None: out.copy_(ctx.sim["dof_pos"]) @@ -258,7 +256,6 @@ class _FakeObservationsCfg(ManagerObservationsCfg): policy: _FakePolicyObsCfg = _FakePolicyObsCfg() -@dispatch def _unit_reward(ctx: ManagerContext) -> torch.Tensor: return torch.ones_like(ctx.env_id, dtype=torch.float32) @@ -276,7 +273,6 @@ class _FakeRewardsCfg(ManagerRewardsCfg): unit: _UnitRewardCfg = _UnitRewardCfg() -@dispatch def _never_termination(ctx: ManagerContext) -> torch.Tensor: return torch.zeros_like(ctx.env_id, dtype=torch.bool) diff --git a/motrix_envs/src/motrix_envs/locomotion/ball_balance/mdp/observations.py b/motrix_envs/src/motrix_envs/locomotion/ball_balance/mdp/observations.py index 0f83e90d..07407af8 100644 --- a/motrix_envs/src/motrix_envs/locomotion/ball_balance/mdp/observations.py +++ b/motrix_envs/src/motrix_envs/locomotion/ball_balance/mdp/observations.py @@ -12,12 +12,10 @@ ObservationTermCfg, ObsTerm, ) -from motrix_env_core.manager.dispatch import dispatch from motrix_env_core.math.quaternion_array import rotate_inverse as rotate_inverse_batch from motrix_env_core.mdp.observations import UniformNoiseCfg -@dispatch def ball_relative_position_obs(ctx: ManagerContext, out: torch.Tensor, noise_amplitude: float) -> None: relative = ctx.sim["ball_pos"] - ctx.sim["robot_base_pos"] out.copy_(rotate_inverse_batch(ctx.sim["robot_base_quat"], relative)) @@ -35,7 +33,6 @@ def __call__(self, ctx) -> ObsTerm: return ObsTerm(3, ball_relative_position_obs, np.float32(self.noise.amplitude)) -@dispatch def ball_relative_velocity_obs(ctx: ManagerContext, out: torch.Tensor, noise_amplitude: float) -> None: out.copy_(rotate_inverse_batch(ctx.sim["robot_base_quat"], ctx.sim["ball_lin_vel"])) if noise_amplitude != 0.0: @@ -52,7 +49,6 @@ def __call__(self, ctx) -> ObsTerm: return ObsTerm(3, ball_relative_velocity_obs, np.float32(self.noise.amplitude)) -@dispatch def ball_position_obs(ctx: ManagerContext, out: torch.Tensor) -> None: out.copy_(ctx.sim["ball_pos"]) @@ -65,7 +61,6 @@ def __call__(self, ctx) -> ObsTerm: return ObsTerm(3, ball_position_obs) -@dispatch def ball_velocity_obs(ctx: ManagerContext, out: torch.Tensor) -> None: out.copy_(ctx.sim["ball_lin_vel"]) diff --git a/motrix_envs/src/motrix_envs/locomotion/ball_balance/mdp/reset.py b/motrix_envs/src/motrix_envs/locomotion/ball_balance/mdp/reset.py index 915a15ed..df29480d 100644 --- a/motrix_envs/src/motrix_envs/locomotion/ball_balance/mdp/reset.py +++ b/motrix_envs/src/motrix_envs/locomotion/ball_balance/mdp/reset.py @@ -14,7 +14,6 @@ ResetTerm, ResetTermCfg, ) -from motrix_env_core.manager.dispatch import dispatch from motrix_env_core.manager.kernel_data import Map, SharedArray, kernel_data from motrix_env_core.math.quaternion_array import from_euler from motrix_env_core.sim import ( @@ -27,7 +26,6 @@ ) -@dispatch def _reset_body_pos( ctx: ManagerContext, sim_writes: Map[torch.Tensor], spawn: torch.Tensor, noise_amplitude: torch.Tensor ) -> None: @@ -53,7 +51,6 @@ def __call__(self, ctx) -> ResetTerm: ) -@dispatch def _reset_body_rot(ctx: ManagerContext, sim_writes: Map[torch.Tensor], noise_amplitude: torch.Tensor) -> None: # Fixed promotion order: f32 angles, f64 trig, f32 quaternion, # sequential f32 sum of squares, sqrt via f64 rounded back to f32, and an @@ -84,7 +81,6 @@ def __call__(self, ctx) -> ResetTerm: ) -@dispatch def _reset_body_lin_vel(ctx: ManagerContext, sim_writes: Map[torch.Tensor], noise_amplitude: torch.Tensor) -> None: noise = ctx.rand.next_uniform(3) sim_writes["linear_velocity"][:, 0].copy_(noise * noise_amplitude) @@ -107,7 +103,6 @@ def __call__(self, ctx) -> ResetTerm: ) -@dispatch def _reset_body_rot_vel(ctx: ManagerContext, sim_writes: Map[torch.Tensor], noise_amplitude: torch.Tensor) -> None: noise = ctx.rand.next_uniform(3) sim_writes["angular_velocity"][:, 0].copy_(noise * noise_amplitude) @@ -138,7 +133,6 @@ class DofResetParams: noise_scale: np.float32 -@dispatch def _reset_body_dof_pos(ctx: ManagerContext, sim_writes: Map[torch.Tensor], params: DofResetParams) -> None: noise = ctx.rand.next_uniform(params.default_pos.shape[0]) value = params.default_pos + noise * params.noise_scale @@ -183,7 +177,6 @@ def __call__(self, ctx) -> ResetTerm: ) -@dispatch def _reset_ball( ctx: ManagerContext, sim_writes: Map[torch.Tensor], diff --git a/motrix_envs/src/motrix_envs/locomotion/ball_balance/mdp/rewards.py b/motrix_envs/src/motrix_envs/locomotion/ball_balance/mdp/rewards.py index 4696c7db..97737ef2 100644 --- a/motrix_envs/src/motrix_envs/locomotion/ball_balance/mdp/rewards.py +++ b/motrix_envs/src/motrix_envs/locomotion/ball_balance/mdp/rewards.py @@ -8,12 +8,10 @@ from motrix_env_core.config import configclass from motrix_env_core.manager import ManagerContext, RewardTerm, RewardTermCfg -from motrix_env_core.manager.dispatch import dispatch from motrix_env_core.manager.kernel_data import SharedArray, kernel_data from motrix_env_core.math.quaternion_array import rotate_inverse as rotate_inverse_batch -@dispatch def upright_reward(ctx: ManagerContext, sigma: float) -> torch.Tensor: base_quat = ctx.sim["robot_base_quat"] gravity = base_quat.new_zeros(base_quat.shape[0], 3) @@ -35,7 +33,6 @@ def __call__(self, ctx) -> RewardTerm: return RewardTerm(upright_reward, np.float32(self.sigma)) -@dispatch def base_height_reward(ctx: ManagerContext, target_z: float, sigma: float) -> torch.Tensor: error = ctx.sim["robot_base_pos"][:, 2] - target_z return torch.exp(-(error * error) / (sigma * sigma)) @@ -52,7 +49,6 @@ def __call__(self, ctx) -> RewardTerm: return RewardTerm(base_height_reward, np.float32(self.target_z), np.float32(self.sigma)) -@dispatch def ball_under_feet_reward(ctx: ManagerContext, sigma: float) -> torch.Tensor: feet_pos = ctx.sim["feet_pos"] ball_pos = ctx.sim["ball_pos"] @@ -80,7 +76,6 @@ class DofDefaultParams: sigma: np.float32 -@dispatch def dof_default_reward(ctx: ManagerContext, params: DofDefaultParams) -> torch.Tensor: error = ctx.sim["robot_dof_pos"] - params.reference error_sq = (error * error).sum(dim=1) / error.shape[1] diff --git a/motrix_envs/src/motrix_envs/locomotion/ball_balance/mdp/terminations.py b/motrix_envs/src/motrix_envs/locomotion/ball_balance/mdp/terminations.py index 793f506e..fb2f4bee 100644 --- a/motrix_envs/src/motrix_envs/locomotion/ball_balance/mdp/terminations.py +++ b/motrix_envs/src/motrix_envs/locomotion/ball_balance/mdp/terminations.py @@ -8,7 +8,6 @@ from motrix_env_core.config import configclass from motrix_env_core.manager import ManagerContext, TerminationTerm, TerminationTermCfg -from motrix_env_core.manager.dispatch import dispatch from motrix_env_core.math.quaternion_array import rotate_inverse as rotate_inverse_batch @@ -17,7 +16,6 @@ class _BallBalanceTerminationCfg(TerminationTermCfg): threshold: float -@dispatch def bad_base_z_termination(ctx: ManagerContext, threshold: float) -> torch.Tensor: error = ctx.sim["robot_base_pos"][:, 2] ctx.metrics["base_z"][:, 0] = error @@ -36,7 +34,6 @@ def __call__(self, ctx) -> TerminationTerm: ) -@dispatch def bad_orientation_termination(ctx: ManagerContext, threshold: float) -> torch.Tensor: base_quat = ctx.sim["robot_base_quat"] gravity = base_quat.new_zeros(base_quat.shape[0], 3) @@ -59,7 +56,6 @@ def __call__(self, ctx) -> TerminationTerm: ) -@dispatch def ball_escaped_termination(ctx: ManagerContext, threshold: float) -> torch.Tensor: feet_pos = ctx.sim["feet_pos"] ball_pos = ctx.sim["ball_pos"] diff --git a/motrix_envs/src/motrix_envs/locomotion/humanoid/walk_manager_mdp/observations.py b/motrix_envs/src/motrix_envs/locomotion/humanoid/walk_manager_mdp/observations.py index db958408..265b7712 100644 --- a/motrix_envs/src/motrix_envs/locomotion/humanoid/walk_manager_mdp/observations.py +++ b/motrix_envs/src/motrix_envs/locomotion/humanoid/walk_manager_mdp/observations.py @@ -8,11 +8,9 @@ from motrix_env_core.config import configclass from motrix_env_core.manager.context import BuildContext, ManagerContext -from motrix_env_core.manager.dispatch import dispatch from motrix_env_core.manager.observations import ObservationTermCfg, ObsTerm -@dispatch def gait_phase_obs(ctx: ManagerContext, out: torch.Tensor, offset: np.int64, size: np.int64) -> None: walk = ctx.commands["walk"] out.copy_(walk.sin_cos[:, offset : offset + size]) diff --git a/motrix_envs/src/motrix_envs/locomotion/humanoid/walk_manager_mdp/reset.py b/motrix_envs/src/motrix_envs/locomotion/humanoid/walk_manager_mdp/reset.py index 9cdc24ee..a4dd0de8 100644 --- a/motrix_envs/src/motrix_envs/locomotion/humanoid/walk_manager_mdp/reset.py +++ b/motrix_envs/src/motrix_envs/locomotion/humanoid/walk_manager_mdp/reset.py @@ -15,7 +15,6 @@ kernel_data, ) from motrix_env_core.manager.context import BuildContext -from motrix_env_core.manager.dispatch import dispatch from motrix_env_core.manager.kernel_data.map import Map from motrix_env_core.mdp.terrain import HeightFieldGrid, heightfield_lookup from motrix_env_core.sim.write import ( @@ -158,7 +157,6 @@ def _write_spawn_state( sim_writes["joints_velocity"][:] = 0.0 -@dispatch def reset_walk_state(ctx: ManagerContext, sim_writes: Map[torch.Tensor], params: WalkResetParams) -> None: """Spawn-state-only reset for configs without dynamics randomization.""" pose = params.init_pose @@ -168,7 +166,6 @@ def reset_walk_state(ctx: ManagerContext, sim_writes: Map[torch.Tensor], params: _write_spawn_state(sim_writes, params, x, y, z) -@dispatch def reset_walk_state_randomized(ctx: ManagerContext, sim_writes: Map[torch.Tensor], params: WalkResetParams) -> None: """Reset plus in-kernel dynamics randomization (kp/damping/friction/mass/com).""" pose = params.init_pose diff --git a/motrix_envs/src/motrix_envs/locomotion/humanoid/walk_manager_mdp/rewards.py b/motrix_envs/src/motrix_envs/locomotion/humanoid/walk_manager_mdp/rewards.py index 628e58fd..f6c9fd67 100644 --- a/motrix_envs/src/motrix_envs/locomotion/humanoid/walk_manager_mdp/rewards.py +++ b/motrix_envs/src/motrix_envs/locomotion/humanoid/walk_manager_mdp/rewards.py @@ -15,7 +15,6 @@ RewardTermCfg, ) from motrix_env_core.manager.context import BuildContext -from motrix_env_core.manager.dispatch import dispatch from motrix_env_core.math.quaternion_array import rotate_inverse_components from motrix_env_core.mdp.terrain import HeightFieldGrid, heightfield_lookup from motrix_env_core.sim import ( @@ -39,7 +38,6 @@ def _expected_foot_height(phi: torch.Tensor, swing_height: float) -> torch.Tenso return y_start + (y_end - y_start) * (t**3 + 3.0 * (t**2 * (1.0 - t))) -@dispatch def penalty_ang_vel_xy_reward(ctx: ManagerContext, base_quat: torch.Tensor, base_ang_vel: torch.Tensor) -> torch.Tensor: vx, vy, _ = rotate_inverse_components(base_quat, base_ang_vel) walk = ctx.commands["walk"] @@ -57,7 +55,6 @@ def __call__(self, ctx) -> RewardTerm: ) -@dispatch def penalty_orientation_reward(ctx: ManagerContext, base_quat: torch.Tensor) -> torch.Tensor: gravity = base_quat.new_zeros(base_quat.shape[0], 3) gravity[:, 2] = -1.0 @@ -76,7 +73,6 @@ def __call__(self, ctx) -> RewardTerm: ) -@dispatch def penalty_action_rate_reward(ctx: ManagerContext) -> torch.Tensor: action = ctx.actions["joint_position"] delta = action.current() - action.previous() @@ -90,7 +86,6 @@ def __call__(self, ctx) -> RewardTerm: return RewardTerm(penalty_action_rate_reward) -@dispatch def feet_phase_reward( ctx: ManagerContext, swing_height: float, @@ -138,7 +133,6 @@ def __call__(self, ctx) -> RewardTerm: ) -@dispatch def pose_reward( ctx: ManagerContext, default_joint_angles: torch.Tensor, pose_weights: torch.Tensor, joint_pos: torch.Tensor ) -> torch.Tensor: @@ -175,7 +169,6 @@ def __call__(self, ctx) -> RewardTerm: ) -@dispatch def penalty_close_feet_xy_reward( ctx: ManagerContext, close_feet_threshold: float, base_quat: torch.Tensor, foot_pos: torch.Tensor ) -> torch.Tensor: @@ -208,7 +201,6 @@ def __call__(self, ctx) -> RewardTerm: ) -@dispatch def penalty_feet_ori_reward( ctx: ManagerContext, default_foot_gravity: torch.Tensor, foot_quat: torch.Tensor ) -> torch.Tensor: diff --git a/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/observations.py b/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/observations.py index 1f6addef..3b5715e4 100644 --- a/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/observations.py +++ b/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/observations.py @@ -14,7 +14,6 @@ ObsTerm, kernel_data, ) -from motrix_env_core.manager.dispatch import dispatch from motrix_env_core.manager.kernel_data import SharedArray from motrix_env_core.math import quaternion_array as batch_math from motrix_env_core.math.quaternion_array import relative_orientation, rotate_inverse @@ -25,7 +24,6 @@ from motrix_envs.locomotion.wbt.mdp.command import WbtMotionCommand -@dispatch def motion_joint_obs(ctx: ManagerContext, out: torch.Tensor) -> None: motion = ctx.commands["motion"] steps = motion.steps[:, 0] @@ -43,7 +41,6 @@ def __call__(self, ctx) -> ObsTerm: return ObsTerm(2 * motion.clip.joint_pos.shape[1], motion_joint_obs) -@dispatch def motion_reference_position_obs(ctx: ManagerContext, out: torch.Tensor) -> None: motion = ctx.commands["motion"] reference_index = motion.reference_index @@ -62,7 +59,6 @@ def __call__(self, ctx) -> ObsTerm: return ObsTerm(3, motion_reference_position_obs) -@dispatch def motion_reference_orientation_obs( ctx: ManagerContext, out: torch.Tensor, @@ -83,7 +79,6 @@ def __call__(self, ctx) -> ObsTerm: return ObsTerm(6, motion_reference_orientation_obs, np.float32(self.noise.amplitude)) -@dispatch def robot_body_position_in_reference_frame_obs(ctx: ManagerContext, out: torch.Tensor) -> None: motion = ctx.commands["motion"] reference_index = motion.reference_index @@ -102,7 +97,6 @@ def __call__(self, ctx) -> ObsTerm: return ObsTerm(3 * num_tracked, robot_body_position_in_reference_frame_obs) -@dispatch def robot_body_orientation_obs(ctx: ManagerContext, out: torch.Tensor) -> None: # Bug-compatible with the historical per-lane oracle: body zero writes its # 6D matrix over the shared out[:4] inverse-quaternion scratch, and every @@ -129,7 +123,6 @@ class RelativePositionParams: noise_amplitude: np.float32 -@dispatch def dof_pos_rel_obs(ctx: ManagerContext, out: torch.Tensor, params: RelativePositionParams) -> None: out.copy_(ctx.sim["robot_dof_pos"] - params.reference) add_uniform_noise(ctx, out, params.noise_amplitude) @@ -172,7 +165,6 @@ def __call__(self, ctx) -> ObsTerm: return ObsTerm(reference.shape[0], dof_pos_rel_obs, params) -@dispatch def dof_vel_obs(ctx: ManagerContext, out: torch.Tensor, noise_amplitude: float) -> None: out.copy_(ctx.sim["robot_dof_vel"]) add_uniform_noise(ctx, out, noise_amplitude) diff --git a/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/reset.py b/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/reset.py index 50b2dd94..6926dfd3 100644 --- a/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/reset.py +++ b/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/reset.py @@ -14,7 +14,6 @@ ResetTerm, ResetTermCfg, ) -from motrix_env_core.manager.dispatch import dispatch from motrix_env_core.manager.kernel_data import Map from motrix_env_core.math.quaternion_array import mul from motrix_env_core.sim import ( @@ -27,7 +26,6 @@ ) -@dispatch def _reset_body_pos( ctx: ManagerContext, sim_writes: Map[torch.Tensor], @@ -58,7 +56,6 @@ def __call__(self, ctx) -> ResetTerm: ) -@dispatch def _reset_body_rot( ctx: ManagerContext, sim_writes: Map[torch.Tensor], @@ -112,7 +109,6 @@ def __call__(self, ctx) -> ResetTerm: ) -@dispatch def _reset_body_lin_vel( ctx: ManagerContext, sim_writes: Map[torch.Tensor], @@ -141,7 +137,6 @@ def __call__(self, ctx) -> ResetTerm: ) -@dispatch def _reset_body_rot_vel( ctx: ManagerContext, sim_writes: Map[torch.Tensor], @@ -170,7 +165,6 @@ def __call__(self, ctx) -> ResetTerm: ) -@dispatch def _reset_body_dof_pos(ctx: ManagerContext, sim_writes: Map[torch.Tensor], noise_scale: float) -> None: motion = ctx.commands["motion"] steps = motion.steps[:, 0] diff --git a/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/rewards.py b/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/rewards.py index 17047fd5..ff40192b 100644 --- a/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/rewards.py +++ b/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/rewards.py @@ -10,13 +10,11 @@ from motrix_env_core.config import configclass from motrix_env_core.manager import ManagerContext, RewardTerm, RewardTermCfg -from motrix_env_core.manager.dispatch import dispatch from motrix_env_core.manager.kernel_data import SharedArray, kernel_data from motrix_env_core.math.quaternion_array import rotation_distance from motrix_env_core.mdp.action import JointPositionActionState -@dispatch def global_ref_position_reward(ctx: ManagerContext, sigma: float) -> torch.Tensor: motion = ctx.commands["motion"] steps = motion.steps[:, 0] @@ -32,7 +30,6 @@ def __call__(self, ctx) -> RewardTerm: return RewardTerm(global_ref_position_reward, np.float32(self.sigma)) -@dispatch def global_ref_orientation_reward(ctx: ManagerContext, sigma: float) -> torch.Tensor: motion = ctx.commands["motion"] steps = motion.steps[:, 0] @@ -50,7 +47,6 @@ def __call__(self, ctx) -> RewardTerm: return RewardTerm(global_ref_orientation_reward, np.float32(self.sigma)) -@dispatch def relative_body_position_reward(ctx: ManagerContext, sigma: float) -> torch.Tensor: motion = ctx.commands["motion"] diff = motion.target_body_position_relative - ctx.sim["tracked_body_pos"] @@ -66,7 +62,6 @@ def __call__(self, ctx) -> RewardTerm: return RewardTerm(relative_body_position_reward, np.float32(self.sigma)) -@dispatch def ee_body_pos_reward(ctx: ManagerContext, body_indices: tuple[int, ...], sigma: float) -> torch.Tensor: motion = ctx.commands["motion"] position = ctx.sim["tracked_body_pos"] @@ -101,7 +96,6 @@ def __call__(self, ctx) -> RewardTerm: return RewardTerm(ee_body_pos_reward, body_indices, np.float32(self.sigma)) -@dispatch def relative_body_orientation_reward(ctx: ManagerContext, sigma: float) -> torch.Tensor: motion = ctx.commands["motion"] quaternion = ctx.sim["tracked_body_quat"] @@ -118,7 +112,6 @@ def __call__(self, ctx) -> RewardTerm: return RewardTerm(relative_body_orientation_reward, np.float32(self.sigma)) -@dispatch def global_body_linear_velocity_reward(ctx: ManagerContext, sigma: float) -> torch.Tensor: motion = ctx.commands["motion"] steps = motion.steps[:, 0] @@ -136,7 +129,6 @@ def __call__(self, ctx) -> RewardTerm: return RewardTerm(global_body_linear_velocity_reward, np.float32(self.sigma)) -@dispatch def global_body_angular_velocity_reward(ctx: ManagerContext, sigma: float) -> torch.Tensor: motion = ctx.commands["motion"] steps = motion.steps[:, 0] @@ -162,7 +154,6 @@ class DofLimitParams: cap: np.float32 -@dispatch def dof_limit_reward(ctx: ManagerContext, params: DofLimitParams) -> torch.Tensor: dof_position = ctx.sim["robot_dof_pos"] violation = torch.abs(dof_position - params.midpoint) - params.half_range * params.soft_limit @@ -186,7 +177,6 @@ def __call__(self, ctx) -> RewardTerm: return RewardTerm(dof_limit_reward, params) -@dispatch def undesired_contacts_reward(ctx: ManagerContext, threshold: float) -> torch.Tensor: contact_forces = ctx.sim["undesired_contact_forces"] forces = contact_forces[..., :3] diff --git a/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/terminations.py b/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/terminations.py index 03b7194f..601af67a 100644 --- a/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/terminations.py +++ b/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/terminations.py @@ -8,7 +8,6 @@ from motrix_env_core.config import configclass from motrix_env_core.manager import ManagerContext, TerminationTerm, TerminationTermCfg -from motrix_env_core.manager.dispatch import dispatch @configclass(kw_only=True) @@ -16,7 +15,6 @@ class _WbtTerminationCfg(TerminationTermCfg): threshold: float -@dispatch def bad_ref_z_termination(ctx: ManagerContext, threshold: float) -> torch.Tensor: motion = ctx.commands["motion"] steps = motion.steps[:, 0] @@ -37,7 +35,6 @@ def __call__(self, ctx) -> TerminationTerm: ) -@dispatch def bad_ref_orientation_termination(ctx: ManagerContext, threshold: float) -> torch.Tensor: motion = ctx.commands["motion"] steps = motion.steps[:, 0] @@ -58,7 +55,6 @@ def __call__(self, ctx) -> TerminationTerm: ) -@dispatch def bad_motion_body_position_termination( ctx: ManagerContext, body_indices: tuple[int, ...], @@ -98,7 +94,6 @@ def __call__(self, ctx) -> TerminationTerm: ) -@dispatch def bad_body_z_termination( ctx: ManagerContext, body_indices: tuple[int, ...], @@ -130,7 +125,6 @@ def __call__(self, ctx) -> TerminationTerm: ) -@dispatch def bad_dof_position_termination(ctx: ManagerContext, threshold: float) -> torch.Tensor: action = ctx.actions["joint_position"] position = ctx.sim["robot_dof_pos"] diff --git a/wiki/design/manager/kernel-abi.md b/wiki/design/manager/kernel-abi.md index cba36dca..fd741047 100644 --- a/wiki/design/manager/kernel-abi.md +++ b/wiki/design/manager/kernel-abi.md @@ -95,7 +95,7 @@ Term 可以读取其他 term 的数据,但不应调用其他 term 的行为方 - flatten concrete runtime terms 和 `ManagerContext`; - 收集 read program 的 sim input map; - 解析 context、term receiver 和 literal Map lookup 的 slot/表达式; -- 校验 `@dispatch` 标记、方法签名、返回注解和 layout; +- 校验方法签名、注解契约和 layout; - 产出 `ManagerLayout`、`ReadPlan`、`KernelInputSource` 与 `PreparedInvocation`,供 Torch 执行器调度。 构造期 warm-up 以克隆输入执行一次 evaluate/observe,确认输入/输出契约(shape、dtype、batch 维、可写性); diff --git a/wiki/design/manager/torch-executor.md b/wiki/design/manager/torch-executor.md index 422f7f76..236e4c03 100644 --- a/wiki/design/manager/torch-executor.md +++ b/wiki/design/manager/torch-executor.md @@ -24,7 +24,7 @@ RNG draw 顺序与 eager 完全一致;Inductor 重结合仅带来 ≤1e-8 量 - `ManagerProgramResolver`(backend-neutral)负责输入准备:flatten `@kernel_data` term 与 `ManagerContext`,构建 `ManagerLayout`、 `KernelInputSource`/`ReadPlan`、`PreparedInvocation`(每个 term 的 - dispatcher、参数表达式/值/prepared index),并校验 `@dispatch` 标记与 + dispatcher、参数表达式/值/prepared index),并校验签名注解契约与 签名契约。它不做任何代码生成、缓存或 dispatcher 编译。 - `TorchManagerRuntime` 消费这些元数据并执行 term 调度;所有写入落回 host 持有的 numpy 缓冲(reward terms/masks/obs/command state/metrics), From 8942eb19a7602922a6456300c02b80ad0a42e96f Mon Sep 17 00:00:00 2001 From: zilch <147668916@qq.com> Date: Thu, 5 Nov 2026 23:37:51 +0800 Subject: [PATCH 06/41] refactor(mdp): consume the joint-default argument as the tensor view 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. --- motrix_env_core/src/motrix_env_core/mdp/observations.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/motrix_env_core/src/motrix_env_core/mdp/observations.py b/motrix_env_core/src/motrix_env_core/mdp/observations.py index b9e43787..00408c3c 100644 --- a/motrix_env_core/src/motrix_env_core/mdp/observations.py +++ b/motrix_env_core/src/motrix_env_core/mdp/observations.py @@ -223,8 +223,7 @@ def body_joint_pos_rel_obs( scale: float, noise_amplitude: float, ) -> None: - out.copy_(joint_pos - torch.as_tensor(np.asarray(defaults), dtype=out.dtype)) - out.mul_(scale) + out.copy_((joint_pos - defaults.to(out.dtype)) * scale) add_uniform_noise(ctx, out, noise_amplitude) From 93e5a025c516a597922fae042c997d90c9261951 Mon Sep 17 00:00:00 2001 From: zilch <147668916@qq.com> Date: Thu, 5 Nov 2026 23:40:46 +0800 Subject: [PATCH 07/41] refactor(terrain): tidy heightfield_lookup placement, typing, and view use --- .../src/motrix_env_core/mdp/terrain.py | 34 +++++++++---------- 1 file changed, 17 insertions(+), 17 deletions(-) diff --git a/motrix_env_core/src/motrix_env_core/mdp/terrain.py b/motrix_env_core/src/motrix_env_core/mdp/terrain.py index c395137c..0caff9c1 100644 --- a/motrix_env_core/src/motrix_env_core/mdp/terrain.py +++ b/motrix_env_core/src/motrix_env_core/mdp/terrain.py @@ -14,8 +14,7 @@ class HeightFieldGrid: """Static terrain grid for batched bilinear height lookup. ``enabled=False`` degenerates to the flat-ground ``constant`` height. - Torch-side lookups (for example - ``motrix_envs.locomotion.humanoid.walk_manager_mdp.terrain.heightfield_lookup_batch``) + Torch-side lookups (:func:`heightfield_lookup`) mirror the engine's hfield sample convention in float64: bilinear interpolation of the four surrounding samples after mapping world ``xy`` to fractional grid coordinates. @@ -43,24 +42,19 @@ class HeightFieldGrid: enabled: bool -__all__ = [ - "heightfield_lookup", - "HeightFieldGrid", -] - - -def heightfield_lookup(grid, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: - """Batched torch bilinear height lookup over a :class:`HeightFieldGrid`. +def heightfield_lookup(grid: HeightFieldGrid, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: + """Batched bilinear height lookup in float64 over a height-field grid. - ``x``/``y`` are ``(N,)`` world coordinates. The interpolation runs in - float64, matching the engine's hfield sample convention. + ``x``/``y`` are ``(N,)`` world coordinates. ``grid`` is the batch + context's kernel-data view of a :class:`HeightFieldGrid`; its fields + arrive as zero-copy tensor views of the shared grid arrays. """ - if not bool(grid.enabled): + if not grid.enabled: return torch.full_like(x.double(), float(grid.constant)) - heights = torch.as_tensor(np.asarray(grid.heights)).to(torch.float64) - origin = torch.as_tensor(np.asarray(grid.origin)).to(torch.float64) - spacing = torch.as_tensor(np.asarray(grid.spacing)).to(torch.float64) - z0 = torch.as_tensor(np.asarray(grid.z0)).to(torch.float64) + heights = grid.heights.to(torch.float64) + origin = grid.origin.to(torch.float64) + spacing = grid.spacing.to(torch.float64) + z0 = grid.z0.to(torch.float64) nrow, ncol = heights.shape fx = ((x.to(torch.float64) - origin[0]) / spacing[0]).clamp(0.0, ncol - 1.0 - 1e-6) fy = ((y.to(torch.float64) - origin[1]) / spacing[1]).clamp(0.0, nrow - 1.0 - 1e-6) @@ -75,3 +69,9 @@ def heightfield_lookup(grid, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: top = h00 * (1.0 - tx) + h01 * tx bottom = h10 * (1.0 - tx) + h11 * tx return z0[0] + top * (1.0 - ty) + bottom * ty + + +__all__ = [ + "heightfield_lookup", + "HeightFieldGrid", +] From 7e971515129c089d53caedcc679ce7f43d2e23de Mon Sep 17 00:00:00 2001 From: zilch <147668916@qq.com> Date: Thu, 5 Nov 2026 23:43:54 +0800 Subject: [PATCH 08/41] refactor(manager): remove the unused MdpContext/ContextMap protocols 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. --- .../src/motrix_env_core/manager/__init__.py | 5 ++--- .../src/motrix_env_core/manager/context.py | 19 ++----------------- 2 files changed, 4 insertions(+), 20 deletions(-) diff --git a/motrix_env_core/src/motrix_env_core/manager/__init__.py b/motrix_env_core/src/motrix_env_core/manager/__init__.py index 5f3c9752..a78f7061 100644 --- a/motrix_env_core/src/motrix_env_core/manager/__init__.py +++ b/motrix_env_core/src/motrix_env_core/manager/__init__.py @@ -11,7 +11,7 @@ ManagerActionsCfg, ) from motrix_env_core.manager.commands import CommandCfg, CommandTerm, ManagerCommandsCfg -from motrix_env_core.manager.context import ContextMap, ManagerContext, MdpContext +from motrix_env_core.manager.context import BuildContext, ManagerContext from motrix_env_core.manager.dispatch import dispatch from motrix_env_core.manager.env import ManagerBasedEnvCfg, ManagerEnv from motrix_env_core.manager.kernel import ManagerWarmupResult @@ -40,11 +40,10 @@ "ActionState", "CommandCfg", "CommandTerm", - "ContextMap", - "MdpContext", "ManagerActionsCfg", "ManagerBasedEnvCfg", "ManagerCommandsCfg", + "BuildContext", "ManagerContext", "dispatch", "ManagerEnv", diff --git a/motrix_env_core/src/motrix_env_core/manager/context.py b/motrix_env_core/src/motrix_env_core/manager/context.py index 4b19ceb4..75871f42 100644 --- a/motrix_env_core/src/motrix_env_core/manager/context.py +++ b/motrix_env_core/src/motrix_env_core/manager/context.py @@ -5,7 +5,7 @@ from __future__ import annotations -from typing import TYPE_CHECKING, Any, Protocol, TypeVar +from typing import TYPE_CHECKING, Any, TypeVar import numpy as np @@ -20,21 +20,6 @@ T = TypeVar("T") -class ContextMap(Protocol[T]): - def __getitem__(self, key: str) -> T: ... - - -class MdpContext(Protocol): - """Logical context surface shared by NumPy/Numba and Torch MDP lowering.""" - - env_id: Any - actions: ContextMap[Any] - commands: ContextMap[Any] - metrics: ContextMap[Any] - sim: ContextMap[Any] - dt: Any - - @kernel_data class ManagerContext: """Framework-owned manager data available to every compiled term.""" @@ -96,4 +81,4 @@ def command_terms(self) -> Any: return self._env.command_terms -__all__ = ["BuildContext", "ContextMap", "ManagerContext", "MdpContext"] +__all__ = ["BuildContext", "ManagerContext"] From dbb1429620685c460f3489bd4443c4074707207b Mon Sep 17 00:00:00 2001 From: zilch <147668916@qq.com> Date: Thu, 5 Nov 2026 23:46:27 +0800 Subject: [PATCH 09/41] refactor(wbt): drop the batch_math module alias for quaternion_array --- .../src/motrix_envs/locomotion/wbt/mdp/observations.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/observations.py b/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/observations.py index 3b5715e4..3a5cfaaf 100644 --- a/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/observations.py +++ b/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/observations.py @@ -15,7 +15,7 @@ kernel_data, ) from motrix_env_core.manager.kernel_data import SharedArray -from motrix_env_core.math import quaternion_array as batch_math +from motrix_env_core.math import quaternion_array from motrix_env_core.math.quaternion_array import relative_orientation, rotate_inverse from motrix_env_core.mdp.observations import ( UniformNoiseCfg, @@ -105,8 +105,10 @@ def robot_body_orientation_obs(ctx: ManagerContext, out: torch.Tensor) -> None: # explicitly. tracked_body_quat = ctx.sim["tracked_body_quat"] reference_quat = tracked_body_quat[:, ctx.commands["motion"].reference_index] - first = batch_math.matrix6d(batch_math.mul(batch_math.inverse(reference_quat), tracked_body_quat[:, 0, :])) - rest = batch_math.matrix6d(batch_math.mul(first[:, None, :4], tracked_body_quat[:, 1:, :])) + first = quaternion_array.matrix6d( + quaternion_array.mul(quaternion_array.inverse(reference_quat), tracked_body_quat[:, 0, :]) + ) + rest = quaternion_array.matrix6d(quaternion_array.mul(first[:, None, :4], tracked_body_quat[:, 1:, :])) out.copy_(torch.cat((first[:, None, :], rest), dim=1).reshape(tracked_body_quat.shape[0], -1)) From 1ebb9a12b16a90d5603596b2766b4e289802f462 Mon Sep 17 00:00:00 2001 From: zilch <147668916@qq.com> Date: Thu, 5 Nov 2026 23:48:38 +0800 Subject: [PATCH 10/41] refactor(manager): type the task program as TorchTaskProgram 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. --- motrix_env_core/src/motrix_env_core/manager/env.py | 7 +++++-- .../src/motrix_env_core/torch_manager/runtime.py | 7 ++++--- 2 files changed, 9 insertions(+), 5 deletions(-) diff --git a/motrix_env_core/src/motrix_env_core/manager/env.py b/motrix_env_core/src/motrix_env_core/manager/env.py index 249bc0a1..0f898a67 100644 --- a/motrix_env_core/src/motrix_env_core/manager/env.py +++ b/motrix_env_core/src/motrix_env_core/manager/env.py @@ -57,6 +57,9 @@ SimQueriesCfg, ) from motrix_env_core.sim.backend import RenderConfig, SimBackend, SimRenderer + +if TYPE_CHECKING: + from motrix_env_core.torch_manager.runtime import TorchTaskProgram from motrix_env_core.sim.model import ActuatorSpec from motrix_env_core.sim.write import CtrlTargetsWrite, SimWrite @@ -343,7 +346,7 @@ def __init__( self._sim_reset_runtime = SimResetRuntime.create(self, cfg.sim_reset, self.sim) from motrix_env_core.mdp.state import _create_rand_value - self._task_program: Any | None = None + self._task_program: TorchTaskProgram | None = None self._kernel_buffers: tuple[Any, ...] | None = None self.metrics: dict[str, np.ndarray] = {} self._kernel_outputs: KernelOutputs | None = None @@ -778,7 +781,7 @@ def manager_layout(self) -> ManagerLayout: raise RuntimeError("Manager layout is available after init_state() builds the task program.") return self._compiled_manager_program.layout - def _build_task_program(self) -> Any: + def _build_task_program(self) -> TorchTaskProgram: from motrix_env_core.torch_manager.binding import build_manager_binding from motrix_env_core.torch_manager.runtime import TorchManagerRuntime diff --git a/motrix_env_core/src/motrix_env_core/torch_manager/runtime.py b/motrix_env_core/src/motrix_env_core/torch_manager/runtime.py index 1c80a877..ef6fc3a5 100644 --- a/motrix_env_core/src/motrix_env_core/torch_manager/runtime.py +++ b/motrix_env_core/src/motrix_env_core/torch_manager/runtime.py @@ -27,6 +27,7 @@ import os import re import warnings +from collections.abc import Callable from dataclasses import dataclass from typing import Any @@ -69,9 +70,9 @@ def _compile_options() -> dict[str, Any]: class TorchTaskProgram: """Drop-in replacement task program backed by Torch kernels.""" - evaluate_kernel: Any - observe_kernel: Any - reset_kernel: Any + evaluate_kernel: Callable[[tuple[Any, ...], np.ndarray, tuple[np.ndarray, ...], Any], None] + observe_kernel: Callable[[tuple[Any, ...], Any], None] + reset_kernel: Callable[[tuple[Any, ...], np.ndarray, tuple[np.ndarray, ...]], None] reward_weights: np.ndarray From 3b7e5861053880ead641213ddf7e06910076960e Mon Sep 17 00:00:00 2001 From: zilch <147668916@qq.com> Date: Thu, 5 Nov 2026 23:49:58 +0800 Subject: [PATCH 11/41] refactor(manager): type SimResetRuntime.apply with TorchTaskProgram The single-attribute structural Protocol was another dual-executor leftover; TorchTaskProgram is the only implementation. --- .../src/motrix_env_core/manager/sim_reset.py | 10 +++------- 1 file changed, 3 insertions(+), 7 deletions(-) diff --git a/motrix_env_core/src/motrix_env_core/manager/sim_reset.py b/motrix_env_core/src/motrix_env_core/manager/sim_reset.py index 442619ba..b771f5fa 100644 --- a/motrix_env_core/src/motrix_env_core/manager/sim_reset.py +++ b/motrix_env_core/src/motrix_env_core/manager/sim_reset.py @@ -18,18 +18,14 @@ from motrix_env_core.sim.write import SimWrite, WriteProgram if TYPE_CHECKING: - from typing import Protocol - from motrix_env_core.manager.env import ManagerEnv from motrix_env_core.sim.backend import SimBackend - - class _TaskProgram(Protocol): - reset_kernel: Any + from motrix_env_core.torch_manager.runtime import TorchTaskProgram @dataclass(frozen=True, slots=True, init=False) class ResetTerm(BaseTerm): - """Immutable reset dispatch descriptor and its static Numba-compatible arguments.""" + """Immutable reset dispatch descriptor and its static positional arguments.""" writes: dict[str, SimWrite] @@ -99,7 +95,7 @@ def create( def apply( self, - task_program: _TaskProgram, + task_program: TorchTaskProgram, inputs: tuple[Any, ...], env_ids: np.ndarray, ) -> None: From 249c6d6403918cb024514405873b7785ac1d0ed3 Mon Sep 17 00:00:00 2001 From: zilch <147668916@qq.com> Date: Thu, 5 Nov 2026 23:54:20 +0800 Subject: [PATCH 12/41] refactor(manager): collapse term-arg leaf checks into one whitelist 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. --- .../src/motrix_env_core/manager/terms.py | 24 +++++++------------ 1 file changed, 9 insertions(+), 15 deletions(-) diff --git a/motrix_env_core/src/motrix_env_core/manager/terms.py b/motrix_env_core/src/motrix_env_core/manager/terms.py index 88110c43..e5bd55f8 100644 --- a/motrix_env_core/src/motrix_env_core/manager/terms.py +++ b/motrix_env_core/src/motrix_env_core/manager/terms.py @@ -14,35 +14,29 @@ from motrix_env_core.manager.kernel_data import canonicalize_kernel_data, is_kernel_data from motrix_env_core.sim.read import SimDataQuery - -def _is_torch_tensor(value: object) -> bool: - return isinstance(value, torch.Tensor) +# Static term-argument leaves accepted as-is; anything else is a cfg-authoring +# error caught at term construction. +_STATIC_LEAVES = (bool, int, float, str, np.generic, np.ndarray, torch.Tensor, SimDataQuery) def canonicalize_term_args(args: tuple[Any, ...], *, context: str) -> tuple[Any, ...]: - """Validate and canonicalize positional Numba-compatible term arguments. + """Validate and canonicalize positional term arguments. Accepted leaves are scalars (including strings for compile-time Map-name - lookups), scalar tuples, ndarrays, simulator data queries, and - ``@kernel_data`` carriers. + lookups), scalar tuples, ndarrays, tensors, simulator data queries, and + ``@kernel_data`` carriers (which are normalized in place). """ values: list[Any] = [] for index, value in enumerate(args): if is_kernel_data(value): values.append(canonicalize_kernel_data(value, context=f"{context} args[{index}]")) - elif isinstance(value, tuple) and all(isinstance(item, (bool, int, float, np.generic)) for item in value): - values.append(value) - elif isinstance(value, (bool, int, float, str, np.generic)): + elif isinstance(value, _STATIC_LEAVES): values.append(value) - elif isinstance(value, np.ndarray): - values.append(value) - elif _is_torch_tensor(value): - values.append(value) - elif isinstance(value, SimDataQuery): + elif isinstance(value, tuple) and all(isinstance(item, (bool, int, float, np.generic)) for item in value): values.append(value) else: raise TypeError( - f"{context} args[{index}] must be a scalar, a scalar tuple, an ndarray, " + f"{context} args[{index}] must be a scalar, a scalar tuple, an ndarray, a tensor, " f"a @kernel_data value, or a simulator data query; got {type(value).__name__}." ) return tuple(values) From 71cf72cc347776bd49180504da301cea4804022e Mon Sep 17 00:00:00 2001 From: zilch <147668916@qq.com> Date: Thu, 5 Nov 2026 23:57:29 +0800 Subject: [PATCH 13/41] refactor(math): inline the array-namespace shim into quaternion_array 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. --- .../src/motrix_env_core/math/__init__.py | 4 ++-- .../src/motrix_env_core/math/namespace.py | 13 ---------- .../motrix_env_core/math/quaternion_array.py | 24 ++++++++++++------- 3 files changed, 18 insertions(+), 23 deletions(-) delete mode 100644 motrix_env_core/src/motrix_env_core/math/namespace.py diff --git a/motrix_env_core/src/motrix_env_core/math/__init__.py b/motrix_env_core/src/motrix_env_core/math/__init__.py index aad97831..e12fef54 100644 --- a/motrix_env_core/src/motrix_env_core/math/__init__.py +++ b/motrix_env_core/src/motrix_env_core/math/__init__.py @@ -1,6 +1,6 @@ # Copyright Motphys Technology Co., Ltd. 2025, 2026 # SPDX-License-Identifier: Apache-2.0 -from motrix_env_core.math import namespace, quaternion, quaternion_array, utils +from motrix_env_core.math import quaternion, quaternion_array, utils -__all__ = ["namespace", "quaternion", "quaternion_array", "utils"] +__all__ = ["quaternion", "quaternion_array", "utils"] diff --git a/motrix_env_core/src/motrix_env_core/math/namespace.py b/motrix_env_core/src/motrix_env_core/math/namespace.py deleted file mode 100644 index ea417ca1..00000000 --- a/motrix_env_core/src/motrix_env_core/math/namespace.py +++ /dev/null @@ -1,13 +0,0 @@ -# Copyright Motphys Technology Co., Ltd. 2025, 2026 -# SPDX-License-Identifier: Apache-2.0 - -"""Type-specialized namespace selection for compiled batch computations.""" - -import array_api_compat.numpy as numpy_namespace -import array_api_compat.torch as torch_namespace -import torch - - -def array_namespace(value, *others): - # Dynamo specializes this type check; no runtime backend switch in kernels. - return torch_namespace if isinstance(value, torch.Tensor) else numpy_namespace diff --git a/motrix_env_core/src/motrix_env_core/math/quaternion_array.py b/motrix_env_core/src/motrix_env_core/math/quaternion_array.py index e347e896..0498478a 100644 --- a/motrix_env_core/src/motrix_env_core/math/quaternion_array.py +++ b/motrix_env_core/src/motrix_env_core/math/quaternion_array.py @@ -3,17 +3,25 @@ """Shared NumPy/Torch batch quaternion operations, using xyzw convention.""" -from motrix_env_core.math.namespace import array_namespace +import array_api_compat.numpy as _numpy_namespace +import array_api_compat.torch as _torch_namespace +import torch as _torch # noqa: F401 (namespace check only) + + +def _array_namespace(value, *others): + # A single isinstance check keeps Dynamo specialization possible inside + # compiled graphs; array_api_compat's own dispatcher is too dynamic here. + return _torch_namespace if isinstance(value, _torch.Tensor) else _numpy_namespace def inverse(q): - xp = array_namespace(q) + xp = _array_namespace(q) norm = xp.maximum(xp.sum(q * q, axis=-1, keepdims=True), xp.full_like(q[..., :1], 1e-12)) return xp.concat((-q[..., :3], q[..., 3:]), axis=-1) / norm def mul(lhs, rhs): - xp = array_namespace(lhs, rhs) + xp = _array_namespace(lhs, rhs) x1, y1, z1, w1 = (lhs[..., i] for i in range(4)) x2, y2, z2, w2 = (rhs[..., i] for i in range(4)) return xp.stack( @@ -28,7 +36,7 @@ def mul(lhs, rhs): def rotate_inverse(q, vector): - xp = array_namespace(q, vector) + xp = _array_namespace(q, vector) x, y, z, w = (q[..., i] for i in range(4)) vx, vy, vz = (vector[..., i] for i in range(3)) cx = y * vz - z * vy - w * vx @@ -38,7 +46,7 @@ def rotate_inverse(q, vector): def matrix6d(q): - xp = array_namespace(q) + xp = _array_namespace(q) x, y, z, w = (q[..., i] for i in range(4)) s = 2 / xp.sum(q * q, axis=-1) return xp.stack( @@ -60,7 +68,7 @@ def relative_orientation(reference, target): def align_targets(robot_position, robot_quat, reference_position, reference_quat, body_position, body_quat): """Yaw-align corpus bodies, anchoring robot XY and corpus reference Z.""" - xp = array_namespace(robot_position) + xp = _array_namespace(robot_position) delta = mul(robot_quat, inverse(reference_quat)) x, y, z, w = (delta[..., i] for i in range(4)) yaw = xp.atan2(2 * (w * z + x * y), 1 - 2 * (y * y + z * z)) @@ -75,7 +83,7 @@ def align_targets(robot_position, robot_quat, reference_position, reference_quat def rotation_distance(lhs, rhs): """Geodesic angle between unit quaternions, without normalization.""" - xp = array_namespace(lhs, rhs) + xp = _array_namespace(lhs, rhs) x1, y1, z1, w1 = (lhs[..., i] for i in range(4)) x2, y2, z2, w2 = (rhs[..., i] for i in range(4)) dx = -w1 * x2 + x1 * w2 - y1 * z2 + z1 * y2 @@ -88,7 +96,7 @@ def rotation_distance(lhs, rhs): def from_euler(euler): """Quaternion (xyzw) from ``(..., 3)`` roll/pitch/yaw angles; trig keeps the input dtype.""" - xp = array_namespace(euler) + xp = _array_namespace(euler) roll, pitch, yaw = (euler[..., i] for i in range(3)) cy = xp.cos(yaw * 0.5) sy = xp.sin(yaw * 0.5) From fae967254e4dff6e2c7ec8a06d7c1b2f336886fb Mon Sep 17 00:00:00 2001 From: zilch <147668916@qq.com> Date: Thu, 5 Nov 2026 23:59:38 +0800 Subject: [PATCH 14/41] refactor(manager): inline the manager binding shim into env binding.py collapsed to a one-line ManagerProgramResolver delegation after the resolver extraction; call it directly from _build_task_program. --- .../src/motrix_env_core/manager/env.py | 4 ++-- .../motrix_env_core/torch_manager/binding.py | 20 ------------------- 2 files changed, 2 insertions(+), 22 deletions(-) delete mode 100644 motrix_env_core/src/motrix_env_core/torch_manager/binding.py diff --git a/motrix_env_core/src/motrix_env_core/manager/env.py b/motrix_env_core/src/motrix_env_core/manager/env.py index 0f898a67..3ec0218c 100644 --- a/motrix_env_core/src/motrix_env_core/manager/env.py +++ b/motrix_env_core/src/motrix_env_core/manager/env.py @@ -782,10 +782,10 @@ def manager_layout(self) -> ManagerLayout: return self._compiled_manager_program.layout def _build_task_program(self) -> TorchTaskProgram: - from motrix_env_core.torch_manager.binding import build_manager_binding + from motrix_env_core.manager.resolver import ManagerProgramResolver from motrix_env_core.torch_manager.runtime import TorchManagerRuntime - self._compiled_manager_program = build_manager_binding(self) + self._compiled_manager_program = ManagerProgramResolver(self).build() self._torch_runtime = TorchManagerRuntime(self, self._compiled_manager_program) return self._torch_runtime.task_program() diff --git a/motrix_env_core/src/motrix_env_core/torch_manager/binding.py b/motrix_env_core/src/motrix_env_core/torch_manager/binding.py deleted file mode 100644 index 828ea31d..00000000 --- a/motrix_env_core/src/motrix_env_core/torch_manager/binding.py +++ /dev/null @@ -1,20 +0,0 @@ -# Copyright Motphys Technology Co., Ltd. 2025, 2026 -# SPDX-License-Identifier: Apache-2.0 - -"""Manager metadata binding for the Torch executor.""" - -from motrix_env_core.manager.env import ManagerEnv -from motrix_env_core.manager.resolver import ManagerProgramResolver, ResolvedManagerProgram - - -def build_manager_binding(env: ManagerEnv) -> ResolvedManagerProgram: - """Resolve an already constructed manager environment before ``init_state``. - - The result supplies the read plan, layouts, prepared inputs, and Python - term entries to the Torch runtime; this helper neither initializes the - environment nor compiles any numerical kernels. - """ - return ManagerProgramResolver(env).build() - - -__all__ = ["build_manager_binding"] From 572f7391b76a7f24dc05970e7730620620c08b27 Mon Sep 17 00:00:00 2001 From: zilch <147668916@qq.com> Date: Fri, 6 Nov 2026 00:02:02 +0800 Subject: [PATCH 15/41] refactor(ball): build reset noise tensors directly from tuples torch.tensor accepts sequences; routing them through numpy was a migration habit from the numba scalar-tuple literals. --- .../locomotion/ball_balance/mdp/reset.py | 20 +++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/motrix_envs/src/motrix_envs/locomotion/ball_balance/mdp/reset.py b/motrix_envs/src/motrix_envs/locomotion/ball_balance/mdp/reset.py index df29480d..6c0e690a 100644 --- a/motrix_envs/src/motrix_envs/locomotion/ball_balance/mdp/reset.py +++ b/motrix_envs/src/motrix_envs/locomotion/ball_balance/mdp/reset.py @@ -45,8 +45,8 @@ def __call__(self, ctx) -> ResetTerm: body = ctx.cfg.scene.objs.robot.resolved_base_link_name return ResetTerm( _reset_body_pos, - torch.from_numpy(np.ascontiguousarray(self.spawn, dtype=np.float32)), - torch.from_numpy(np.ascontiguousarray(self.noise, dtype=np.float32) * np.float32(self.noise_scale)), + torch.tensor(self.spawn, dtype=torch.float32), + torch.tensor(self.noise, dtype=torch.float32) * self.noise_scale, writes={"position": BodyPositionWrite((body,))}, ) @@ -73,10 +73,10 @@ class BodyRotResetCfg(ResetTermCfg): def __call__(self, ctx) -> ResetTerm: body = ctx.cfg.scene.objs.robot.resolved_base_link_name - amplitude = np.asarray(self.noise, dtype=np.float32) * np.float32(self.noise_scale) + amplitude = torch.tensor(self.noise, dtype=torch.float32) * self.noise_scale return ResetTerm( _reset_body_rot, - torch.from_numpy(amplitude), + amplitude, writes={"rotation": BodyRotationWrite((body,))}, ) @@ -95,10 +95,10 @@ class BodyLinVelResetCfg(ResetTermCfg): def __call__(self, ctx) -> ResetTerm: body = ctx.cfg.scene.objs.robot.resolved_base_link_name - amplitude = np.asarray(self.noise, dtype=np.float32) * np.float32(self.noise_scale) + amplitude = torch.tensor(self.noise, dtype=torch.float32) * self.noise_scale return ResetTerm( _reset_body_lin_vel, - torch.from_numpy(amplitude), + amplitude, writes={"linear_velocity": BodyLinearVelocityWrite((body,))}, ) @@ -117,10 +117,10 @@ class BodyRotVelResetCfg(ResetTermCfg): def __call__(self, ctx) -> ResetTerm: body = ctx.cfg.scene.objs.robot.resolved_base_link_name - amplitude = np.asarray(self.noise, dtype=np.float32) * np.float32(self.noise_scale) + amplitude = torch.tensor(self.noise, dtype=torch.float32) * self.noise_scale return ResetTerm( _reset_body_rot_vel, - torch.from_numpy(amplitude), + amplitude, writes={"angular_velocity": BodyAngularVelocityWrite((body,))}, ) @@ -205,8 +205,8 @@ class BallResetCfg(ResetTermCfg): def __call__(self, ctx) -> ResetTerm: return ResetTerm( _reset_ball, - torch.from_numpy(np.ascontiguousarray(self.spawn, dtype=np.float32)), - torch.from_numpy(np.ascontiguousarray(self.noise, dtype=np.float32) * np.float32(self.noise_scale)), + torch.tensor(self.spawn, dtype=torch.float32), + torch.tensor(self.noise, dtype=torch.float32) * self.noise_scale, torch.from_numpy( np.ascontiguousarray(self.velocity_noise, dtype=np.float32) * np.float32(self.noise_scale) ), From f18d3725e5bdd48cb4adb66c109879bb3b416676 Mon Sep 17 00:00:00 2001 From: zilch <147668916@qq.com> Date: Fri, 6 Nov 2026 00:07:46 +0800 Subject: [PATCH 16/41] refactor: drop the last numpy stepping stones in reset/reward term args --- .../src/motrix_envs/locomotion/ball_balance/mdp/reset.py | 4 +--- .../locomotion/humanoid/walk_manager_mdp/rewards.py | 2 +- 2 files changed, 2 insertions(+), 4 deletions(-) diff --git a/motrix_envs/src/motrix_envs/locomotion/ball_balance/mdp/reset.py b/motrix_envs/src/motrix_envs/locomotion/ball_balance/mdp/reset.py index 6c0e690a..c50ccb23 100644 --- a/motrix_envs/src/motrix_envs/locomotion/ball_balance/mdp/reset.py +++ b/motrix_envs/src/motrix_envs/locomotion/ball_balance/mdp/reset.py @@ -207,9 +207,7 @@ def __call__(self, ctx) -> ResetTerm: _reset_ball, torch.tensor(self.spawn, dtype=torch.float32), torch.tensor(self.noise, dtype=torch.float32) * self.noise_scale, - torch.from_numpy( - np.ascontiguousarray(self.velocity_noise, dtype=np.float32) * np.float32(self.noise_scale) - ), + torch.tensor(self.velocity_noise, dtype=torch.float32) * self.noise_scale, writes={ "position": BodyPositionWrite((self.ball_link_name,)), "linear_velocity": BodyLinearVelocityWrite((self.ball_link_name,)), diff --git a/motrix_envs/src/motrix_envs/locomotion/humanoid/walk_manager_mdp/rewards.py b/motrix_envs/src/motrix_envs/locomotion/humanoid/walk_manager_mdp/rewards.py index f6c9fd67..ebee3d8d 100644 --- a/motrix_envs/src/motrix_envs/locomotion/humanoid/walk_manager_mdp/rewards.py +++ b/motrix_envs/src/motrix_envs/locomotion/humanoid/walk_manager_mdp/rewards.py @@ -209,7 +209,7 @@ def penalty_feet_ori_reward( gravity = foot_quat.new_zeros(*foot_quat.shape[:2], 3) gravity[..., 2] = -1.0 gx, gy, gz = rotate_inverse_components(foot_quat.reshape(-1, 4), gravity.reshape(-1, 3)) - reference = torch.as_tensor(np.asarray(default_foot_gravity)) + reference = torch.as_tensor(default_foot_gravity) total = None for foot in range(num_feet): # Flat layout is lane-major (lane * F + foot): strided slices pick one foot. From c05ea6fa27c1101846932bc8855b4ad339906479 Mon Sep 17 00:00:00 2001 From: zilch <147668916@qq.com> Date: Fri, 6 Nov 2026 00:17:15 +0800 Subject: [PATCH 17/41] refactor(torch): type TorchManagerRuntime with ManagerEnv and ResolvedManagerProgram --- .../src/motrix_env_core/torch_manager/runtime.py | 16 +++++++++++++--- 1 file changed, 13 insertions(+), 3 deletions(-) diff --git a/motrix_env_core/src/motrix_env_core/torch_manager/runtime.py b/motrix_env_core/src/motrix_env_core/torch_manager/runtime.py index ef6fc3a5..86fdff97 100644 --- a/motrix_env_core/src/motrix_env_core/torch_manager/runtime.py +++ b/motrix_env_core/src/motrix_env_core/torch_manager/runtime.py @@ -29,7 +29,7 @@ import warnings from collections.abc import Callable from dataclasses import dataclass -from typing import Any +from typing import TYPE_CHECKING, Any import numpy as np import torch @@ -37,6 +37,10 @@ from motrix_env_core.manager.kernel_data import KernelDataScope from motrix_env_core.torch_manager.rand import draw_uniform +if TYPE_CHECKING: + from motrix_env_core.manager.env import ManagerEnv + from motrix_env_core.manager.resolver import ResolvedManagerProgram + _INPUT_EXPRESSION = re.compile(r"^input_(\d+)(\[env_id\])?$") # Simulator-backed read-only arrays are wrapped as tensors but never written @@ -321,7 +325,13 @@ class _BatchContext: context operate on whole batches. """ - def __init__(self, inputs: tuple[Any, ...], cache: _TensorCache, paths: dict[tuple[str, ...], int], env: Any): + def __init__( + self, + inputs: tuple[Any, ...], + cache: _TensorCache, + paths: dict[tuple[str, ...], int], + env: ManagerEnv, + ): get = _leaf_getter(cache, inputs) self._paths = paths self.env_id = get(paths[("env_id",)]) @@ -358,7 +368,7 @@ class TorchManagerRuntime: per-term eager schedule for debugging. """ - def __init__(self, env: Any, program: Any, compile: bool = True) -> None: + def __init__(self, env: ManagerEnv, program: ResolvedManagerProgram, compile: bool = True) -> None: self.env = env self.program = program self.layout = program.layout From 411f425dc306c9528628a28cf1c0d83f518d16ef Mon Sep 17 00:00:00 2001 From: zilch <147668916@qq.com> Date: Fri, 6 Nov 2026 00:22:38 +0800 Subject: [PATCH 18/41] refactor(torch): make the compiled kernel the single implementation per 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. --- .../motrix_env_core/torch_manager/runtime.py | 155 ++++-------------- 1 file changed, 28 insertions(+), 127 deletions(-) diff --git a/motrix_env_core/src/motrix_env_core/torch_manager/runtime.py b/motrix_env_core/src/motrix_env_core/torch_manager/runtime.py index 86fdff97..ef543ded 100644 --- a/motrix_env_core/src/motrix_env_core/torch_manager/runtime.py +++ b/motrix_env_core/src/motrix_env_core/torch_manager/runtime.py @@ -449,136 +449,32 @@ def task_program(self) -> TorchTaskProgram: ) def evaluate(self, inputs: tuple[Any, ...], reward_weights: np.ndarray, buffers, outputs) -> None: - if self.compile_enabled: - kernel = self._compiled_kernels.get(("evaluate", id(inputs), id(buffers))) - if kernel is None: - kernel = self._build_compiled_evaluate(inputs, reward_weights, buffers, outputs) - self._compiled_kernels[("evaluate", id(inputs), id(buffers))] = kernel - kernel(torch.from_numpy(outputs.reward), torch.from_numpy(outputs.terminated)) - return - self._evaluate_eager(inputs, reward_weights, buffers, outputs) - - def _evaluate_eager(self, inputs: tuple[Any, ...], reward_weights: np.ndarray, buffers, outputs) -> None: - ctx = self._context_for(inputs) - weights = torch.from_numpy(np.asarray(reward_weights, dtype=np.float32)) - raw_terms = torch.from_numpy(buffers[0]) - weighted_terms = torch.from_numpy(buffers[1]) - termination_masks = torch.from_numpy(buffers[2]) - dt = ctx.dt - total = None - terminated = None - for invocation, args in self._evaluate_plan(inputs): - kind = invocation.kind - if kind in ("command_update", "command_advance"): - receiver = ctx.commands[invocation.receiver_key] - invocation.dispatcher(receiver, ctx) - elif kind == "reward": - value = invocation.dispatcher(ctx, *args).to(torch.float32) - index = self._reward_indices[invocation.context] - weighted = value * weights[index] * dt - raw_terms[:, index] = value - weighted_terms[:, index] = weighted - total = weighted if total is None else total + weighted - elif kind == "termination": - mask = invocation.dispatcher(ctx, *args).to(torch.bool) - index = self._termination_indices[invocation.context] - termination_masks[:, index] = mask - terminated = mask if terminated is None else terminated | mask - else: # pragma: no cover - schedule is fixed at resolve time - raise RuntimeError(f"Unsupported evaluate invocation kind {kind!r}.") - reward_view = torch.from_numpy(outputs.reward) - reward_view.copy_(total if total is not None else torch.zeros_like(reward_view)) - terminated_view = torch.from_numpy(outputs.terminated) - terminated_view.copy_(terminated if terminated is not None else torch.zeros_like(terminated_view)) + kernel = self._compiled_kernels.get(("evaluate", id(inputs), id(buffers))) + if kernel is None: + kernel = self._build_compiled_evaluate(inputs, reward_weights, buffers, outputs) + self._compiled_kernels[("evaluate", id(inputs), id(buffers))] = kernel + kernel(torch.from_numpy(outputs.reward), torch.from_numpy(outputs.terminated)) def observe(self, inputs: tuple[Any, ...], outputs) -> None: - if self.compile_enabled: - kernel = self._compiled_kernels.get(("observe", id(inputs))) - if kernel is None: - kernel = self._build_compiled_observe(inputs, outputs) - self._compiled_kernels[("observe", id(inputs))] = kernel - value_obs = outputs.value_obs - kernel( - torch.from_numpy(outputs.policy_obs), - torch.from_numpy(value_obs) if value_obs is not None and value_obs.size else None, - ) - return - self._observe_eager(inputs, outputs) - - def _observe_eager(self, inputs: tuple[Any, ...], outputs) -> None: - ctx = self._context_for(inputs) - group_tensors = {"policy": torch.from_numpy(outputs.policy_obs)} - if outputs.value_obs is not None and outputs.value_obs.size: - group_tensors["value"] = torch.from_numpy(outputs.value_obs) - for group_name, group in self.layout.observations.items(): - group_tensor = group_tensors[group_name] - for term in group.terms: - invocation = self._observation_by_name[f"{group_name}.{term.name}"] - out = group_tensor[:, term.output_slice] - invocation.dispatcher(ctx, out, *self._invocation_args_for(invocation, inputs)) + kernel = self._compiled_kernels.get(("observe", id(inputs))) + if kernel is None: + kernel = self._build_compiled_observe(inputs, outputs) + self._compiled_kernels[("observe", id(inputs))] = kernel + value_obs = outputs.value_obs + kernel( + torch.from_numpy(outputs.policy_obs), + torch.from_numpy(value_obs) if value_obs is not None and value_obs.size else None, + ) def reset(self, inputs: tuple[Any, ...], env_ids: np.ndarray, reset_buffers) -> None: if not self._reset_env_invocations and not self._reset_terms: return ids = torch.from_numpy(np.ascontiguousarray(env_ids, dtype=np.int64)) - if self.compile_enabled: - kernel = self._compiled_kernels.get(("reset", id(inputs), id(reset_buffers))) - if kernel is None: - kernel = self._build_compiled_reset(inputs, reset_buffers) - self._compiled_kernels[("reset", id(inputs), id(reset_buffers))] = kernel - kernel(ids) - return - self._reset_eager(inputs, ids, reset_buffers) - - def _reset_eager(self, inputs: tuple[Any, ...], ids: torch.Tensor, reset_buffers) -> None: - ctx = self._context_for(inputs) - - gathered: dict[int, torch.Tensor] = {} - - def gather(slot: int, tensor: torch.Tensor) -> torch.Tensor: - cached = gathered.get(slot) - if cached is None: - cached = tensor.index_select(0, ids) - gathered[slot] = cached - return cached - - get = _leaf_getter(self._cache, inputs) - per_env = self._per_env_slots - - def lane_get(slot: int) -> torch.Tensor: - full = get(slot) - return gather(slot, full) if slot in per_env else full - - lane_rand = _LaneRand(self._context_paths[("rand", "state")], ctx, gather) - lane_ctx = _LaneContext( - ctx, - _CommandsMap(self._context_paths, self.env.command_terms, lane_get), - lane_rand, - _SimMap( - lane_get, - {path[1]: slot for path, slot in self._context_paths.items() if len(path) == 2 and path[0] == "sim"}, - ), - _MetricsMap(self._context_paths, lane_get), - ids, - actions=_ActionsMap(self._context_paths, self.env.action_terms, lane_get), - ) - for invocation in self._reset_env_invocations: - receiver = lane_ctx.commands[invocation.receiver_key] - invocation.dispatcher(receiver, lane_ctx) - for invocation, offset, fields in self._reset_terms: - sim_writes = { - name: torch.from_numpy(reset_buffers[offset + index][: ids.shape[0]]) - for index, name in enumerate(fields) - } - invocation.dispatcher(lane_ctx, _WritesMap(sim_writes), *self._invocation_args_for(invocation, inputs)) - for slot, tensor in gathered.items(): - full = self._cache.tensor(inputs, slot) - if tensor.shape == full.shape: - full.copy_(tensor) - else: - full.index_copy_(0, ids, tensor) - - # -- compiled schedules ------------------------------------------------- + kernel = self._compiled_kernels.get(("reset", id(inputs), id(reset_buffers))) + if kernel is None: + kernel = self._build_compiled_reset(inputs, reset_buffers) + self._compiled_kernels[("reset", id(inputs), id(reset_buffers))] = kernel + kernel(ids) def _build_compiled_evaluate(self, inputs, reward_weights, buffers, outputs): ctx = self._context_for(inputs) @@ -613,7 +509,9 @@ def kernel(reward_view: torch.Tensor, terminated_view: torch.Tensor) -> None: reward_view.copy_(total if total is not None else torch.zeros_like(reward_view)) terminated_view.copy_(terminated if terminated is not None else torch.zeros_like(terminated_view)) - return torch.compile(kernel, dynamic=False, options=_compile_options()) + if self.compile_enabled: + return torch.compile(kernel, dynamic=False, options=_compile_options()) + return kernel def _build_compiled_observe(self, inputs, outputs): ctx = self._context_for(inputs) @@ -662,7 +560,9 @@ def group_kernel(target: torch.Tensor) -> None: invocation.dispatcher(ctx, buf, *args) target.copy_(torch.cat([buf for _inv, _args, buf in entries], dim=1)) - return torch.compile(group_kernel, dynamic=False, options=_compile_options()) + if self.compile_enabled: + return torch.compile(group_kernel, dynamic=False, options=_compile_options()) + return group_kernel group_kernels = {group_name: make_group_kernel(entries) for group_name, entries in per_term_plan} policy_kernel = group_kernels.get("policy") @@ -747,7 +647,8 @@ def lane_get(slot: int) -> torch.Tensor: lane_mask = mask.view((num_envs,) + (1,) * (tensor.dim() - 1)) full.copy_(torch.where(lane_mask, tensor, full)) - compiled = torch.compile(kernel, dynamic=False, options=_compile_options()) + if self.compile_enabled: + kernel = torch.compile(kernel, dynamic=False, options=_compile_options()) mask_np = mask.numpy() compact_plan = tuple( (buffer.numpy(), stage.numpy()) @@ -759,7 +660,7 @@ def wrapper(ids: torch.Tensor) -> None: ids_np = ids.numpy() mask_np.fill(False) mask_np[ids_np] = True - compiled() + kernel() count = ids_np.shape[0] for buffer, stage in compact_plan: buffer[:count] = stage[ids_np] From d100ceae506f69ef110e63689348b21d00968bd5 Mon Sep 17 00:00:00 2001 From: zilch <147668916@qq.com> Date: Fri, 6 Nov 2026 00:30:13 +0800 Subject: [PATCH 19/41] refactor(torch): annotate the runtime's kernel and helper signatures --- .../motrix_env_core/torch_manager/runtime.py | 69 ++++++++++++++----- 1 file changed, 53 insertions(+), 16 deletions(-) diff --git a/motrix_env_core/src/motrix_env_core/torch_manager/runtime.py b/motrix_env_core/src/motrix_env_core/torch_manager/runtime.py index ef543ded..20abed49 100644 --- a/motrix_env_core/src/motrix_env_core/torch_manager/runtime.py +++ b/motrix_env_core/src/motrix_env_core/torch_manager/runtime.py @@ -34,12 +34,13 @@ import numpy as np import torch +from motrix_env_core.manager.kernel import KernelOutputs from motrix_env_core.manager.kernel_data import KernelDataScope from motrix_env_core.torch_manager.rand import draw_uniform if TYPE_CHECKING: from motrix_env_core.manager.env import ManagerEnv - from motrix_env_core.manager.resolver import ResolvedManagerProgram + from motrix_env_core.manager.resolver import PreparedInvocation, ResolvedManagerProgram _INPUT_EXPRESSION = re.compile(r"^input_(\d+)(\[env_id\])?$") @@ -139,7 +140,7 @@ def tensor(self, inputs: tuple[Any, ...], slot: int) -> torch.Tensor: return tensor -def _leaf_getter(cache: _TensorCache, inputs: tuple[Any, ...]): +def _leaf_getter(cache: _TensorCache, inputs: tuple[Any, ...]) -> Callable[[int], torch.Tensor]: """Leaf resolver over one input tuple. All tensor views are materialized eagerly so the getter body is a plain @@ -165,7 +166,13 @@ class _Namespace: tensor views stay valid for the lifetime of the enclosing input tuple. """ - def __init__(self, paths: dict[tuple[str, ...], int], prefix: tuple[str, ...], host: Any, get) -> None: + def __init__( + self, + paths: dict[tuple[str, ...], int], + prefix: tuple[str, ...], + host: Any, + get: Callable[[int], torch.Tensor], + ) -> None: self._paths = paths self._prefix = prefix self._host = host @@ -264,7 +271,7 @@ def __getattr__(self, name: str) -> Any: class _SimMap: - def __init__(self, get, slots: dict[str, int]) -> None: + def __init__(self, get: Callable[[int], torch.Tensor], slots: dict[str, int]) -> None: self._get = get self._slots = slots @@ -273,7 +280,12 @@ def __getitem__(self, key: str) -> torch.Tensor: class _CommandsMap: - def __init__(self, paths: dict[tuple[str, ...], int], command_terms: dict[str, Any], get) -> None: + def __init__( + self, + paths: dict[tuple[str, ...], int], + command_terms: dict[str, Any], + get: Callable[[int], torch.Tensor], + ) -> None: self._paths = paths self._command_terms = command_terms self._get = get @@ -291,7 +303,12 @@ def __getitem__(self, name: str) -> _Namespace: class _ActionsMap: - def __init__(self, paths: dict[tuple[str, ...], int], action_terms: dict[str, Any], get) -> None: + def __init__( + self, + paths: dict[tuple[str, ...], int], + action_terms: dict[str, Any], + get: Callable[[int], torch.Tensor], + ) -> None: self._paths = paths self._action_terms = action_terms self._get = get @@ -310,7 +327,7 @@ def __getitem__(self, name: str) -> _BatchActions: class _MetricsMap: - def __init__(self, paths: dict[tuple[str, ...], int], get) -> None: + def __init__(self, paths: dict[tuple[str, ...], int], get: Callable[[int], torch.Tensor]) -> None: self._paths = paths self._get = get @@ -448,14 +465,20 @@ def task_program(self) -> TorchTaskProgram: reward_weights=self.reward_weights, ) - def evaluate(self, inputs: tuple[Any, ...], reward_weights: np.ndarray, buffers, outputs) -> None: + def evaluate( + self, + inputs: tuple[Any, ...], + reward_weights: np.ndarray, + buffers: tuple[np.ndarray, np.ndarray, np.ndarray], + outputs: KernelOutputs, + ) -> None: kernel = self._compiled_kernels.get(("evaluate", id(inputs), id(buffers))) if kernel is None: kernel = self._build_compiled_evaluate(inputs, reward_weights, buffers, outputs) self._compiled_kernels[("evaluate", id(inputs), id(buffers))] = kernel kernel(torch.from_numpy(outputs.reward), torch.from_numpy(outputs.terminated)) - def observe(self, inputs: tuple[Any, ...], outputs) -> None: + def observe(self, inputs: tuple[Any, ...], outputs: KernelOutputs) -> None: kernel = self._compiled_kernels.get(("observe", id(inputs))) if kernel is None: kernel = self._build_compiled_observe(inputs, outputs) @@ -466,7 +489,7 @@ def observe(self, inputs: tuple[Any, ...], outputs) -> None: torch.from_numpy(value_obs) if value_obs is not None and value_obs.size else None, ) - def reset(self, inputs: tuple[Any, ...], env_ids: np.ndarray, reset_buffers) -> None: + def reset(self, inputs: tuple[Any, ...], env_ids: np.ndarray, reset_buffers: tuple[np.ndarray, ...]) -> None: if not self._reset_env_invocations and not self._reset_terms: return ids = torch.from_numpy(np.ascontiguousarray(env_ids, dtype=np.int64)) @@ -476,7 +499,13 @@ def reset(self, inputs: tuple[Any, ...], env_ids: np.ndarray, reset_buffers) -> self._compiled_kernels[("reset", id(inputs), id(reset_buffers))] = kernel kernel(ids) - def _build_compiled_evaluate(self, inputs, reward_weights, buffers, outputs): + def _build_compiled_evaluate( + self, + inputs: tuple[Any, ...], + reward_weights: np.ndarray, + buffers: tuple[np.ndarray, np.ndarray, np.ndarray], + outputs: KernelOutputs, + ) -> Callable[[torch.Tensor, torch.Tensor], None]: ctx = self._context_for(inputs) weights = torch.from_numpy(np.asarray(reward_weights, dtype=np.float32)) raw_terms = torch.from_numpy(buffers[0]) @@ -513,7 +542,9 @@ def kernel(reward_view: torch.Tensor, terminated_view: torch.Tensor) -> None: return torch.compile(kernel, dynamic=False, options=_compile_options()) return kernel - def _build_compiled_observe(self, inputs, outputs): + def _build_compiled_observe( + self, inputs: tuple[Any, ...], outputs: KernelOutputs + ) -> Callable[[torch.Tensor, torch.Tensor | None], None]: ctx = self._context_for(inputs) observation_by_name = self._observation_by_name plan = tuple( @@ -554,7 +585,9 @@ def _build_compiled_observe(self, inputs, outputs): # keep their code hot through the interleaved physics/evaluate passes, # which measurably beats a single fused graph on many-core hosts. - def make_group_kernel(entries): + def make_group_kernel( + entries: tuple[tuple[PreparedInvocation, tuple[Any, ...], torch.Tensor], ...], + ) -> Callable[[torch.Tensor], None]: def group_kernel(target: torch.Tensor) -> None: for invocation, args, buf in entries: invocation.dispatcher(ctx, buf, *args) @@ -576,7 +609,9 @@ def kernel(policy_obs: torch.Tensor, value_obs: torch.Tensor | None) -> None: return kernel - def _build_compiled_reset(self, inputs, reset_buffers): + def _build_compiled_reset( + self, inputs: tuple[Any, ...], reset_buffers: tuple[np.ndarray, ...] + ) -> Callable[[torch.Tensor], None]: """Compile the reset schedule as a static full-batch masked graph. The reset math runs for ALL lanes with the selected lanes chosen by a @@ -701,7 +736,7 @@ def _invocation_args_for(self, invocation: Any, inputs: tuple[Any, ...]) -> tupl cache[key] = args return args - def _invocation_args(self, invocation: Any, inputs: tuple[Any, ...]) -> tuple[Any, ...]: + def _invocation_args(self, invocation: PreparedInvocation, inputs: tuple[Any, ...]) -> tuple[Any, ...]: args: list[Any] = [] for value, prepared_index, expression in zip( invocation.args_values, invocation.args_prepared_indices, invocation.args_expressions, strict=True @@ -728,7 +763,9 @@ def _invocation_args(self, invocation: Any, inputs: tuple[Any, ...]) -> tuple[An class _LaneRand: """SplitMix64 draws on gathered reset rows, updated in place.""" - def __init__(self, rand_slot: int, base: _BatchContext, gather) -> None: + def __init__( + self, rand_slot: int, base: _BatchContext, gather: Callable[[int, torch.Tensor], torch.Tensor] + ) -> None: self._rand_slot = rand_slot self._base = base self._gather = gather From 90b3f60c6acbbfe9d7d99b65a393f9107c8910ef Mon Sep 17 00:00:00 2001 From: zilch <147668916@qq.com> Date: Fri, 6 Nov 2026 00:35:15 +0800 Subject: [PATCH 20/41] refactor(torch): drop the unused outputs parameter from kernel builders --- .../src/motrix_env_core/torch_manager/runtime.py | 11 ++++------- 1 file changed, 4 insertions(+), 7 deletions(-) diff --git a/motrix_env_core/src/motrix_env_core/torch_manager/runtime.py b/motrix_env_core/src/motrix_env_core/torch_manager/runtime.py index 20abed49..0c0c571d 100644 --- a/motrix_env_core/src/motrix_env_core/torch_manager/runtime.py +++ b/motrix_env_core/src/motrix_env_core/torch_manager/runtime.py @@ -474,14 +474,14 @@ def evaluate( ) -> None: kernel = self._compiled_kernels.get(("evaluate", id(inputs), id(buffers))) if kernel is None: - kernel = self._build_compiled_evaluate(inputs, reward_weights, buffers, outputs) + kernel = self._build_compiled_evaluate(inputs, reward_weights, buffers) self._compiled_kernels[("evaluate", id(inputs), id(buffers))] = kernel kernel(torch.from_numpy(outputs.reward), torch.from_numpy(outputs.terminated)) def observe(self, inputs: tuple[Any, ...], outputs: KernelOutputs) -> None: kernel = self._compiled_kernels.get(("observe", id(inputs))) if kernel is None: - kernel = self._build_compiled_observe(inputs, outputs) + kernel = self._build_compiled_observe(inputs) self._compiled_kernels[("observe", id(inputs))] = kernel value_obs = outputs.value_obs kernel( @@ -504,7 +504,6 @@ def _build_compiled_evaluate( inputs: tuple[Any, ...], reward_weights: np.ndarray, buffers: tuple[np.ndarray, np.ndarray, np.ndarray], - outputs: KernelOutputs, ) -> Callable[[torch.Tensor, torch.Tensor], None]: ctx = self._context_for(inputs) weights = torch.from_numpy(np.asarray(reward_weights, dtype=np.float32)) @@ -542,9 +541,7 @@ def kernel(reward_view: torch.Tensor, terminated_view: torch.Tensor) -> None: return torch.compile(kernel, dynamic=False, options=_compile_options()) return kernel - def _build_compiled_observe( - self, inputs: tuple[Any, ...], outputs: KernelOutputs - ) -> Callable[[torch.Tensor, torch.Tensor | None], None]: + def _build_compiled_observe(self, inputs: tuple[Any, ...]) -> Callable[[torch.Tensor, torch.Tensor | None], None]: ctx = self._context_for(inputs) observation_by_name = self._observation_by_name plan = tuple( @@ -558,7 +555,7 @@ def _build_compiled_observe( # slice views of the group buffer forces per-term unfusable store # regions in the generated code (aliased memory), which multiplies # the OpenMP fork/join cost of the fused kernel. - num_envs = outputs.policy_obs.shape[0] + num_envs = ctx.env_id.shape[0] term_bufs = { group_name: tuple( torch.empty(num_envs, term.output_slice.stop - term.output_slice.start, dtype=torch.float32) From c7757ec2f340fe0c2543f72ef51abf6627bf8d41 Mon Sep 17 00:00:00 2001 From: zilch <147668916@qq.com> Date: Fri, 6 Nov 2026 01:06:49 +0800 Subject: [PATCH 21/41] refactor(manager)!: resolve the manager data path as trees, not a flat 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. --- .../src/motrix_env_core/manager/env.py | 127 ++---- .../src/motrix_env_core/manager/kernel.py | 50 ++- .../manager/kernel_data/__init__.py | 16 +- .../manager/kernel_data/lowering.py | 150 +------ .../manager/kernel_data/map.py | 56 +-- .../manager/kernel_data/resolved.py | 43 ++ .../src/motrix_env_core/manager/resolver.py | 408 +++++++----------- .../motrix_env_core/torch_manager/runtime.py | 315 +++++++------- motrix_env_core/tests/test_kernel_data.py | 18 +- motrix_env_core/tests/test_manager_env.py | 42 +- motrix_envs/tests/test_walk_randomization.py | 2 +- wiki/design/manager/kernel-abi.md | 32 +- wiki/design/manager/torch-executor.md | 27 +- 13 files changed, 510 insertions(+), 776 deletions(-) create mode 100644 motrix_env_core/src/motrix_env_core/manager/kernel_data/resolved.py diff --git a/motrix_env_core/src/motrix_env_core/manager/env.py b/motrix_env_core/src/motrix_env_core/manager/env.py index 3ec0218c..70ea2f5f 100644 --- a/motrix_env_core/src/motrix_env_core/manager/env.py +++ b/motrix_env_core/src/motrix_env_core/manager/env.py @@ -22,15 +22,12 @@ KernelOutputs, ManagerWarmupResult, clone_kernel_value, + clone_resolved_trees, validate_kernel_context, ) from motrix_env_core.manager.kernel_data import ( - KernelDataLayout, - KernelDataScope, - KernelLeafLayout, + ResolvedTree, canonicalize_kernel_data, - flatten_kernel_data, - iter_layout_leaves, ) if TYPE_CHECKING: @@ -182,78 +179,24 @@ def termination_cfgs(self) -> dict[str, TerminationTermCfg]: @dataclass(frozen=True) class ReadPlan: + """Simulator read program bound to the resolved context tree. + + ``read`` executes the compiled read program; the simulator-backed ndarray + leaves of the resolved context tree are rewritten in place, so every other + resolved leaf and the tree objects themselves stay valid without + rebuilding. + """ + sim_data: PhysicsReadProgram - sources: tuple[KernelInputSource, ...] - flat_inputs: tuple[Any, ...] + tree: ResolvedTree def read( self, env: ManagerEnv, env_ids: np.ndarray | None = None, - ) -> tuple[Any, ...]: + ) -> None: with env.perf.scope("sim_read"): self.sim_data.execute(env_ids) - return self.flat_inputs - - -@dataclass(frozen=True) -class KernelInputSource: - """Compile-time description of one manager-kernel input source. - - Manager inputs are flattened and validated once while the Numba program is - compiled. ``values`` then keeps those flattened leaves by reference so - runtime reads only refresh simulator-backed arrays and can return the same - tuple without rebuilding the input tree. ``source_name`` is a stable, - human-readable label used in manager-layout diagnostics and validation - errors; it is intentionally independent of a source implementation class. - - Attributes: - source_name: Stable semantic label for layout entries and validation - messages (for example, ``"manager_context"`` or ``"sim_reset"``). - value_type: Original structured input type reconstructed by the kernel - compiler when preparing a manager context or reset term. - layout: Lowered tree layout describing the order, scope, and types of - the flattened input leaves. - values: Flattened leaves captured at compile time. Array leaves are - retained by reference and are refreshed in place by their owning - runtime objects; scalar/shared leaves are reused unchanged. - """ - - source_name: str - value_type: type[Any] - layout: KernelDataLayout - values: tuple[Any, ...] - - @property - def fields(self) -> tuple[KernelLeafLayout, ...]: - return iter_layout_leaves(self.layout) - - @property - def field_names(self) -> tuple[str, ...]: - return tuple(".".join(field.path) or "data" for field in self.fields) - - @classmethod - def prepare( - cls, - source_name: str, - value_type: type[Any], - layout: KernelDataLayout, - value: Any, - ) -> KernelInputSource: - """Flatten and validate one static manager-kernel input at compile time.""" - values, tree_def = flatten_kernel_data(value) - if tree_def.fingerprint != layout.tree_def.fingerprint: - raise TypeError(f"{source_name} value has a different KernelData schema.") - fields = iter_layout_leaves(layout) - for prepared_field, field_value in zip(fields, values, strict=True): - if prepared_field.value_type is np.ndarray and not isinstance(field_value, np.ndarray): - raise TypeError(f"{source_name} field {'.'.join(prepared_field.path)!r} must be np.ndarray.") - if prepared_field.value_type is not np.ndarray and type(field_value) is not prepared_field.value_type: - raise TypeError( - f"{source_name} field {'.'.join(prepared_field.path)!r} must return " - f"{prepared_field.value_type.__name__}, got {type(field_value).__name__}." - ) - return cls(source_name, value_type, layout, values) @dataclass(frozen=True) @@ -492,7 +435,7 @@ def init_state(self) -> ArrayEnvState: final_obs = self._zeros_obs(self.observation_space) self._final_obs_slot = _ObsBufferSlot.bind(final_obs, state.reward, state.terminated) self._refresh_sim_reads() - inputs = self._kernel_inputs + inputs = self._resolved_inputs self._validate_kernel_context(inputs) warmup_started = perf_counter() self._execute_observe_kernel(inputs) @@ -511,7 +454,7 @@ def _prev_physics_step(self) -> None: obs, self._kernel_outputs = self._observation_buffer.swap() self._state = self._state.replace(obs=obs) - def _execute_evaluate_kernel(self, inputs: tuple[Any, ...]) -> None: + def _execute_evaluate_kernel(self, inputs: tuple[ResolvedTree, ...]) -> None: """Evaluate rewards and terminations against the current command.""" assert self._task_program is not None assert self._kernel_buffers is not None @@ -523,7 +466,7 @@ def _execute_evaluate_kernel(self, inputs: tuple[Any, ...]) -> None: self._kernel_outputs, ) - def _execute_observe_kernel(self, inputs: tuple[Any, ...]) -> None: + def _execute_observe_kernel(self, inputs: tuple[ResolvedTree, ...]) -> None: """Evaluate observations against the command for the next action.""" assert self._task_program is not None assert self._kernel_outputs is not None @@ -538,7 +481,7 @@ def _capture_final_obs(self, env_ids: np.ndarray) -> FinalObs: assert self._task_program is not None assert self._final_obs_slot is not None slot = self._final_obs_slot - self._task_program.observe_kernel(self._kernel_inputs, slot.kernel_outputs) + self._task_program.observe_kernel(self._resolved_inputs, slot.kernel_outputs) obs = slot.observation return FinalObs( env_ids=env_ids, @@ -649,7 +592,7 @@ def apply_action(self, actions: np.ndarray, state: ArrayEnvState) -> ArrayEnvSta self._action_writes.execute() return state - def _reset_sim_rows(self, env_ids: np.ndarray, inputs: tuple[Any, ...]) -> None: + def _reset_sim_rows(self, env_ids: np.ndarray, inputs: tuple[ResolvedTree, ...]) -> None: """Run the simulator-only stages of the reset pipeline for selected rows. Clears frontend ctrl targets to mirror the backend's reset-row ctrl @@ -721,7 +664,7 @@ def reset(self, env_ids: np.ndarray, sim_reset_ids: np.ndarray | None = None) -> if sim_reset_ids is not None and sim_reset_ids.size: env_ids = np.concatenate([env_ids, sim_reset_ids]) if env_ids.size: - self._reset_sim_rows(env_ids, self._kernel_inputs) + self._reset_sim_rows(env_ids, self._resolved_inputs) def _make_metrics_view(self) -> dict[str, Any]: """Assemble the persistent live metrics view for the current state. @@ -751,7 +694,7 @@ def compute_transition(self, state: ArrayEnvState) -> ArrayEnvState: with self.perf.scope("read_inputs"): self._refresh_sim_reads() with self.perf.scope("evaluate"): - self._execute_evaluate_kernel(self._kernel_inputs) + self._execute_evaluate_kernel(self._resolved_inputs) with self.perf.scope("command_on_transition"): for command_term in self._command_terms.values(): command_term.on_transition() @@ -764,7 +707,7 @@ def compute_observation(self, state: ArrayEnvState) -> ArrayEnvState: if self._kernel_outputs is None: return state with self.perf.scope("observe"): - self._execute_observe_kernel(self._kernel_inputs) + self._execute_observe_kernel(self._resolved_inputs) return state @property @@ -810,15 +753,15 @@ def _make_kernel_buffers(self, state: ArrayEnvState) -> tuple[np.ndarray, ...]: return buffers @property - def _kernel_inputs(self) -> tuple[Any, ...]: - """Persistent flat kernel-input bundle shared by all three kernels. + def _resolved_inputs(self) -> tuple[ResolvedTree, ...]: + """Persistent resolved-input trees shared by all three kernels. - The tuple is built once by the read plan. Simulator-backed arrays are - rewritten in place by :meth:`_refresh_sim_reads`; every other entry - aliases host state, so the bundle is never rebuilt. + The trees are built once by the resolver. Simulator-backed array + leaves are rewritten in place by :meth:`_refresh_sim_reads`; every + other leaf aliases host state, so the trees are never rebuilt. """ assert self._compiled_manager_program is not None - return self._compiled_manager_program.read_plan.flat_inputs + return self._compiled_manager_program.prepared_trees def _refresh_sim_reads(self, env_ids: np.ndarray | None = None) -> None: """Refresh the simulator-backed kernel inputs in place. @@ -829,7 +772,7 @@ def _refresh_sim_reads(self, env_ids: np.ndarray | None = None) -> None: assert self._compiled_manager_program is not None self._compiled_manager_program.read_plan.read(self, env_ids) - def _validate_kernel_context(self, inputs: tuple[Any, ...]) -> None: + def _validate_kernel_context(self, inputs: tuple[ResolvedTree, ...]) -> None: assert self._task_program is not None assert self._kernel_buffers is not None assert self._kernel_outputs is not None @@ -840,20 +783,11 @@ def _validate_kernel_context(self, inputs: tuple[Any, ...]) -> None: self._kernel_buffers, self._kernel_outputs, self.num_envs, - self._kernel_input_batch_axes(), ) except (TypeError, ValueError) as error: message = f"Invalid kernel context for {type(self).__name__} task kernel: {error}" raise type(error)(message) from error - def _kernel_input_batch_axes(self) -> tuple[bool, ...]: - assert self._compiled_manager_program is not None - return tuple( - field.scope is KernelDataScope.PER_ENV - for source in self._compiled_manager_program.read_plan.sources - for field in source.fields - ) - def warmup(self) -> ManagerWarmupResult: """Run one isolated evaluate/observe pass to exercise the executor graphs. @@ -870,18 +804,19 @@ def warmup(self) -> ManagerWarmupResult: assert self._kernel_outputs is not None self._refresh_sim_reads() - inputs = self._kernel_inputs + inputs = self._resolved_inputs self._validate_kernel_context(inputs) import torch started = perf_counter() + warmup_inputs = clone_resolved_trees(inputs) self._task_program.evaluate_kernel( - clone_kernel_value(inputs), + warmup_inputs, self._task_program.reward_weights, clone_kernel_value(self._kernel_buffers), clone_kernel_value(self._kernel_outputs), ) - self._task_program.observe_kernel(clone_kernel_value(inputs), clone_kernel_value(self._kernel_outputs)) + self._task_program.observe_kernel(warmup_inputs, clone_kernel_value(self._kernel_outputs)) return ManagerWarmupResult( compile_seconds=0.0, signatures=("torch eager",), diff --git a/motrix_env_core/src/motrix_env_core/manager/kernel.py b/motrix_env_core/src/motrix_env_core/manager/kernel.py index 0897a183..471fdf92 100644 --- a/motrix_env_core/src/motrix_env_core/manager/kernel.py +++ b/motrix_env_core/src/motrix_env_core/manager/kernel.py @@ -6,6 +6,9 @@ import numpy as np +from motrix_env_core.manager.kernel_data.lowering import KernelDataScope +from motrix_env_core.manager.kernel_data.resolved import ResolvedLeaf, ResolvedTree + T = TypeVar("T") @@ -32,28 +35,51 @@ class ManagerWarmupResult: def validate_kernel_context( - inputs: Any, + trees: Any, reward_weights: np.ndarray, buffers: Any, outputs: KernelOutputs, num_envs: int, - input_batch_axes: tuple[bool, ...] | None = None, ) -> None: - """Validate the numeric kernel contract once when a context is bound.""" - if input_batch_axes is None: - _validate_value(inputs, "inputs", num_envs=num_envs, batched=True, writable=False) - else: - if not isinstance(inputs, tuple): - raise TypeError(f"inputs must be tuple, got {type(inputs).__name__}.") - if len(inputs) != len(input_batch_axes): - raise ValueError(f"inputs must contain {len(input_batch_axes)} values, got {len(inputs)}.") - for index, (value, batched) in enumerate(zip(inputs, input_batch_axes, strict=True)): - _validate_value(value, f"inputs.{index}", num_envs=num_envs, batched=batched, writable=False) + """Validate the numeric kernel contract once when a context is bound. + + ``trees`` is the resolved program input: one :class:`ResolvedTree` per + prepared kernel-data source. Each leaf is validated against its declared + scope (``PER_ENV`` leaves must lead with the batch dimension). + """ + if not isinstance(trees, tuple): + raise TypeError(f"inputs must be a tuple of resolved trees, got {type(trees).__name__}.") + for tree in trees: + for path, leaf in tree.leaves.items(): + _validate_value( + leaf.value, + f"inputs.{'.'.join(path)}", + num_envs=num_envs, + batched=leaf.scope is KernelDataScope.PER_ENV, + writable=False, + ) _validate_value(reward_weights, "reward_weights", num_envs=num_envs, batched=False, writable=False) _validate_value(buffers, "buffers", num_envs=num_envs, batched=True, writable=True) _validate_outputs(outputs, num_envs) +def clone_resolved_trees(trees: tuple[ResolvedTree, ...]) -> tuple[ResolvedTree, ...]: + """Clone every leaf of the resolved trees for isolated warmup execution. + + Leaves are cloned with :func:`clone_kernel_value` semantics (dtype, + strides, and read-only flag preserved); the host root objects stay shared + exactly as in real execution, where namespace fallback reads immutable + config attributes from the original records. + """ + return tuple( + ResolvedTree( + tree.value, + {path: ResolvedLeaf(clone_kernel_value(leaf.value), leaf.scope) for path, leaf in tree.leaves.items()}, + ) + for tree in trees + ) + + def clone_kernel_value(value: T) -> T: """Recursively clone arrays while preserving task-specific tuple types.""" if isinstance(value, np.ndarray): diff --git a/motrix_env_core/src/motrix_env_core/manager/kernel_data/__init__.py b/motrix_env_core/src/motrix_env_core/manager/kernel_data/__init__.py index 8bc4666b..ae6780b7 100644 --- a/motrix_env_core/src/motrix_env_core/manager/kernel_data/__init__.py +++ b/motrix_env_core/src/motrix_env_core/manager/kernel_data/__init__.py @@ -1,5 +1,4 @@ # Copyright Motphys Technology Co., Ltd. 2025, 2026 -# SPDX-License-Identifier: Apache-2.0 from motrix_env_core.manager.kernel_data.canonical import canonicalize_kernel_data from motrix_env_core.manager.kernel_data.lowering import ( @@ -11,12 +10,9 @@ KernelRecordLayout, SharedArray, iter_layout_leaves, - lane_expression, - proxy_symbol, - proxy_types, - rebuild_lowered, ) -from motrix_env_core.manager.kernel_data.map import Map, map_keys, map_proxy +from motrix_env_core.manager.kernel_data.map import Map +from motrix_env_core.manager.kernel_data.resolved import ResolvedLeaf, ResolvedTree from motrix_env_core.manager.kernel_data.tree import ( BaseDef, FieldDef, @@ -47,6 +43,8 @@ "KernelRecordLayout", "TreeClassDef", "SharedArray", + "ResolvedLeaf", + "ResolvedTree", "canonicalize_kernel_data", "flatten_kernel_data", "is_kernel_data", @@ -54,11 +52,5 @@ "iter_layout_leaves", "iter_tree_leaves", "kernel_data", - "map_keys", - "map_proxy", - "lane_expression", - "proxy_symbol", - "proxy_types", - "rebuild_lowered", "unflatten_kernel_data", ] diff --git a/motrix_env_core/src/motrix_env_core/manager/kernel_data/lowering.py b/motrix_env_core/src/motrix_env_core/manager/kernel_data/lowering.py index d8adc007..bfa2ade6 100644 --- a/motrix_env_core/src/motrix_env_core/manager/kernel_data/lowering.py +++ b/motrix_env_core/src/motrix_env_core/manager/kernel_data/lowering.py @@ -4,19 +4,17 @@ from __future__ import annotations import hashlib -import sys from dataclasses import dataclass from enum import Enum -from typing import Annotated, Any, NamedTuple, TypeAlias, get_args, get_origin +from typing import Annotated, Any, TypeAlias, get_args, get_origin import numpy as np from motrix_env_core.manager.fingerprint import function_fingerprint -from motrix_env_core.manager.kernel_data.map import map_proxy from motrix_env_core.manager.kernel_data.methods import dispatch_methods from motrix_env_core.manager.kernel_data.tree import KernelMapDef, LeafDef, TreeClassDef -_LOWERING_SCHEMA_VERSION = 2 +_LOWERING_SCHEMA_VERSION = 3 _LAYOUT_CACHE: dict[ tuple[str, bool], KernelRecordLayout, @@ -24,11 +22,12 @@ class KernelDataScope(str, Enum): - """How one flattened KernelData leaf is exposed to a Numba execution lane. + """How one resolved KernelData leaf is exposed to execution. Attributes: - PER_ENV: Select ``leaf[env_id]`` while reconstructing each execution lane. - SHARED: Pass the complete leaf unchanged to every execution lane. + PER_ENV: Leaf is an ``(N, ...)`` array covering the whole batch; the + reset schedule clones it for selected lanes. + SHARED: Pass the complete leaf unchanged to every environment lane. """ PER_ENV = "per_env" @@ -45,9 +44,9 @@ class KernelLeafLayout: Attributes: tree_def: Logical leaf definition that owns the path and annotation. - value_type: Concrete Python/NumPy type accepted by the kernel input slot. - scope: Whether lane reconstruction indexes by environment or shares the leaf. - slot_index: Zero-based slot relative to the root layout's flattened inputs. + value_type: Concrete Python/NumPy type accepted for the leaf value. + scope: Whether the leaf is per-environment batched or shared. + slot_index: Zero-based declaration index relative to the root layout. """ tree_def: LeafDef @@ -93,13 +92,11 @@ class KernelMapLayout: Attributes: tree_def: Value-derived logical map definition. - lowered_type: Schema-specific NamedTuple proxy constructed in Numba lanes. entries: Lowered children in the same order as the logical map entries. - fingerprint: Deterministic identity including child types, scopes, and slots. + fingerprint: Deterministic identity including keys and child schemas. """ tree_def: KernelMapDef - lowered_type: type[tuple[Any, ...]] entries: tuple[KernelMapEntryLayout, ...] fingerprint: str @@ -110,13 +107,12 @@ class KernelRecordLayout: Attributes: tree_def: Logical record definition being lowered. - lowered_type: Schema-specific NamedTuple proxy constructed in Numba lanes. fields: Lowered children in logical dataclass declaration order. - fingerprint: Deterministic identity including child types, scopes, and slots. + fingerprint: Deterministic identity including child types, scopes, + and dispatch-method code fingerprints. """ tree_def: TreeClassDef - lowered_type: type[tuple[Any, ...]] fields: tuple[KernelFieldLayout, ...] fingerprint: str @@ -139,60 +135,6 @@ def iter_layout_leaves(layout: KernelDataLayout) -> tuple[KernelLeafLayout, ...] return tuple(leaf for field in layout.fields for leaf in iter_layout_leaves(field.child)) -def rebuild_lowered(layout: KernelDataLayout, leaves: tuple[Any, ...]) -> Any: - """Rebuild compiler-owned lowered proxies from ordered leaves.""" - value, next_index = _rebuild_lowered(layout, leaves, 0) - if next_index != len(leaves): - raise ValueError("KernelData lowered rebuild received extra leaves.") - return value - - -def _rebuild_lowered(layout: KernelDataLayout, leaves: tuple[Any, ...], index: int) -> tuple[Any, int]: - if isinstance(layout, KernelLeafLayout): - if index >= len(leaves): - raise ValueError("KernelData lowered rebuild received too few leaves.") - return leaves[index], index + 1 - if isinstance(layout, KernelMapLayout): - values = [] - for entry in layout.entries: - value, index = _rebuild_lowered(entry.child, leaves, index) - values.append(value) - return layout.lowered_type(*values), index - assert isinstance(layout, KernelRecordLayout) - values = [] - for field in layout.fields: - value, index = _rebuild_lowered(field.child, leaves, index) - values.append(value) - return layout.lowered_type(*values), index - - -def lane_expression(layout: KernelDataLayout, input_offset: int, env_symbol: str = "env_id") -> str: - """Generate a bottom-up lane-local reconstruction expression.""" - if isinstance(layout, KernelLeafLayout): - input_symbol = f"input_{input_offset + layout.slot_index}" - return f"{input_symbol}[{env_symbol}]" if layout.scope is KernelDataScope.PER_ENV else input_symbol - if isinstance(layout, KernelMapLayout): - arguments = ", ".join(lane_expression(entry.child, input_offset, env_symbol) for entry in layout.entries) - return f"{proxy_symbol(layout.lowered_type)}({arguments})" - assert isinstance(layout, KernelRecordLayout) - arguments = ", ".join(lane_expression(field.child, input_offset, env_symbol) for field in layout.fields) - return f"{proxy_symbol(layout.lowered_type)}({arguments})" - - -def proxy_types(layout: KernelDataLayout) -> tuple[type[tuple[Any, ...]], ...]: - """Return all lowered record proxy types in child-first order.""" - if isinstance(layout, KernelLeafLayout): - return () - if isinstance(layout, KernelMapLayout): - return tuple(proxy for entry in layout.entries for proxy in proxy_types(entry.child)) + (layout.lowered_type,) - assert isinstance(layout, KernelRecordLayout) - return tuple(proxy for field in layout.fields for proxy in proxy_types(field.child)) + (layout.lowered_type,) - - -def proxy_symbol(proxy: type[tuple[Any, ...]]) -> str: - return proxy.__name__ - - def _tree_method_fingerprint(tree_def: TreeClassDef | KernelMapDef | LeafDef) -> Any: """Include nested records and map values before consulting the layout cache.""" if isinstance(tree_def, LeafDef): @@ -209,7 +151,13 @@ def _tree_method_fingerprint(tree_def: TreeClassDef | KernelMapDef | LeafDef) -> class KernelDataLowering: - """Lower one logical ``TreeClassDef`` into a fixed flat Numba ABI layout.""" + """Validate one logical ``TreeClassDef`` and derive its leaf layout. + + The lowering checks that every leaf annotation is a supported + ndarray/scalar type, assigns each leaf its data scope, and produces a + deterministic structural fingerprint (field names, types, scopes, and + dispatch-method code) used for resolved-program invalidation. + """ def lower( self, @@ -217,7 +165,7 @@ def lower( *, context: str, force_shared: bool = False, - ) -> KernelDataLayout: + ) -> KernelRecordLayout: method_key = hashlib.sha256(repr(_tree_method_fingerprint(tree_def)).encode()).hexdigest() cache_key = (f"{tree_def.fingerprint}:{method_key}", force_shared) cached = _LAYOUT_CACHE.get(cache_key) @@ -251,7 +199,6 @@ def _lower_record( slot_index, force_shared, context, - tree_def.logical_type.__module__, ) else: assert isinstance(field.tree_def, LeafDef) @@ -259,7 +206,6 @@ def _lower_record( field.tree_def.annotation, context=f"{context} KernelData {tree_def.logical_type.__name__}.{'.'.join(field.tree_def.path)}", force_shared=force_shared, - force_per_env=False, ) if value_type is None: raise TypeError( @@ -280,12 +226,7 @@ def _lower_record( (_LOWERING_SCHEMA_VERSION, tree_def.fingerprint, identity, tuple(fingerprint_fields), method_fingerprints) ) fingerprint = hashlib.sha256(raw_fingerprint.encode()).hexdigest() - lowered_type = self._lowered_proxy( - tree_def.logical_type, - tuple(field.name for field in field_layouts), - fingerprint, - ) - return KernelRecordLayout(tree_def, lowered_type, tuple(field_layouts), fingerprint), slot_index + return KernelRecordLayout(tree_def, tuple(field_layouts), fingerprint), slot_index def _lower_map( self, @@ -293,7 +234,6 @@ def _lower_map( slot_index: int, force_shared: bool, context: str, - module_name: str, ) -> tuple[KernelMapLayout, int]: entry_layouts: list[KernelMapEntryLayout] = [] fingerprint_entries: list[Any] = [] @@ -311,7 +251,6 @@ def _lower_map( slot_index, force_shared, context, - module_name, ) else: assert isinstance(entry.tree_def, LeafDef) @@ -319,7 +258,6 @@ def _lower_map( entry.tree_def.annotation, context=f"{context} Map {'.'.join(entry.tree_def.path)}", force_shared=force_shared, - force_per_env=False, ) if value_type is None: raise TypeError( @@ -332,14 +270,7 @@ def _lower_map( fingerprint_entries.append((entry.key, self._fingerprint_part(child))) raw_fingerprint = repr((_LOWERING_SCHEMA_VERSION, tree_def.fingerprint, tuple(fingerprint_entries))) fingerprint = hashlib.sha256(raw_fingerprint.encode()).hexdigest() - namespace = ".".join(tree_def.path) or "root" - lowered_type = map_proxy( - namespace, - tuple(entry.key for entry in entry_layouts), - module_name=module_name, - schema_fingerprint=fingerprint, - ) - return KernelMapLayout(tree_def, lowered_type, tuple(entry_layouts), fingerprint), slot_index + return KernelMapLayout(tree_def, tuple(entry_layouts), fingerprint), slot_index @staticmethod def _fingerprint_part(layout: KernelDataLayout) -> Any: @@ -366,41 +297,12 @@ def _fingerprint_part(layout: KernelDataLayout) -> Any: tuple((field.name, KernelDataLowering._fingerprint_part(field.child)) for field in layout.fields), ) - @staticmethod - def _lowered_proxy( - logical_type: type, - field_names: tuple[str, ...], - fingerprint: str, - ) -> type[tuple[Any, ...]]: - sanitized = "_".join( - part for part in logical_type.__qualname__.replace("", "locals").split(".") if part - ) - name = f"_{sanitized}KernelData_{fingerprint[:12]}" - module = sys.modules.get(logical_type.__module__) - if module is None: - raise TypeError(f"KernelData module {logical_type.__module__!r} is not loaded.") - existing = getattr(module, name, None) - if existing is not None: - if ( - not isinstance(existing, type) - or not issubclass(existing, tuple) - or getattr(existing, "_fields", None) != field_names - ): - raise TypeError(f"KernelData lowered proxy name collision: {logical_type.__module__}.{name}.") - return existing - lowered = NamedTuple(name, [(field_name, Any) for field_name in field_names]) # type: ignore[misc] - lowered.__module__ = logical_type.__module__ - lowered.__qualname__ = name - setattr(module, name, lowered) - return lowered - @staticmethod def _leaf_type( annotation: Any, *, context: str, force_shared: bool, - force_per_env: bool = False, ) -> tuple[type[Any] | None, KernelDataScope]: metadata: tuple[Any, ...] = () if get_origin(annotation) is Annotated: @@ -409,13 +311,7 @@ def _leaf_type( if metadata and metadata != (KernelDataScope.SHARED,): raise TypeError(f"{context} has unsupported or conflicting metadata {metadata!r}.") if annotation is np.ndarray: - if metadata and force_per_env: - raise TypeError(f"{context} cannot be both SharedArray and explicitly per-environment.") - scope = ( - KernelDataScope.PER_ENV - if force_per_env - else (KernelDataScope.SHARED if metadata or force_shared else KernelDataScope.PER_ENV) - ) + scope = KernelDataScope.SHARED if metadata or force_shared else KernelDataScope.PER_ENV return annotation, scope if metadata: raise TypeError(f"{context} scope metadata is only valid on np.ndarray leaves.") diff --git a/motrix_env_core/src/motrix_env_core/manager/kernel_data/map.py b/motrix_env_core/src/motrix_env_core/manager/kernel_data/map.py index 1b23402f..9e278e0a 100644 --- a/motrix_env_core/src/motrix_env_core/manager/kernel_data/map.py +++ b/motrix_env_core/src/motrix_env_core/manager/kernel_data/map.py @@ -1,15 +1,9 @@ # Copyright Motphys Technology Co., Ltd. 2025, 2026 -# SPDX-License-Identifier: Apache-2.0 from __future__ import annotations -import hashlib -import sys from collections.abc import Iterator, Mapping -from typing import Any, Generic, NamedTuple, TypeVar - -_MAP_KEYS = "__motrix_kernel_data_map_keys__" -_MAP_SCHEMA = "__motrix_kernel_data_map_schema__" +from typing import Generic, TypeVar ValueT = TypeVar("ValueT") @@ -43,50 +37,4 @@ def values_tuple(self) -> tuple[ValueT, ...]: return self._values -def map_proxy( - namespace: str, - keys: tuple[str, ...], - *, - module_name: str, - schema_fingerprint: str = "", -) -> type[tuple[Any, ...]]: - """Return a stable lowered NamedTuple type for one static map schema.""" - identity = (module_name, namespace, keys, schema_fingerprint) - fingerprint = hashlib.sha256(repr(identity).encode()).hexdigest()[:12] - sanitized = "_".join(part for part in namespace.replace(".", "_").split("_") if part) - name = f"_{sanitized}Map_{fingerprint}" - module = sys.modules.get(module_name) - if module is None: - raise TypeError(f"Map proxy module {module_name!r} is not loaded.") - existing = getattr(module, name, None) - field_names = tuple(f"value_{index}" for index in range(len(keys))) - if existing is not None: - if ( - not isinstance(existing, type) - or not issubclass(existing, tuple) - or getattr(existing, "_fields", None) != field_names - or getattr(existing, _MAP_KEYS, None) != keys - or getattr(existing, _MAP_SCHEMA, None) != schema_fingerprint - ): - raise TypeError(f"Map lowered proxy name collision: {module_name}.{name}.") - return existing - proxy = NamedTuple(name, [(field_name, Any) for field_name in field_names]) # type: ignore[misc] - proxy.__module__ = module_name - proxy.__qualname__ = name - setattr(proxy, _MAP_KEYS, keys) - setattr(proxy, _MAP_SCHEMA, schema_fingerprint) - setattr(module, name, proxy) - return proxy - - -def map_keys(value_type: type[Any]) -> tuple[str, ...] | None: - """Return static map keys for a lowered proxy type.""" - keys = getattr(value_type, _MAP_KEYS, None) - return keys if isinstance(keys, tuple) else None - - -__all__ = [ - "Map", - "map_keys", - "map_proxy", -] +__all__ = ["Map"] diff --git a/motrix_env_core/src/motrix_env_core/manager/kernel_data/resolved.py b/motrix_env_core/src/motrix_env_core/manager/kernel_data/resolved.py new file mode 100644 index 00000000..e3b29535 --- /dev/null +++ b/motrix_env_core/src/motrix_env_core/manager/kernel_data/resolved.py @@ -0,0 +1,43 @@ +# Copyright Motphys Technology Co., Ltd. 2025, 2026 +# SPDX-License-Identifier: Apache-2.0 + +"""Resolved tree values consumed by manager executors.""" + +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import dataclass +from typing import Any + +from motrix_env_core.manager.kernel_data.lowering import KernelDataScope + + +@dataclass(frozen=True) +class ResolvedLeaf: + """One resolved leaf of a kernel-data tree. + + Attributes: + value: Live leaf value reference — an ndarray held and refreshed in + place by its owning runtime object, or an immutable scalar. + scope: Whether the leaf is an ``(N, ...)`` batched array or shared. + """ + + value: Any + scope: KernelDataScope + + +@dataclass(frozen=True) +class ResolvedTree: + """Path-keyed resolution of one kernel-data root value. + + ``value`` keeps the host root object (used as the attribute fallback for + executor namespace views); ``leaves`` maps every leaf path to its live + value reference and scope. Array leaves are rewritten in place by their + owning runtime objects, so the tree itself never needs rebuilding. + """ + + value: Any + leaves: Mapping[tuple[str, ...], ResolvedLeaf] + + +__all__ = ["ResolvedLeaf", "ResolvedTree"] diff --git a/motrix_env_core/src/motrix_env_core/manager/resolver.py b/motrix_env_core/src/motrix_env_core/manager/resolver.py index ccbbba23..556ebbbf 100644 --- a/motrix_env_core/src/motrix_env_core/manager/resolver.py +++ b/motrix_env_core/src/motrix_env_core/manager/resolver.py @@ -4,16 +4,18 @@ """Backend-neutral resolver for manager task programs. Lowers one constructed manager environment into its execution metadata: the -simulator read plan, the flat kernel-input layout, prepared term arguments, -and one invocation per manager term (command hooks, observations, rewards, -terminations, simulator resets). Numerical executors — currently the Torch -executor in :mod:`motrix_env_core.torch_manager` — consume this metadata and -schedule the invocations; no code generation or dispatcher compilation -happens here. +simulator read plan, the resolved context/term-argument trees (path-keyed +leaf descriptors with live value references and data scopes), prepared term +arguments, and one invocation per manager term (command hooks, observations, +rewards, terminations, simulator resets). Numerical executors — currently the +Torch executor in :mod:`motrix_env_core.torch_manager` — consume this +metadata and schedule the invocations; no code generation or dispatcher +compilation happens here. """ from __future__ import annotations +import hashlib import inspect import logging from collections.abc import Mapping @@ -26,22 +28,17 @@ import torch from motrix_env_core.manager.context import ManagerContext -from motrix_env_core.manager.env import ( - KernelInputSource, - ManagerEnv, - ReadPlan, -) +from motrix_env_core.manager.env import ManagerEnv, ReadPlan from motrix_env_core.manager.kernel_data import ( KernelDataLayout, KernelDataLowering, KernelDataScope, Map, - SharedArray, + ResolvedLeaf, + ResolvedTree, flatten_kernel_data, is_kernel_data, iter_layout_leaves, - kernel_data, - lane_expression, ) from motrix_env_core.manager.observations import ObservationGroupEntry, ObsTerm from motrix_env_core.manager.rewards import RewardTerm @@ -56,16 +53,8 @@ # --------------------------------------------------------------------------- -@dataclass(frozen=True) -class InputSlotLayout: - index: int - source: str - scope: KernelDataScope - - @dataclass(frozen=True) class SimInputLayout: - slot: int key: str shape: tuple[int, ...] consumers: tuple[str, ...] @@ -106,7 +95,6 @@ class MetricFieldLayout: @dataclass(frozen=True) class ManagerLayout: - inputs: tuple[InputSlotLayout, ...] sim_inputs: tuple[SimInputLayout, ...] actions: tuple[ActionTermLayout, ...] observations: dict[str, ObservationGroupLayout] @@ -116,12 +104,10 @@ class ManagerLayout: def dump(self) -> str: lines = [ - "inputs:", + "sim_inputs:", ] - lines.extend(f" [{slot.index}] scope={slot.scope.value} source={slot.source}" for slot in self.inputs) - lines.append("sim_inputs:") lines.extend( - f" [{sim_input.slot}] {sim_input.key}: shape={sim_input.shape} consumers={sim_input.consumers}" + f" {sim_input.key}: shape={sim_input.shape} consumers={sim_input.consumers}" for sim_input in self.sim_inputs ) lines.append("actions:") @@ -144,6 +130,42 @@ def dump(self) -> str: return "\n".join(lines) +# --------------------------------------------------------------------------- +# Prepared arguments +# --------------------------------------------------------------------------- + + +@dataclass(frozen=True) +class ConstantArg: + """Literal term argument passed through unchanged (strings, tensors...).""" + + value: Any + + +@dataclass(frozen=True) +class TreeArg: + """Term argument resolved as a structured kernel-data tree.""" + + tree_index: int + + +@dataclass(frozen=True) +class LeafArg: + """Term argument resolved as one leaf value (numeric/ndarray argument).""" + + leaf: ResolvedLeaf + + +@dataclass(frozen=True) +class SimArg: + """Term argument resolving to one sim read-program batched array.""" + + key: str + + +PreparedArg = ConstantArg | TreeArg | LeafArg | SimArg + + @dataclass(frozen=True) class PreparedInvocation: """Runtime metadata required to execute one manager term.""" @@ -152,16 +174,13 @@ class PreparedInvocation: kind: str dispatcher: Any receiver_key: str | None - args_expressions: tuple[str, ...] = () - args_values: tuple[Any, ...] = () - args_prepared_indices: tuple[int | None, ...] = () + args: tuple[PreparedArg, ...] = () output_size: int = 0 @dataclass(frozen=True) class ResolvedManagerContext: - prepared_index: int - expression: str + tree_index: int @dataclass(frozen=True) @@ -179,57 +198,37 @@ class ResolvedManagerProgram: read_plan: ReadPlan layout: ManagerLayout invocations: tuple[PreparedInvocation, ...] - prepared_terms: tuple[KernelInputSource, ...] - input_offsets: tuple[int, ...] + prepared_trees: tuple[ResolvedTree, ...] context: ResolvedManagerContext + fingerprint: str __all__ = [ "ActionTermLayout", - "InputSlotLayout", + "ConstantArg", + "LeafArg", "ManagerLayout", "ManagerProgramResolver", "MetricFieldLayout", "ObservationGroupLayout", "ObservationTermLayout", + "PreparedArg", "PreparedInvocation", "ResolvedManagerContext", "ResolvedManagerProgram", "ResolvedSimReset", "ScalarTermLayout", + "SimArg", "SimInputLayout", + "TreeArg", ] # --------------------------------------------------------------------------- -# Term-argument buffer wrappers +# Resolver # --------------------------------------------------------------------------- -@kernel_data -class TermArrayBuffer: - """Shared array slot for one ndarray term argument. - - Arrays passed directly as manager term arguments are lowered with shared - scope (every execution lane sees the full array): they hold env-invariant - reference data such as default poses, weights, or lookup tables. - """ - - value: SharedArray - - -@kernel_data -class TermScalarBuffer: - """Runtime buffer slot for one numeric term argument. - - Numeric tuning arguments (thresholds, scales, gains, ...) are lowered - into a fixed-dtype shared input slot so executors read the value at - runtime instead of baking it into a compiled artifact. - """ - - value: np.float32 - - @dataclass(frozen=True) class ManagerSimInput: """One compiled simulator input exposed through ``ManagerContext.sim``.""" @@ -246,9 +245,18 @@ def _manager_sim_inputs(sim_data: PhysicsReadProgram) -> tuple[Mapping[str, Mana return MappingProxyType(inputs), fingerprint -# --------------------------------------------------------------------------- -# Resolver -# --------------------------------------------------------------------------- +def _tree_leaf_fingerprint(tree: ResolvedTree) -> tuple[Any, ...]: + """Per-leaf content fingerprint: scope plus array shape/dtype identity.""" + return tuple( + ( + path, + leaf.scope.value, + None + if not isinstance(leaf.value, np.ndarray) + else (leaf.value.shape, leaf.value.dtype.str, leaf.value.strides), + ) + for path, leaf in tree.leaves.items() + ) class ManagerProgramResolver: @@ -258,10 +266,7 @@ def __init__(self, env: ManagerEnv): self._env = env self._sim_inputs: Mapping[str, ManagerSimInput] self._sim_inputs, sim_input_fingerprint = _manager_sim_inputs(env.sim_data) - self._sim_slots_by_key: dict[str, int] = {} - self._prepared_terms: list[KernelInputSource] = [] - self._input_offsets: list[int] = [] - self._flat_input_count = 0 + self._prepared_trees: list[ResolvedTree] = [] self._kernel_data_lowering = KernelDataLowering() self._layout_identity: list[Any] = [sim_input_fingerprint] @@ -287,16 +292,14 @@ def build(self) -> ResolvedManagerProgram: self._validate_observation_spaces(observation_groups) layout = ManagerLayout( - inputs=tuple( - InputSlotLayout( - offset + field_index, - f"{prepared.source_name}.{field_name}", - prepared.fields[field_index].scope, + sim_inputs=tuple( + SimInputLayout( + key=key, + shape=binding.value.shape[1:], + consumers=("ManagerContext.sim",), ) - for prepared, offset in zip(self._prepared_terms, self._input_offsets, strict=True) - for field_index, field_name in enumerate(prepared.field_names) + for key, binding in self._sim_inputs.items() ), - sim_inputs=self._sim_input_layout(), actions=tuple( ActionTermLayout( name, @@ -321,24 +324,22 @@ def build(self) -> ResolvedManagerProgram: + termination_invocations + tuple(reset.invocation for reset in reset_invocations) ) + leaf_count = sum(len(tree.leaves) for tree in self._prepared_trees) logger.info( - "Manager startup %s: resolve finished in %.3fs (%d invocations, %d inputs)", + "Manager startup %s: resolve finished in %.3fs (%d invocations, %d resolved leaves)", type(self._env).__name__, perf_counter() - build_started, len(invocations), - self._flat_input_count, + leaf_count, ) + context_tree = self._prepared_trees[context.tree_index] return ResolvedManagerProgram( - read_plan=ReadPlan( - self._env.sim_data, - tuple(self._prepared_terms), - tuple(value for prepared in self._prepared_terms for value in prepared.values), - ), + read_plan=ReadPlan(self._env.sim_data, context_tree), layout=layout, invocations=invocations, - prepared_terms=tuple(self._prepared_terms), - input_offsets=tuple(self._input_offsets), + prepared_trees=tuple(self._prepared_trees), context=context, + fingerprint=hashlib.sha256(repr(tuple(self._layout_identity)).encode()).hexdigest(), ) # -- resets ------------------------------------------------------------- @@ -379,42 +380,24 @@ def _resolve_reset( if annotations.get(parameters[0].name) is not ManagerContext: raise TypeError( f"Manager simulator reset term {name!r} dispatch must annotate its first parameter as " - f"ManagerContext (conventionally named ctx)." + "ManagerContext (conventionally named ctx)." ) sim_writes_annotation = annotations.get(parameters[1].name) if get_origin(sim_writes_annotation) is not Map or get_args(sim_writes_annotation) != (torch.Tensor,): raise TypeError( f"Manager simulator reset term {name!r} dispatch must annotate its second parameter as " - f"Map[torch.Tensor] (conventionally named sim_writes)." + "Map[torch.Tensor] (conventionally named sim_writes)." ) dispatcher = function - expressions = [] - prepared_indices = [] - for index, value in enumerate(term.args): - if is_kernel_data(value): - _, tree_def = flatten_kernel_data(value) - layout = self._kernel_data_lowering.lower( - tree_def, - context=f"Simulator reset {name} args[{index}]", - force_shared=True, - ) - prepared_index, expression = self._register_prepared( - f"sim_reset.{name}.args[{index}]", value, type(value), layout - ) - prepared_indices.append(prepared_index) - expressions.append(expression) - else: - expression, _, prepared_index = self._resolve_plain_arg(f"sim_reset.{name}.args[{index}]", value) - prepared_indices.append(prepared_index) - expressions.append(expression) + args = tuple( + self._resolve_term_arg(f"sim_reset.{name}.args[{index}]", value) for index, value in enumerate(term.args) + ) invocation = PreparedInvocation( "sim_reset." + name, "reset", dispatcher, None, - args_expressions=tuple(expressions), - args_values=term.args, - args_prepared_indices=tuple(prepared_indices), + args=args, ) return ResolvedSimReset( invocation, @@ -425,18 +408,6 @@ def _resolve_reset( # -- layouts ------------------------------------------------------------ - def _sim_input_layout(self) -> tuple[SimInputLayout, ...]: - slots_by_key = self._sim_slots_by_key - return tuple( - SimInputLayout( - slot=slots_by_key[key], - key=key, - shape=binding.value.shape[1:], - consumers=("ManagerContext.sim",), - ) - for key, binding in self._sim_inputs.items() - ) - def _resolve_runtime_terms(self) -> None: for name, action_term in self._env.action_terms.items(): _, tree_def = flatten_kernel_data(action_term.state) @@ -488,25 +459,16 @@ def _resolve_manager_context( tree_def, context="ManagerContext", ) - prepared_index, expression = self._register_prepared("manager_context", value, ManagerContext, layout) - self._layout_identity.append(("manager_context", layout.fingerprint)) + tree_index = self._register_tree("manager_context", value, layout) + tree = self._prepared_trees[tree_index] + self._layout_identity.append(("manager_context", layout.fingerprint, _tree_leaf_fingerprint(tree))) self._layout_identity.append( ( "metrics", tuple((name, value.shape, value.dtype.str) for name, value in self._env.metrics.items()), ) ) - # Map each declared sim key to its global input slot inside the - # flattened context, so term args naming a sim key resolve to the - # key's batched view instead of crossing the executor as a string. - context_source = self._prepared_terms[prepared_index] - context_offset = self._input_offsets[prepared_index] - self._sim_slots_by_key = { - field.path[1]: context_offset + field_index - for field_index, field in enumerate(context_source.fields) - if len(field.path) == 2 and field.path[0] == "sim" - } - return ResolvedManagerContext(prepared_index, expression) + return ResolvedManagerContext(tree_index) def _resolve_command_hooks(self, function_name: str) -> tuple[PreparedInvocation, ...]: hooks = [] @@ -566,7 +528,8 @@ def _resolve_observation_dispatch( if len(parameters) != 2 + len(term.args): raise TypeError( f"Manager observation {group}.{name} dispatch must take a ManagerContext parameter, an " - f"torch.Tensor output parameter, and {len(term.args)} positional args; got " + "torch.Tensor output parameter, and " + f"{len(term.args)} positional args; got " f"{tuple(parameter.name for parameter in parameters)}." ) annotations = get_type_hints(function, include_extras=True) @@ -576,40 +539,19 @@ def _resolve_observation_dispatch( ): raise TypeError( f"Manager observation {group}.{name} dispatch must annotate its first two parameters as " - f"ManagerContext and torch.Tensor (conventionally named ctx and out)." + "ManagerContext and torch.Tensor (conventionally named ctx and out)." ) - expressions = [] - prepared_indices = [] - warmup_values = [] - for index, value in enumerate(term.args): - if is_kernel_data(value): - _, tree_def = flatten_kernel_data(value) - layout = self._kernel_data_lowering.lower( - tree_def, context=f"Observation {group}.{name} args[{index}]", force_shared=True - ) - prepared_index, expression = self._register_prepared( - f"observation.{group}.{name}.args[{index}]", value, type(value), layout - ) - prepared_indices.append(prepared_index) - expressions.append(expression) - warmup_values.append(value) - else: - expression, warmup, prepared_index = self._resolve_plain_arg( - f"observation.{group}.{name}.args[{index}]", value - ) - prepared_indices.append(prepared_index) - expressions.append(expression) - warmup_values.append(warmup) - expressions = tuple(expressions) + args = tuple( + self._resolve_term_arg(f"observation.{group}.{name}.args[{index}]", value) + for index, value in enumerate(term.args) + ) dispatcher = function return PreparedInvocation( f"{group}.{name}", "observation", dispatcher, None, - args_expressions=expressions, - args_values=tuple(warmup_values), - args_prepared_indices=tuple(prepared_indices), + args=args, output_size=output_size, ) @@ -682,35 +624,16 @@ def _resolve_dispatch_term( f"Manager {group} {name} dispatch must annotate its first parameter as ManagerContext " f"(conventionally named ctx)." ) - expressions = [] - prepared_indices = [] - warmup_values = [] - for index, value in enumerate(term.args): - if is_kernel_data(value): - _, tree_def = flatten_kernel_data(value) - layout = self._kernel_data_lowering.lower( - tree_def, context=f"{group}.{name} args[{index}]", force_shared=True - ) - prepared_index, expression = self._register_prepared( - f"{group}.{name}.args[{index}]", value, type(value), layout - ) - prepared_indices.append(prepared_index) - expressions.append(expression) - warmup_values.append(value) - else: - expression, warmup, prepared_index = self._resolve_plain_arg(f"{group}.{name}.args[{index}]", value) - prepared_indices.append(prepared_index) - expressions.append(expression) - warmup_values.append(warmup) + args = tuple( + self._resolve_term_arg(f"{group}.{name}.args[{index}]", value) for index, value in enumerate(term.args) + ) dispatcher = function return PreparedInvocation( f"{group}.{name}", group, dispatcher, None, - args_expressions=tuple(expressions), - args_values=tuple(warmup_values), - args_prepared_indices=tuple(prepared_indices), + args=args, ) def _resolve_termination( @@ -751,7 +674,7 @@ def _resolve_receiver_term( if annotations.get(parameters[0].name) is not ManagerContext: raise TypeError( f"Manager term {group}.{name} {function_name} must annotate its first parameter after the " - f"receiver as ManagerContext (conventionally named ctx)." + "receiver as ManagerContext (conventionally named ctx)." ) if argument_parameters: raise TypeError( @@ -764,83 +687,62 @@ def _resolve_receiver_term( dispatcher = function return PreparedInvocation(context, kind, dispatcher, receiver_key, output_size=output_size) - # -- input preparation ---------------------------------------------------- - - def _register_prepared( - self, - source_name: str, - value: Any, - value_type: type[Any], - layout: KernelDataLayout, - ) -> tuple[int, str]: - prepared_index = len(self._prepared_terms) - input_offset = self._flat_input_count - self._prepared_terms.append(KernelInputSource.prepare(source_name, value_type, layout, value)) - self._input_offsets.append(input_offset) - self._flat_input_count += len(iter_layout_leaves(layout)) - return prepared_index, lane_expression(layout, input_offset) - - def _scalar_buffer_arg(self, source_name: str, value: float) -> tuple[str, np.float32]: - """Lower one numeric term argument into a runtime scalar buffer slot. - - Returns the lane argument expression and the normalized warmup value. - Only the fixed float32 layout enters the resolved metadata; the value - itself is read from the input slot at runtime. - """ - scalar = np.float32(value) - wrapped = TermScalarBuffer(scalar) - _, tree_def = flatten_kernel_data(wrapped) - layout = self._kernel_data_lowering.lower( - tree_def, - context=source_name, - force_shared=True, - ) - prepared_index, _ = self._register_prepared(source_name, wrapped, TermScalarBuffer, layout) - leaf = iter_layout_leaves(layout)[0] - return f"input_{self._input_offsets[prepared_index] + leaf.slot_index}", scalar - - def _array_buffer_arg(self, source_name: str, value: np.ndarray) -> tuple[str, np.ndarray]: - """Lower one ndarray term argument into a per-environment input slot. - - Returns the lane argument expression (the lane row) and the warmup - value. - """ - wrapped = TermArrayBuffer(value=value) - _, tree_def = flatten_kernel_data(wrapped) - layout = self._kernel_data_lowering.lower( - tree_def, - context=source_name, - ) - prepared_index, _ = self._register_prepared(source_name, wrapped, TermArrayBuffer, layout) - leaf = iter_layout_leaves(layout)[0] - return f"input_{self._input_offsets[prepared_index] + leaf.slot_index}", value + # -- argument resolution --------------------------------------------------- - def _resolve_plain_arg(self, source_name: str, value: Any) -> tuple[str | None, Any, int | None]: - """Resolve one non-kernel-data term argument. + def _resolve_term_arg(self, source_name: str, value: Any) -> PreparedArg: + """Resolve one manager term argument into a prepared descriptor. - Float-like scalars become runtime scalar buffers; every other scalar - stays a plain constant. Returns the lane argument expression, the - warmup value, and the prepared index (``None`` for both plain scalars - and buffers). + KernelData arguments become structured trees; float-like scalars and + ndarrays become shared leaf slots read by reference; simulator queries + resolve to their read-program batched array; every other value is a + literal constant. """ + if is_kernel_data(value): + _, tree_def = flatten_kernel_data(value) + layout = self._kernel_data_lowering.lower( + tree_def, + context=source_name, + force_shared=True, + ) + return TreeArg(self._register_tree(source_name, value, layout)) if isinstance(value, (float, np.floating)): - expression, warmup = self._scalar_buffer_arg(source_name, value) - return expression, warmup, None + return LeafArg(ResolvedLeaf(np.float32(value), KernelDataScope.SHARED)) if isinstance(value, np.ndarray): - expression, warmup = self._array_buffer_arg(source_name, value) - return expression, warmup, None - if isinstance(value, torch.Tensor): - return f"torch.Tensor({value.dtype}{tuple(value.shape)})", value, None - if isinstance(value, str): - return repr(value), value, None + return LeafArg(ResolvedLeaf(value, KernelDataScope.SHARED)) if isinstance(value, SimDataQuery): # A query passed as a term argument resolves to its registered # read-program batched view; the dispatch receives the array in # the argument's position. - key = self._env._term_query_keys[value] - slot = self._sim_slots_by_key[key] - return f"input_{slot}[env_id]", self._sim_inputs[key].value[0], None - return repr(value), value, None + return SimArg(self._env._term_query_keys[value]) + return ConstantArg(value) + + def _register_tree( + self, + source_name: str, + value: Any, + layout: KernelDataLayout, + ) -> int: + """Validate one kernel-data value against its layout and record its tree. + + Leaves are captured by reference: array leaves are refreshed in place + by their owning runtime objects and scalar/shared leaves are reused + unchanged. + """ + values, tree_def = flatten_kernel_data(value) + if tree_def.fingerprint != layout.tree_def.fingerprint: + raise TypeError(f"{source_name} value has a different KernelData schema.") + leaves: dict[tuple[str, ...], ResolvedLeaf] = {} + for prepared_field, field_value in zip(iter_layout_leaves(layout), values, strict=True): + if prepared_field.value_type is np.ndarray and not isinstance(field_value, np.ndarray): + raise TypeError(f"{source_name} field {'.'.join(prepared_field.path)!r} must be np.ndarray.") + if prepared_field.value_type is not np.ndarray and type(field_value) is not prepared_field.value_type: + raise TypeError( + f"{source_name} field {'.'.join(prepared_field.path)!r} must return " + f"{prepared_field.value_type.__name__}, got {type(field_value).__name__}." + ) + leaves[prepared_field.path] = ResolvedLeaf(field_value, prepared_field.scope) + self._prepared_trees.append(ResolvedTree(value, leaves)) + return len(self._prepared_trees) - 1 def _validate_observation_spaces(self, groups: dict[str, ObservationGroupEntry]) -> None: expected_policy = self._env.policy_observation_space.shape diff --git a/motrix_env_core/src/motrix_env_core/torch_manager/runtime.py b/motrix_env_core/src/motrix_env_core/torch_manager/runtime.py index 0c0c571d..b354ba4d 100644 --- a/motrix_env_core/src/motrix_env_core/torch_manager/runtime.py +++ b/motrix_env_core/src/motrix_env_core/torch_manager/runtime.py @@ -1,17 +1,17 @@ # Copyright Motphys Technology Co., Ltd. 2025, 2026 -# SPDX-License: Apache-2.0 +# SPDX-License-Identifier: Apache-2.0 """Generic Torch executor for resolved Manager programs. The runtime consumes the metadata produced by the manager resolver (read -plan, layout, and prepared invocations) and executes the same term schedule -as the previous fused Numba kernels — command ``update`` hooks, rewards, -terminations, command ``advance`` hooks for evaluate; per-group observation -terms for observe; command ``reset_env`` hooks plus simulator reset terms -for reset — with one batched tensor call per term instead of a per-lane -loop. All writes go back into the persistent numpy buffers the host -lifecycle already owns, so ``ManagerEnv`` keeps its current API and state -surfaces unchanged. +plan, layout, resolved context/term-argument trees, and prepared invocations) +and executes the same term schedule as the previous fused Numba kernels — +command ``update`` hooks, rewards, terminations, command ``advance`` hooks +for evaluate; per-group observation terms for observe; command ``reset_env`` +hooks plus simulator reset terms for reset — with one batched tensor call per +term instead of a per-lane loop. All writes go back into the persistent numpy +buffers the host lifecycle already owns, so ``ManagerEnv`` keeps its current +API and state surfaces unchanged. By default each kernel schedule is additionally fused with ``torch.compile`` into a compiled closure over the captured persistent tensors (one fused @@ -25,7 +25,6 @@ from __future__ import annotations import os -import re import warnings from collections.abc import Callable from dataclasses import dataclass @@ -36,14 +35,13 @@ from motrix_env_core.manager.kernel import KernelOutputs from motrix_env_core.manager.kernel_data import KernelDataScope +from motrix_env_core.manager.resolver import LeafArg, SimArg, TreeArg from motrix_env_core.torch_manager.rand import draw_uniform if TYPE_CHECKING: from motrix_env_core.manager.env import ManagerEnv from motrix_env_core.manager.resolver import PreparedInvocation, ResolvedManagerProgram -_INPUT_EXPRESSION = re.compile(r"^input_(\d+)(\[env_id\])?$") - # Simulator-backed read-only arrays are wrapped as tensors but never written # by the executor; suppress PyTorch's non-writable-array warning once instead # of copying every read-only leaf on every view build. @@ -82,28 +80,29 @@ class TorchTaskProgram: class _TensorCache: - """Tensor views over one persistent flat-input tuple. + """Tensor views over one resolved-input tree anchor. - Simulator-backed arrays are rewritten in place by the read program, so - views stay valid for the lifetime of the input tuple. The cache rebuilds - only when the tuple identity changes. + Leaves are held by reference in the resolved trees and simulator-backed + arrays are rewritten in place by the read program, so views stay valid for + the lifetime of the anchor. The cache rebuilds only when the anchor + identity changes; tensors are keyed by leaf-value identity. """ def __init__(self) -> None: - self._inputs: tuple[Any, ...] | None = None - self._tensors: dict[int, torch.Tensor] = {} + self._anchor: Any = None + self._tensors: dict[int, Any] = {} - def tensors(self, inputs: tuple[Any, ...]) -> dict[int, torch.Tensor]: - if self._inputs is not inputs: - self._inputs = inputs + def tensors(self, anchor: Any) -> dict[int, Any]: + if self._anchor is not anchor: + self._anchor = anchor self._tensors = {} return self._tensors - def tensor(self, inputs: tuple[Any, ...], slot: int) -> torch.Tensor: - tensors = self.tensors(inputs) - tensor = tensors.get(slot) + def tensor(self, anchor: Any, value: Any) -> Any: + tensors = self.tensors(anchor) + key = id(value) + tensor = tensors.get(key) if tensor is None: - value = inputs[slot] if isinstance(value, np.ndarray): # 0-dim scalar leaves are configuration constants (body # indices, sizes, noise amplitudes, mode flags); as tensors @@ -114,66 +113,74 @@ def tensor(self, inputs: tuple[Any, ...], slot: int) -> torch.Tensor: # trace time. if value.shape == (): if np.issubdtype(value.dtype, np.integer): - tensors[slot] = int(value) - return tensors[slot] + tensors[key] = int(value) + return tensors[key] if np.issubdtype(value.dtype, np.floating): - tensors[slot] = float(value) - return tensors[slot] + tensors[key] = float(value) + return tensors[key] if np.issubdtype(value.dtype, np.bool_): - tensors[slot] = bool(value) - return tensors[slot] + tensors[key] = bool(value) + return tensors[key] if value.dtype == np.uint64: value = value.view(np.int64) tensor = torch.from_numpy(value) elif isinstance(value, np.integer): - tensors[slot] = int(value) - return tensors[slot] + tensors[key] = int(value) + return tensors[key] elif isinstance(value, (np.floating, float)): - tensors[slot] = float(value) - return tensors[slot] + tensors[key] = float(value) + return tensors[key] elif isinstance(value, (np.bool_, bool)): - tensors[slot] = bool(value) - return tensors[slot] + tensors[key] = bool(value) + return tensors[key] else: tensor = torch.as_tensor(value) - tensors[slot] = tensor + tensors[key] = tensor return tensor -def _leaf_getter(cache: _TensorCache, inputs: tuple[Any, ...]) -> Callable[[int], torch.Tensor]: - """Leaf resolver over one input tuple. +def _tensorize_leaves(cache: _TensorCache, anchor: Any, tree: Any) -> dict[tuple[str, ...], Any]: + """Eagerly materialize every leaf of one resolved tree. - All tensor views are materialized eagerly so the getter body is a plain - list index — no per-slot branches for Dynamo to specialize when a graph + All tensor views are built up front so namespace resolution is a plain + dict lookup — no per-leaf branches for Dynamo to specialize when a graph break forces it to compile as its own frame. """ + return {path: cache.tensor(anchor, leaf.value) for path, leaf in tree.leaves.items()} + + +def _leaf_lookup(leaves: dict[tuple[str, ...], Any]) -> Callable[[tuple[str, ...]], Any]: + """Traceable leaf getter over one materialized leaf-view mapping. - tensors = [cache.tensor(inputs, slot) for slot in range(len(inputs))] + ``dict.__getitem__`` bound methods are builtins Dynamo refuses to trace; + a plain closure keeps namespace resolution inside the compiled frame. + """ - def get(slot: int) -> torch.Tensor: - return tensors[slot] + def get(path: tuple[str, ...]) -> Any: + return leaves[path] return get class _Namespace: - """Attribute access over one subtree of the flattened manager context. - - Leaf paths resolve to tensors through ``get``; intermediate paths build - child namespaces; names absent from the layout fall back to the owning - host object (plain Python attributes such as scalar parameters). - Resolved children are memoized per instance: the layout is immutable, and - tensor views stay valid for the lifetime of the enclosing input tuple. + """Attribute access over one subtree of the resolved manager context. + + Leaf paths resolve to tensors through ``get`` (the batch view materializer + or the reset-lane gather); intermediate paths build child namespaces; + names absent from the tree fall back to the owning host object (plain + Python attributes such as scalar parameters). Resolved children are + memoized per instance: the tree is immutable, and tensor views stay valid + for the lifetime of the enclosing input anchor. """ def __init__( self, - paths: dict[tuple[str, ...], int], + leaves: dict[tuple[str, ...], Any], prefix: tuple[str, ...], host: Any, - get: Callable[[int], torch.Tensor], + get: Callable[[tuple[str, ...]], Any], ) -> None: - self._paths = paths + self._leaves = leaves self._prefix = prefix self._host = host self._get = get @@ -191,21 +198,19 @@ def __getattr__(self, name: str) -> Any: def _resolve(self, name: str) -> Any: path = (*self._prefix, name) - if path in self._paths: - return self._get(self._paths[path]) - children = { - rest: slot for p, slot in self._paths.items() if p[: len(path)] == path and (rest := p[len(path) :]) - } + if path in self._leaves: + return self._get(path) + children = {rest: path + rest[:1] for p in self._leaves if p[: len(path)] == path and (rest := p[len(path) :])} if children: - # Single-leaf wrappers (for example scalar/array term buffers) - # resolve directly to their only leaf; structured records with - # nested kernel_data build a child namespace so chained paths - # like ``motion.clip.joint_pos`` stay reachable. + # Single-leaf wrappers resolve directly to their only leaf; + # structured records with nested kernel_data build a child + # namespace so chained paths like ``motion.clip.joint_pos`` stay + # reachable. if len(children) == 1: first = next(iter(children)) if len(first) == 1: return self._get(children[first]) - return _Namespace(self._paths, path, self._host, self._get) + return _Namespace(self._leaves, path, self._host, self._get) if self._host is None: raise AttributeError(name) value = getattr(self._host, name) @@ -271,22 +276,22 @@ def __getattr__(self, name: str) -> Any: class _SimMap: - def __init__(self, get: Callable[[int], torch.Tensor], slots: dict[str, int]) -> None: + def __init__(self, get: Callable[[tuple[str, ...]], Any], paths: dict[str, tuple[str, ...]]) -> None: self._get = get - self._slots = slots + self._paths = paths - def __getitem__(self, key: str) -> torch.Tensor: - return self._get(self._slots[key]) + def __getitem__(self, key: str) -> Any: + return self._get(self._paths[key]) class _CommandsMap: def __init__( self, - paths: dict[tuple[str, ...], int], + leaves: dict[tuple[str, ...], Any], command_terms: dict[str, Any], - get: Callable[[int], torch.Tensor], + get: Callable[[tuple[str, ...]], Any], ) -> None: - self._paths = paths + self._leaves = leaves self._command_terms = command_terms self._get = get self._cache: dict[str, _Namespace] = {} @@ -295,9 +300,9 @@ def __getitem__(self, name: str) -> _Namespace: namespace = self._cache.get(name) if namespace is None: prefix = ("commands", name) - if not any(path[:2] == prefix for path in self._paths): + if not any(path[:2] == prefix for path in self._leaves): raise KeyError(name) - namespace = _Namespace(self._paths, prefix, self._command_terms[name], self._get) + namespace = _Namespace(self._leaves, prefix, self._command_terms[name], self._get) self._cache[name] = namespace return namespace @@ -305,11 +310,11 @@ def __getitem__(self, name: str) -> _Namespace: class _ActionsMap: def __init__( self, - paths: dict[tuple[str, ...], int], + leaves: dict[tuple[str, ...], Any], action_terms: dict[str, Any], - get: Callable[[int], torch.Tensor], + get: Callable[[tuple[str, ...]], Any], ) -> None: - self._paths = paths + self._leaves = leaves self._action_terms = action_terms self._get = get self._cache: dict[str, _BatchActions] = {} @@ -318,21 +323,20 @@ def __getitem__(self, name: str) -> _BatchActions: actions = self._cache.get(name) if actions is None: prefix = ("actions", name) - if not any(path[:2] == prefix for path in self._paths): + if not any(path[:2] == prefix for path in self._leaves): raise KeyError(name) - namespace = _Namespace(self._paths, prefix, self._action_terms[name].state, self._get) + namespace = _Namespace(self._leaves, prefix, self._action_terms[name].state, self._get) actions = _BatchActions(namespace) self._cache[name] = actions return actions class _MetricsMap: - def __init__(self, paths: dict[tuple[str, ...], int], get: Callable[[int], torch.Tensor]) -> None: - self._paths = paths + def __init__(self, get: Callable[[tuple[str, ...]], Any]) -> None: self._get = get - def __getitem__(self, name: str) -> torch.Tensor: - return self._get(self._paths[("metrics", name)]) + def __getitem__(self, name: str) -> Any: + return self._get(("metrics", name)) class _BatchContext: @@ -344,21 +348,22 @@ class _BatchContext: def __init__( self, - inputs: tuple[Any, ...], - cache: _TensorCache, - paths: dict[tuple[str, ...], int], + tree: Any, + tensors: dict[tuple[str, ...], Any], env: ManagerEnv, ): - get = _leaf_getter(cache, inputs) - self._paths = paths - self.env_id = get(paths[("env_id",)]) - self.dt = float(inputs[paths[("dt",)]]) - self.sim_reset_requested = get(paths[("sim_reset_requested",)]) - self.rand = _BatchRand(get(paths[("rand", "state")])) - self.sim = _SimMap(get, {path[1]: slot for path, slot in paths.items() if len(path) == 2 and path[0] == "sim"}) - self.commands = _CommandsMap(paths, env.command_terms, get) - self.actions = _ActionsMap(paths, env.action_terms, get) - self.metrics = _MetricsMap(paths, get) + self.env_id = tensors[("env_id",)] + self.dt = float(tree.leaves[("dt",)].value) + self.sim_reset_requested = tensors[("sim_reset_requested",)] + self.rand = _BatchRand(tensors[("rand", "state")]) + get = _leaf_lookup(tensors) + self.sim = _SimMap( + get, + {path[1]: path for path in tensors if len(path) == 2 and path[0] == "sim"}, + ) + self.commands = _CommandsMap(tensors, env.command_terms, get) + self.actions = _ActionsMap(tensors, env.action_terms, get) + self.metrics = _MetricsMap(get) class _WritesMap: @@ -369,14 +374,6 @@ def __getitem__(self, name: str) -> torch.Tensor: return self._writes[name] -def _parse_input_expression(expression: str) -> tuple[int, bool] | None: - """Return ``(slot, is_lane)`` for ``input_N`` / ``input_N[env_id]``.""" - match = _INPUT_EXPRESSION.match(expression) - if match is None: - return None - return int(match.group(1)), match.group(2) is not None - - class TorchManagerRuntime: """Execute one resolved manager program with batched Torch kernels. @@ -392,14 +389,13 @@ def __init__(self, env: ManagerEnv, program: ResolvedManagerProgram, compile: bo self.compile_enabled = compile self._cache = _TensorCache() - context_source = program.prepared_terms[program.context.prepared_index] - context_base = program.input_offsets[program.context.prepared_index] - self._context_paths: dict[tuple[str, ...], int] = { - tuple(field.path): context_base + field.slot_index for field in context_source.fields - } - self._per_env_slots = frozenset( - context_base + field.slot_index for field in context_source.fields if field.scope is KernelDataScope.PER_ENV + self._context_tree = program.prepared_trees[program.context.tree_index] + self._context_leaves = self._context_tree.leaves + self._per_env_paths = frozenset( + path for path, leaf in self._context_leaves.items() if leaf.scope is KernelDataScope.PER_ENV ) + self._sim_paths = {path[1]: path for path in self._context_leaves if len(path) == 2 and path[0] == "sim"} + self._rand_path = ("rand", "state") # Command hooks are static functions on the CommandTerm class taking # the command's kernel-data namespace as their receiver. @@ -447,11 +443,13 @@ def __init__(self, env: ManagerEnv, program: ResolvedManagerProgram, compile: bo ) self._reset_env_invocations = tuple(i for i in self.invocations if i.kind == "command_reset_env") - # Per-inputs-identity caches: the host reuses one persistent input - # tuple (and at most two kernel-output buffer sets), so the batched + # Per-anchor caches: the host reuses one persistent resolved-input + # anchor (and at most two kernel-output buffer sets), so the batched # context and per-invocation argument tuples are built once. self._batch_ctx: _BatchContext | None = None - self._batch_ctx_inputs: tuple[Any, ...] | None = None + self._context_tensors: dict[tuple[str, ...], Any] | None = None + self._tree_tensors: tuple[dict[tuple[str, ...], Any], ...] | None = None + self._batch_ctx_anchor: Any = None self._invocation_args_cache: dict[int, dict[int, tuple[Any, ...]]] = {} self._compiled_kernels: dict[Any, Any] = {} @@ -621,11 +619,10 @@ def _build_compiled_reset( write field. Static shapes let Inductor parallelize every loop. """ ctx = self._context_for(inputs) - rand_slot = self._context_paths[("rand", "state")] - per_env = self._per_env_slots - get = _leaf_getter(self._cache, inputs) - sim_slots = {path[1]: slot for path, slot in self._context_paths.items() if len(path) == 2 and path[0] == "sim"} - sim_slot_ids = frozenset(sim_slots.values()) + assert self._context_tensors is not None + per_env = self._per_env_paths + sim_paths = frozenset(self._sim_paths.values()) + get = _leaf_lookup(self._context_tensors) command_reset_env = self._reset_env_invocations num_envs = ctx.rand.state.shape[0] mask = torch.zeros(num_envs, dtype=torch.bool) @@ -644,28 +641,28 @@ def _build_compiled_reset( ) def kernel() -> None: - clones: dict[int, torch.Tensor] = {} + clones: dict[tuple[str, ...], torch.Tensor] = {} - def lane_get(slot: int) -> torch.Tensor: - cached = clones.get(slot) + def lane_get(path: tuple[str, ...]) -> torch.Tensor: + cached = clones.get(path) if cached is None: - full = get(slot) + full = get(path) # Simulator reads are never written by reset terms (writes # go through the sim_writes map), so they pass through # uncloned; every other per-env leaf may be mutated. - cached = full.clone() if slot in per_env and slot not in sim_slot_ids else full - clones[slot] = cached + cached = full.clone() if path in per_env and path not in sim_paths else full + clones[path] = cached return cached - lane_rand = _LaneRand(rand_slot, ctx, lambda slot, _tensor: lane_get(slot)) + lane_rand = _LaneRand(self._rand_path, lane_get) lane_ctx = _LaneContext( ctx, - _CommandsMap(self._context_paths, self.env.command_terms, lane_get), + _CommandsMap(self._context_tensors, self.env.command_terms, lane_get), lane_rand, - _SimMap(lane_get, sim_slots), - _MetricsMap(self._context_paths, lane_get), + _SimMap(lane_get, self._sim_paths), + _MetricsMap(lane_get), all_ids, - actions=_ActionsMap(self._context_paths, self.env.action_terms, lane_get), + actions=_ActionsMap(self._context_tensors, self.env.action_terms, lane_get), ) for invocation in command_reset_env: receiver = lane_ctx.commands[invocation.receiver_key] @@ -673,9 +670,9 @@ def lane_get(slot: int) -> torch.Tensor: for invocation, _buffers, staging, fields, args in write_plan: sim_writes = {name: buffer for name, buffer in zip(fields, staging, strict=True)} invocation.dispatcher(lane_ctx, _WritesMap(sim_writes), *args) - for slot, tensor in clones.items(): - if tensor is not get(slot): - full = get(slot) + for path, tensor in clones.items(): + if tensor is not get(path): + full = get(path) lane_mask = mask.view((num_envs,) + (1,) * (tensor.dim() - 1)) full.copy_(torch.where(lane_mask, tensor, full)) @@ -702,9 +699,12 @@ def wrapper(ids: torch.Tensor) -> None: # -- helpers ------------------------------------------------------------ def _context_for(self, inputs: tuple[Any, ...]) -> _BatchContext: - if self._batch_ctx_inputs is not inputs: - self._batch_ctx = _BatchContext(inputs, self._cache, self._context_paths, self.env) - self._batch_ctx_inputs = inputs + if self._batch_ctx_anchor is not inputs: + tensors = tuple(_tensorize_leaves(self._cache, inputs, tree) for tree in inputs) + self._tree_tensors = tensors + self._context_tensors = tensors[self.program.context.tree_index] + self._batch_ctx = _BatchContext(self._context_tree, self._context_tensors, self.env) + self._batch_ctx_anchor = inputs self._invocation_args_cache = {} return self._batch_ctx @@ -734,44 +734,35 @@ def _invocation_args_for(self, invocation: Any, inputs: tuple[Any, ...]) -> tupl return args def _invocation_args(self, invocation: PreparedInvocation, inputs: tuple[Any, ...]) -> tuple[Any, ...]: + assert self._tree_tensors is not None args: list[Any] = [] - for value, prepared_index, expression in zip( - invocation.args_values, invocation.args_prepared_indices, invocation.args_expressions, strict=True - ): - if prepared_index is not None: - source = self.program.prepared_terms[prepared_index] - base = self.program.input_offsets[prepared_index] - paths = {tuple(field.path): base + field.slot_index for field in source.fields} - args.append(_Namespace(paths, (), value, _leaf_getter(self._cache, inputs))) - continue - parsed = _parse_input_expression(expression) - if parsed is not None: - slot, _is_lane = parsed - leaf_value = inputs[slot] - if isinstance(leaf_value, np.ndarray): - args.append(self._cache.tensor(inputs, slot)) - else: - args.append(float(leaf_value)) - continue - args.append(value) + for arg in invocation.args: + if isinstance(arg, TreeArg): + leaves = self._tree_tensors[arg.tree_index] + host = self.program.prepared_trees[arg.tree_index].value + args.append(_Namespace(leaves, (), host, _leaf_lookup(leaves))) + elif isinstance(arg, LeafArg): + args.append(self._cache.tensor(inputs, arg.leaf.value)) + elif isinstance(arg, SimArg): + assert self._context_tensors is not None + args.append(self._context_tensors[self._sim_paths[arg.key]]) + else: + args.append(arg.value) return tuple(args) class _LaneRand: """SplitMix64 draws on gathered reset rows, updated in place.""" - def __init__( - self, rand_slot: int, base: _BatchContext, gather: Callable[[int, torch.Tensor], torch.Tensor] - ) -> None: - self._rand_slot = rand_slot - self._base = base + def __init__(self, rand_path: tuple[str, ...], gather: Callable[[tuple[str, ...]], torch.Tensor]) -> None: + self._rand_path = rand_path self._gather = gather self._state: torch.Tensor | None = None @property def state(self) -> torch.Tensor: if self._state is None: - self._state = self._gather(self._rand_slot, self._base.rand.state) + self._state = self._gather(self._rand_path) return self._state def next_uniform(self, count: int = 1) -> torch.Tensor: @@ -784,7 +775,7 @@ def uniform_range(self, count: int, low: float, high: float) -> torch.Tensor: class _LaneContext: - """Reset-lane context: per-env leaves gathered to the selected rows.""" + """Reset-lane context: per-env leaves cloned to the selected rows.""" def __init__( self, diff --git a/motrix_env_core/tests/test_kernel_data.py b/motrix_env_core/tests/test_kernel_data.py index bdf8e2f1..ecb47b9b 100644 --- a/motrix_env_core/tests/test_kernel_data.py +++ b/motrix_env_core/tests/test_kernel_data.py @@ -14,7 +14,6 @@ flatten_kernel_data, iter_layout_leaves, kernel_data, - rebuild_lowered, unflatten_kernel_data, ) @@ -76,7 +75,7 @@ def test_kernel_data_is_frozen_and_compiles_nested_inherited_leaves() -> None: ] -def test_kernel_data_flatten_is_zero_copy_and_rebuilds_both_views() -> None: +def test_kernel_data_flatten_is_zero_copy_and_rebuilds_host_value() -> None: velocity = np.asarray([[1.0]], dtype=np.float32) position = np.asarray([[2.0]], dtype=np.float32) orientation = np.asarray([3.0], dtype=np.float32) @@ -84,11 +83,10 @@ def test_kernel_data_flatten_is_zero_copy_and_rebuilds_both_views() -> None: leaves, tree_def = flatten_kernel_data(value) assert leaves[:3] == (velocity, position, orientation) rebuilt = unflatten_kernel_data(tree_def, leaves) - layout = KernelDataLowering().lower(tree_def, context="test") - lowered = rebuild_lowered(layout, leaves) + KernelDataLowering().lower(tree_def, context="test") assert rebuilt.robot.pose.position is position - assert lowered.robot.pose.orientation is orientation - assert lowered.phase == np.float32(0.25) + assert rebuilt.robot.pose.orientation is orientation + assert rebuilt.phase == np.float32(0.25) def test_kernel_data_schema_fingerprint_ignores_values() -> None: @@ -107,10 +105,10 @@ def test_kernel_data_schema_fingerprint_ignores_values() -> None: first_layout = KernelDataLowering().lower(first, context="first") second_layout = KernelDataLowering().lower(second, context="second") assert first_layout is second_layout - assert first_layout.lowered_type.__module__ == _Tracking.__module__ + assert first_layout.logical_type.__module__ == _Tracking.__module__ -def test_map_schema_changes_tree_and_proxy_identity() -> None: +def test_map_schema_changes_tree_and_layout_identity() -> None: first = _MappedValues(Map({"value": np.float32(1.0)})) second = _MappedValues(Map({"value": np.asarray([[1.0]], dtype=np.float32)})) _, first_tree = flatten_kernel_data(first) @@ -118,8 +116,8 @@ def test_map_schema_changes_tree_and_proxy_identity() -> None: first_layout = KernelDataLowering().lower(first_tree, context="first") second_layout = KernelDataLowering().lower(second_tree, context="second") assert first_tree.fingerprint != second_tree.fingerprint - assert first_layout.lowered_type is not second_layout.lowered_type - assert first_layout.fields[0].child.lowered_type is not second_layout.fields[0].child.lowered_type + assert first_layout.fingerprint != second_layout.fingerprint + assert first_layout.fields[0].child.fingerprint != second_layout.fields[0].child.fingerprint def test_kernel_data_rejects_unsupported_leaf() -> None: diff --git a/motrix_env_core/tests/test_manager_env.py b/motrix_env_core/tests/test_manager_env.py index ea20051b..43144459 100644 --- a/motrix_env_core/tests/test_manager_env.py +++ b/motrix_env_core/tests/test_manager_env.py @@ -16,7 +16,6 @@ import pytest import torch -import motrix_env_core.manager.env as manager_env_module from motrix_env_core.base import EnvCfg # noqa: E402 from motrix_env_core.config import configclass # noqa: E402 from motrix_env_core.config.scene import SceneCfg # noqa: E402 @@ -438,7 +437,7 @@ class _MultipleActionsCfg(ManagerActionsCfg): np.testing.assert_array_equal(env.action_terms["test"].state.source, actions[:, :1]) np.testing.assert_array_equal(env.action_terms["second"].state.source, actions[:, 1:]) env._refresh_sim_reads() - env._execute_observe_kernel(env._kernel_inputs) + env._execute_observe_kernel(env._resolved_inputs) np.testing.assert_array_equal(state.obs.policy, actions) assert [(term.name, term.input_slice) for term in env.manager_layout.actions] == [ ("test", slice(0, 1)), @@ -528,10 +527,11 @@ def test_manager_context_is_injected_once_and_reused_across_all_term_kinds() -> action = env.action_terms["test"] assert isinstance(action, ActionTerm) command = _counter_command(env) - context_sources = [slot for slot in env.manager_layout.inputs if "manager_context" in slot.source] - scopes = {slot.source: slot.scope.value for slot in context_sources} - assert scopes["manager_context.actions.test.action_queue"] == "per_env" - assert scopes["manager_context.actions.test.action_ptr"] == "shared" + assert env._compiled_manager_program is not None + context_tree = env._compiled_manager_program.prepared_trees[env._compiled_manager_program.context.tree_index] + scopes = {path: leaf.scope for path, leaf in context_tree.leaves.items()} + assert scopes[("actions", "test", "action_queue")].value == "per_env" + assert scopes[("actions", "test", "action_ptr")].value == "shared" assert [term.output_slice for term in env.manager_layout.observations["policy"].terms] == [ slice(0, 2), slice(2, 3), @@ -561,23 +561,21 @@ def test_manager_context_is_injected_once_and_reused_across_all_term_kinds() -> assert warmup.signatures -def test_read_plan_uses_compile_time_flat_inputs_without_runtime_flattening( - monkeypatch: pytest.MonkeyPatch, -) -> None: +def test_read_plan_refreshes_resolved_leaves_in_place_without_rebuilding() -> None: env = _ManagerEnv(num_envs=3) env.init_state() assert env._compiled_manager_program is not None - read_plan = env._compiled_manager_program.read_plan - inputs = read_plan.read(env) - - def fail_runtime_flatten(value: object) -> tuple[object, ...]: - del value - raise AssertionError("flatten_kernel_data() must not run after manager compilation") + program = env._compiled_manager_program + context_tree = program.prepared_trees[program.context.tree_index] + leaves_before = dict(context_tree.leaves) + inputs_before = env._resolved_inputs - monkeypatch.setattr(manager_env_module, "flatten_kernel_data", fail_runtime_flatten) + program.read_plan.read(env) + program.read_plan.read(env, np.asarray([0, 2], dtype=np.int64)) - assert read_plan.read(env) is inputs - assert read_plan.read(env, np.asarray([0, 2], dtype=np.int64)) is inputs + assert program.prepared_trees[program.context.tree_index] is context_tree + assert env._resolved_inputs is inputs_before + assert all(context_tree.leaves[path] is leaf for path, leaf in leaves_before.items()) def test_command_evaluation_updates_state_only_in_evaluate_kernel() -> None: @@ -591,7 +589,7 @@ def test_command_evaluation_updates_state_only_in_evaluate_kernel() -> None: derived.fill(-1.0) env._refresh_sim_reads() - inputs = env._kernel_inputs + inputs = env._resolved_inputs env._execute_observe_kernel(inputs) np.testing.assert_array_equal(derived[:, 0], [-1.0, -1.0]) np.testing.assert_array_equal(state.obs.policy[:, 2], [-1.0, -1.0]) @@ -691,7 +689,7 @@ def noisy_observe(inputs, outputs): np.testing.assert_array_equal(env._rand.state, rng_before + np.uint64(1)) saved_final = final.obs.policy.copy() np.testing.assert_array_equal(final.obs.policy[:, 0], (rng_before[[0], 0] + np.uint64(1)).astype(np.float32)) - env._execute_observe_kernel(env._kernel_inputs) + env._execute_observe_kernel(env._resolved_inputs) np.testing.assert_array_equal(env._rand.state, rng_before + np.uint64(2)) np.testing.assert_array_equal(final.obs.policy, saved_final) @@ -710,8 +708,8 @@ def test_post_reset_observation_does_not_run_command_evaluation() -> None: np.testing.assert_array_equal(state.obs.policy[:, 2], 0.0) env._refresh_sim_reads() - env._execute_evaluate_kernel(env._kernel_inputs) - env._execute_observe_kernel(env._kernel_inputs) + env._execute_evaluate_kernel(env._resolved_inputs) + env._execute_observe_kernel(env._resolved_inputs) np.testing.assert_allclose(state.obs.policy[:, 2], [0.5, 1.0, 1.5]) diff --git a/motrix_envs/tests/test_walk_randomization.py b/motrix_envs/tests/test_walk_randomization.py index 8892e7fb..3ca8d53a 100644 --- a/motrix_envs/tests/test_walk_randomization.py +++ b/motrix_envs/tests/test_walk_randomization.py @@ -55,7 +55,7 @@ def _tiny_delay_manager(num_envs: int = 2, num_actuators: int = 3, lo: int = 0, env._action_writes = SimpleNamespace(buffer=lambda name: controls, execute=lambda: None) env._command_terms = {} env._state = SimpleNamespace() - env._compiled_manager_program = SimpleNamespace(read_plan=SimpleNamespace(flat_inputs=())) + env._compiled_manager_program = SimpleNamespace(prepared_trees=()) env.perf = SimpleNamespace(scope=lambda name: nullcontext()) env._reset_sim_rows = lambda env_ids, inputs: None return env, term.state, controls diff --git a/wiki/design/manager/kernel-abi.md b/wiki/design/manager/kernel-abi.md index fd741047..793b7fe8 100644 --- a/wiki/design/manager/kernel-abi.md +++ b/wiki/design/manager/kernel-abi.md @@ -2,7 +2,8 @@ ## 摘要 -本文定义 Manager term 的数据 ABI:`@kernel_data` logical tree 如何 flatten,`ManagerContext` 如何以整批张量 +本文定义 Manager term 的数据 ABI:`@kernel_data` logical tree 如何解析为 +path-keyed 的 resolved tree,`ManagerContext` 如何以整批张量 提供固定 namespace,以及 `PER_ENV`、`SHARED` 和 metric backing 的 ownership 语义。任务开发者只需遵守 [Manager Task API](./task-authoring.md) 中的 term 协议;数值执行器见 [Torch 执行器设计](./torch-executor.md), backend 读侧和 simulator 边界见 [SimBackend](./sim-backend.md)。 @@ -25,7 +26,7 @@ class ManagerContext: Context 不包含 config、provider、backend handle、dynamic Python container 或 generic `states` / `resources` store。 Persistent runtime data 直接归属 concrete term;framework-owned random state 通过 `rand` 暴露,`dt` 提供当前 -`ctrl_dt` 的固定 float32 标量。执行器把该结构解析为固定 layout,运行时以 namespace 视图(`ctx.sim[key]`、 +`ctrl_dt` 的固定 float32 标量。执行器把该结构解析为固定 layout(path → leaf 引用 + scope),运行时以 namespace 视图(`ctx.sim[key]`、 `ctx.commands[name].field` 等)暴露整批张量。 ## 2. 批量语义 @@ -36,12 +37,11 @@ Host 侧 `env_id` 是: np.arange(num_envs, dtype=np.int64) ``` -所有 `PER_ENV` leaves 是 `(N, ...)` 整批数组;term 一次调用处理整批。partial reset 时执行器按原始 -environment id gather `PER_ENV` leaves(第 0 维 = len(env_ids)),term 写 gathered 行后 scatter 回常驻缓冲, -不会把 reset batch row 当作 environment id。 +所有 `PER_ENV` leaves 是 `(N, ...)` 整批数组;term 一次调用处理整批。partial reset 采用静态全批 masked 设计: +reset 数学对全部 lane 计算,选中的 lane 由布尔掩码选择,`PER_ENV` leaves 按 clone-on-first-access 修改后以 +`where(mask, clone, original)` 合并回常驻缓冲,未选中 lane 保持不变,不会把 reset batch row 当作 environment id。 -所有 Map 的 key 顺序、concrete value schema、field tree、leaf dtype、scope 和 input slot 在 resolve 前固定。 -Map lookup 使用 literal key: +所有 Map 的 key 顺序、concrete value schema、field tree、leaf dtype、scope 和 resolved path 在 resolve 前固定。Map lookup 使用 literal key: ```python motion = ctx.commands["motion"] @@ -54,7 +54,7 @@ dof_pos = ctx.sim["robot_dof_pos"] | 类型 | 默认 scope | term 访问语义 | | --- | --- | --- | -| `np.ndarray` | `PER_ENV` | 整批 `(N, ...)` 数组/张量;reset 场景为 gathered 行 | +| `np.ndarray` | `PER_ENV` | 整批 `(N, ...)` 数组/张量;reset 场景为 masked clone 行 | | `SharedArray` | `SHARED` | 完整可写 ndarray,所有环境共享同一份 | | Python scalar | `SHARED` | 所有环境使用同一个标量 | | 固定 dtype NumPy scalar | `SHARED` | 所有环境使用同一个定 dtype 标量 | @@ -62,7 +62,7 @@ dof_pos = ctx.sim["robot_dof_pos"] `SharedArray` 只描述 scope,不代表 immutable。数组是否可以更新由 concrete term 的 owner/lifecycle 契约决定。 Term receiver 的参数通常按 shared lowering 处理,以免把静态参数误切片。 -## 4. Logical tree 与 flat ABI +## 4. Logical tree 与 resolved tree `@kernel_data` 支持有序、可逆的 record tree: @@ -74,8 +74,8 @@ rebuilt = unflatten_kernel_data(tree_def, leaves) Record 可以嵌套另一个 KernelData record;不支持 multiple inheritance、覆盖 inherited field、普通 dataclass、递归结构、dict、list、字符串、object array 或 backend handle。 Logical tree 与 execution lowering 分离:同一 concrete value 在 host term 和 context 中保留相同的 tree schema, -但 resolver 根据使用位置决定 scope、slot 与运行时 namespace 视图。flat ABI 只允许 ndarray/scalar leaves 跨越 -执行边界;term 内不应产生稳态临时大数组分配。 +但 resolver 根据使用位置决定 scope 并产出 path-keyed 的 resolved tree(leaf 引用 + scope),运行时据此构建 +namespace 视图。只有 ndarray/scalar leaves 可以成为 resolved leaves;term 内不应产生稳态临时大数组分配。 ## 5. Ownership 与 metric @@ -92,11 +92,11 @@ Term 可以读取其他 term 的数据,但不应调用其他 term 的行为方 `ManagerProgramResolver`(backend-neutral)负责: -- flatten concrete runtime terms 和 `ManagerContext`; +- 校验 concrete runtime terms 和 `ManagerContext` 的 KernelData schema(leaf 类型与 scope); - 收集 read program 的 sim input map; -- 解析 context、term receiver 和 literal Map lookup 的 slot/表达式; +- 把 context 与 term 参数解析为 path-keyed resolved trees 和参数描述符(常量/叶子/树/sim key); - 校验方法签名、注解契约和 layout; -- 产出 `ManagerLayout`、`ReadPlan`、`KernelInputSource` 与 `PreparedInvocation`,供 Torch 执行器调度。 +- 产出 `ManagerLayout`、`ReadPlan`、resolved trees 与 `PreparedInvocation`,供 Torch 执行器调度。 构造期 warm-up 以克隆输入执行一次 evaluate/observe,确认输入/输出契约(shape、dtype、batch 维、可写性); schema rename、term type 变化、scope 变化和 query layout 变化都会改变 resolved layout。 @@ -104,8 +104,8 @@ schema rename、term type 变化、scope 变化和 query layout 变化都会改 ## 7. 核心不变量 - Context schema 在一个 environment 生命周期内固定; -- 每个 key 对应唯一 concrete type 和 stable slot; -- 所有 `PER_ENV` leaves 共享同一个第 0 维(整批为 N,reset 为 len(env_ids)); +- 每个 key 对应唯一 concrete type 和 stable resolved path; +- 所有 `PER_ENV` leaves 共享同一个第 0 维(整批为 N;reset 数学以 masked 全批形式运行); - `ctx.env_id` 在 step/reset 中与其他 per-env data 对齐; - simulator snapshots 只来自编译后的 read program; - 不建立 generic `states`、`resources` 或 `ManagerValue` registry; diff --git a/wiki/design/manager/torch-executor.md b/wiki/design/manager/torch-executor.md index 236e4c03..5515813b 100644 --- a/wiki/design/manager/torch-executor.md +++ b/wiki/design/manager/torch-executor.md @@ -21,11 +21,12 @@ RNG draw 顺序与 eager 完全一致;Inductor 重结合仅带来 ≤1e-8 量 ## 解析与执行的分工 -- `ManagerProgramResolver`(backend-neutral)负责输入准备:flatten - `@kernel_data` term 与 `ManagerContext`,构建 `ManagerLayout`、 - `KernelInputSource`/`ReadPlan`、`PreparedInvocation`(每个 term 的 - dispatcher、参数表达式/值/prepared index),并校验签名注解契约与 - 签名契约。它不做任何代码生成、缓存或 dispatcher 编译。 +- `ManagerProgramResolver`(backend-neutral)负责输入准备:把 `@kernel_data` + term 参数与 `ManagerContext` 解析为 path-keyed 的 resolved tree(叶子持有 + live value 引用与 `PER_ENV`/`SHARED` scope),构建 `ManagerLayout`、 + `ReadPlan`、`PreparedInvocation`(每个 term 的 dispatcher 与 resolved + argument 描述符),并校验签名注解契约与签名契约。它不做任何代码生成、 + 缓存或 dispatcher 编译。 - `TorchManagerRuntime` 消费这些元数据并执行 term 调度;所有写入落回 host 持有的 numpy 缓冲(reward terms/masks/obs/command state/metrics), `ManagerEnv` 的公共 API 与状态面保持不变。 @@ -36,14 +37,18 @@ RNG draw 顺序与 eager 完全一致;Inductor 重结合仅带来 ≤1e-8 量 (raw、`weighted = raw * weight * dt`、Σ)→ terminations(masks、OR)→ command `advance`(rewards/terminations 读 pre-advance command state); observe 按 group/term 顺序写 obs 切片;reset = command `reset_env` - (lane-gathered)→ sim-reset terms 写 compact 行。 + (masked)→ sim-reset terms 写 compact 行。 - RNG:per-lane SplitMix64(`torch_manager.rand.draw_uniform`),批量语义为 term 一次 draw `count=k` 等价于 per-lane 连续调用 k 次。 -- 输入元组持久复用:simulator-backed 数组由 read program 原地重写, - `_TensorCache` 在同一输入元组上缓存 tensor 视图。 -- reset 路径对 `PER_ENV` slot 做 lane gather,term 写 gathered 行后 scatter - 回常驻缓冲;compact sim-reset write 缓冲布局与 `SimResetRuntime.buffers` - 一致。 +- 解析树持久复用:resolver 产出的 resolved tree 是三个 kernel 共用的持久 + anchor;simulator-backed 数组叶子由 read program 原地重写, + `_TensorCache` 在同一 anchor 上按叶子值身份缓存 tensor 视图,0 维 + int/float/bool 叶子转为 Python 标量(graph-break-free 编译的关键)。 +- reset 路径是静态全批 masked 设计:选中的 lane 由布尔掩码选择, + `PER_ENV` 叶子在 reset term 首次访问时 clone(sim 叶子除外),term 写 + clone 行后以 `where(mask, clone, original)` 合并回常驻缓冲,未选中 lane + (含 RNG 状态)保持不变;compact sim-reset write 缓冲布局与 + `SimResetRuntime.buffers` 一致。 ## 批量 term 约定 From 973c42d3e61289e8f287d57d26c16b5f246b3ed1 Mon Sep 17 00:00:00 2001 From: zilch <147668916@qq.com> Date: Fri, 6 Nov 2026 01:25:11 +0800 Subject: [PATCH 22/41] refactor(manager): replace PreparedInvocation with per-kind typed term 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. --- .../src/motrix_env_core/manager/resolver.py | 317 +++++++++--------- .../motrix_env_core/torch_manager/runtime.py | 193 +++++------ wiki/design/manager/kernel-abi.md | 2 +- wiki/design/manager/runtime.md | 2 +- wiki/design/manager/torch-executor.md | 7 +- 5 files changed, 248 insertions(+), 273 deletions(-) diff --git a/motrix_env_core/src/motrix_env_core/manager/resolver.py b/motrix_env_core/src/motrix_env_core/manager/resolver.py index 556ebbbf..3562a46c 100644 --- a/motrix_env_core/src/motrix_env_core/manager/resolver.py +++ b/motrix_env_core/src/motrix_env_core/manager/resolver.py @@ -6,11 +6,11 @@ Lowers one constructed manager environment into its execution metadata: the simulator read plan, the resolved context/term-argument trees (path-keyed leaf descriptors with live value references and data scopes), prepared term -arguments, and one invocation per manager term (command hooks, observations, -rewards, terminations, simulator resets). Numerical executors — currently the -Torch executor in :mod:`motrix_env_core.torch_manager` — consume this -metadata and schedule the invocations; no code generation or dispatcher -compilation happens here. +arguments, and one typed call record per manager term (command hooks, +observations, rewards, terminations, simulator resets). Numerical executors — +currently the Torch executor in :mod:`motrix_env_core.torch_manager` — +consume this metadata and schedule the typed calls; no code generation or +dispatcher compilation happens here. """ from __future__ import annotations @@ -18,7 +18,7 @@ import hashlib import inspect import logging -from collections.abc import Mapping +from collections.abc import Callable, Mapping from dataclasses import dataclass from time import perf_counter from types import MappingProxyType @@ -167,15 +167,63 @@ class SimArg: @dataclass(frozen=True) -class PreparedInvocation: - """Runtime metadata required to execute one manager term.""" +class CommandUpdateHook: + """One command term's ``update`` hook (receiver + static function).""" - context: str - kind: str - dispatcher: Any - receiver_key: str | None - args: tuple[PreparedArg, ...] = () - output_size: int = 0 + receiver_key: str + function: Callable + + +@dataclass(frozen=True) +class CommandAdvanceHook: + """One command term's ``advance`` hook (receiver + static function).""" + + receiver_key: str + function: Callable + + +@dataclass(frozen=True) +class CommandResetHook: + """One command term's ``reset_env`` hook (receiver + static function).""" + + receiver_key: str + function: Callable + + +@dataclass(frozen=True) +class RewardTermCall: + """One reward term call with its resolver-assigned buffer column.""" + + index: int + function: Callable + args: tuple[PreparedArg, ...] + + +@dataclass(frozen=True) +class TerminationTermCall: + """One termination term call with its resolver-assigned buffer column.""" + + index: int + function: Callable + args: tuple[PreparedArg, ...] + + +@dataclass(frozen=True) +class ObservationTermCall: + """One observation term call writing its group slice.""" + + group: str + name: str + function: Callable + args: tuple[PreparedArg, ...] + + +@dataclass(frozen=True) +class SimResetTermCall: + """One simulator reset term call (write offsets stay on ``ResolvedSimReset``).""" + + function: Callable + args: tuple[PreparedArg, ...] @dataclass(frozen=True) @@ -185,7 +233,7 @@ class ResolvedManagerContext: @dataclass(frozen=True) class ResolvedSimReset: - invocation: PreparedInvocation + call: SimResetTermCall sim_writes_type: str output_offset: int output_count: int @@ -193,11 +241,25 @@ class ResolvedSimReset: @dataclass(frozen=True) class ResolvedManagerProgram: - """Resolved execution metadata for one manager environment.""" + """Resolved execution metadata for one manager environment. + + The phase sequences encode the execution contract structurally: evaluate + runs command updates, then rewards, then terminations, then command + advances (declaration order within each phase); observe iterates the + per-group observation calls; reset runs the command reset hooks followed + by the simulator reset terms. Reward and termination indices match + ``ManagerLayout.rewards``/``terminations`` ordering. + """ read_plan: ReadPlan layout: ManagerLayout - invocations: tuple[PreparedInvocation, ...] + command_updates: tuple[CommandUpdateHook, ...] + rewards: tuple[RewardTermCall, ...] + terminations: tuple[TerminationTermCall, ...] + command_advances: tuple[CommandAdvanceHook, ...] + command_reset_hooks: tuple[CommandResetHook, ...] + sim_reset_terms: tuple[ResolvedSimReset, ...] + observations: dict[str, tuple[ObservationTermCall, ...]] prepared_trees: tuple[ResolvedTree, ...] context: ResolvedManagerContext fingerprint: str @@ -205,21 +267,27 @@ class ResolvedManagerProgram: __all__ = [ "ActionTermLayout", + "CommandAdvanceHook", + "CommandResetHook", + "CommandUpdateHook", "ConstantArg", "LeafArg", "ManagerLayout", "ManagerProgramResolver", "MetricFieldLayout", "ObservationGroupLayout", + "ObservationTermCall", "ObservationTermLayout", "PreparedArg", - "PreparedInvocation", "ResolvedManagerContext", "ResolvedManagerProgram", "ResolvedSimReset", + "RewardTermCall", "ScalarTermLayout", "SimArg", "SimInputLayout", + "SimResetTermCall", + "TerminationTermCall", "TreeArg", ] @@ -277,13 +345,19 @@ def build(self) -> ResolvedManagerProgram: reward_terms = self._env._reward_terms termination_terms = self._env.termination_manager.terms context = self._resolve_manager_context(observation_groups, reward_terms, termination_terms) - command_updates = self._resolve_command_hooks("update") - command_advances = self._resolve_command_hooks("advance") - command_resets = self._resolve_command_hooks("reset_env") + command_updates = tuple( + CommandUpdateHook(key, function) for key, function in self._resolve_command_hooks("update") + ) + command_advances = tuple( + CommandAdvanceHook(key, function) for key, function in self._resolve_command_hooks("advance") + ) + command_reset_hooks = tuple( + CommandResetHook(key, function) for key, function in self._resolve_command_hooks("reset_env") + ) observation_terms, observation_layout = self._resolve_observations(observation_groups) - reward_invocations, reward_layout, _ = self._resolve_rewards(self._env.cfg, reward_terms) - termination_invocations, termination_layout = self._resolve_terminations(termination_terms) - reset_invocations = self._resolve_resets() + reward_calls, reward_layout, _ = self._resolve_rewards(self._env.cfg, reward_terms) + termination_calls, termination_layout = self._resolve_terminations(termination_terms) + sim_reset_terms = self._resolve_resets() per_env_metric_layout = tuple( MetricFieldLayout(name, index, np.dtype(value.dtype)) for index, (name, value) in enumerate(self._env.metrics.items()) @@ -315,28 +389,34 @@ def build(self) -> ResolvedManagerProgram: terminations=termination_layout, metrics=per_env_metric_layout, ) - invocations = ( - command_updates - + command_advances - + command_resets - + tuple(term for terms in observation_terms.values() for term in terms) - + reward_invocations - + termination_invocations - + tuple(reset.invocation for reset in reset_invocations) + call_count = ( + len(command_updates) + + len(command_advances) + + len(command_reset_hooks) + + sum(len(terms) for terms in observation_terms.values()) + + len(reward_calls) + + len(termination_calls) + + len(sim_reset_terms) ) leaf_count = sum(len(tree.leaves) for tree in self._prepared_trees) logger.info( - "Manager startup %s: resolve finished in %.3fs (%d invocations, %d resolved leaves)", + "Manager startup %s: resolve finished in %.3fs (%d term calls, %d resolved leaves)", type(self._env).__name__, perf_counter() - build_started, - len(invocations), + call_count, leaf_count, ) context_tree = self._prepared_trees[context.tree_index] return ResolvedManagerProgram( read_plan=ReadPlan(self._env.sim_data, context_tree), layout=layout, - invocations=invocations, + command_updates=command_updates, + rewards=reward_calls, + terminations=termination_calls, + command_advances=command_advances, + command_reset_hooks=command_reset_hooks, + sim_reset_terms=sim_reset_terms, + observations=observation_terms, prepared_trees=tuple(self._prepared_trees), context=context, fingerprint=hashlib.sha256(repr(tuple(self._layout_identity)).encode()).hexdigest(), @@ -392,15 +472,8 @@ def _resolve_reset( args = tuple( self._resolve_term_arg(f"sim_reset.{name}.args[{index}]", value) for index, value in enumerate(term.args) ) - invocation = PreparedInvocation( - "sim_reset." + name, - "reset", - dispatcher, - None, - args=args, - ) return ResolvedSimReset( - invocation, + SimResetTermCall(dispatcher, args), f"sim_reset_{name}_writes", output_offset, len(fields), @@ -470,54 +543,68 @@ def _resolve_manager_context( ) return ResolvedManagerContext(tree_index) - def _resolve_command_hooks(self, function_name: str) -> tuple[PreparedInvocation, ...]: + def _resolve_command_hooks(self, function_name: str) -> tuple[tuple[str, Callable], ...]: + """Resolve one hook family over all command terms in declaration order. + + Returns ``(receiver_key, function)`` pairs; the static function + signature is validated against the command hook contract. + """ hooks = [] - for index, (name, command_term) in enumerate(self._env.command_terms.items()): + for name, command_term in self._env.command_terms.items(): function = inspect.getattr_static(type(command_term), function_name) if isinstance(function, staticmethod): function = function.__func__ - hooks.append( - self._resolve_receiver_term( - "command", - name, - command_term, - function_name, - index, - 0, - function=function, - receiver_key=name, - ) - ) + self._validate_command_hook(name, function, function_name) + hooks.append((name, function)) return tuple(hooks) + def _validate_command_hook(self, name: str, function: Any, function_name: str) -> None: + if not inspect.isfunction(function): + raise TypeError(f"Manager term command.{name} {function_name} must be a static function.") + parameters = tuple(inspect.signature(function).parameters.values()) + if not parameters: + raise TypeError(f"Manager term command.{name} {function_name} must declare a receiver parameter.") + ctx_parameters = parameters[1:] + if not ctx_parameters: + raise TypeError(f"Manager term command.{name} {function_name} must declare a ManagerContext parameter.") + argument_parameters = ctx_parameters[1:] + annotations = get_type_hints(function, include_extras=True) + if annotations.get(ctx_parameters[0].name) is not ManagerContext: + raise TypeError( + f"Manager term command.{name} {function_name} must annotate its first parameter after the " + "receiver as ManagerContext (conventionally named ctx)." + ) + if argument_parameters: + raise TypeError( + f"Manager term command.{name} {function_name} has unsupported dependency parameters " + f"{[parameter.name for parameter in argument_parameters]}; read manager and simulator data " + "from the ManagerContext parameter." + ) + # -- observations -------------------------------------------------------- def _resolve_observations( self, groups: dict[str, ObservationGroupEntry], - ) -> tuple[dict[str, tuple[PreparedInvocation, ...]], dict[str, ObservationGroupLayout]]: + ) -> tuple[dict[str, tuple[ObservationTermCall, ...]], dict[str, ObservationGroupLayout]]: resolved_groups = {} layout_groups = {} - term_index = 0 for group_name, group in groups.items(): terms = [] layouts = [] offset = 0 for entry in group.terms: - terms.append(self._resolve_observation(group_name, entry.name, entry.term, term_index, entry.size)) + terms.append(self._resolve_observation(group_name, entry.name, entry.term)) output_slice = slice(offset, offset + entry.size) layouts.append(ObservationTermLayout(entry.name, output_slice)) offset = output_slice.stop - term_index += 1 if offset != group.size: raise RuntimeError(f"Observation group {group_name!r} resolved width changed during resolution.") resolved_groups[group_name] = tuple(terms) layout_groups[group_name] = ObservationGroupLayout(group_name, tuple(layouts), group.size) return resolved_groups, layout_groups - def _resolve_observation_dispatch( - self, group: str, name: str, term: ObsTerm, term_index: int, output_size: int - ) -> PreparedInvocation: + def _resolve_observation(self, group: str, name: str, term: ObsTerm) -> ObservationTermCall: function = term.dispatch parameters = tuple(inspect.signature(function).parameters.values()) if any( @@ -545,25 +632,7 @@ def _resolve_observation_dispatch( self._resolve_term_arg(f"observation.{group}.{name}.args[{index}]", value) for index, value in enumerate(term.args) ) - dispatcher = function - return PreparedInvocation( - f"{group}.{name}", - "observation", - dispatcher, - None, - args=args, - output_size=output_size, - ) - - def _resolve_observation( - self, - group: str, - name: str, - term: Any, - term_index: int, - output_size: int, - ) -> PreparedInvocation: - return self._resolve_observation_dispatch(group, name, term, term_index, output_size) + return ObservationTermCall(group, name, function, args) # -- rewards and terminations -------------------------------------------- @@ -571,13 +640,14 @@ def _resolve_rewards( self, cfg: Any, reward_terms: dict[str, RewardTerm], - ) -> tuple[tuple[PreparedInvocation, ...], tuple[ScalarTermLayout, ...], tuple[float, ...]]: + ) -> tuple[tuple[RewardTermCall, ...], tuple[ScalarTermLayout, ...], tuple[float, ...]]: term_cfgs = cfg.reward_cfgs() terms = [] layout = [] weights = [] for index, (name, term) in enumerate(reward_terms.items()): - terms.append(self._resolve_reward(name, term, index)) + args = self._resolve_dispatch_args("reward", name, term) + terms.append(RewardTermCall(index, term.dispatch, args)) layout.append(ScalarTermLayout(name, index)) try: weights.append(float(term_cfgs[name].weight)) @@ -585,27 +655,19 @@ def _resolve_rewards( raise TypeError(f"Manager reward term reward.{name} weight must be float-compatible.") from error return tuple(terms), tuple(layout), tuple(weights) - def _resolve_reward(self, name: str, term: RewardTerm, term_index: int) -> PreparedInvocation: - return self._resolve_dispatch_term("reward", name, term, term_index) - def _resolve_terminations( self, termination_terms: dict[str, TerminationTerm], - ) -> tuple[tuple[PreparedInvocation, ...], tuple[ScalarTermLayout, ...]]: + ) -> tuple[tuple[TerminationTermCall, ...], tuple[ScalarTermLayout, ...]]: terms = [] layout = [] for index, (name, term) in enumerate(termination_terms.items()): - terms.append(self._resolve_termination(name, term, index)) + args = self._resolve_dispatch_args("termination", name, term) + terms.append(TerminationTermCall(index, term.dispatch, args)) layout.append(ScalarTermLayout(name, index)) return tuple(terms), tuple(layout) - def _resolve_dispatch_term( - self, - group: str, - name: str, - term: Any, - term_index: int, - ) -> PreparedInvocation: + def _resolve_dispatch_args(self, group: str, name: str, term: Any) -> tuple[PreparedArg, ...]: function = term.dispatch parameters = tuple(inspect.signature(function).parameters.values()) if any( @@ -624,68 +686,9 @@ def _resolve_dispatch_term( f"Manager {group} {name} dispatch must annotate its first parameter as ManagerContext " f"(conventionally named ctx)." ) - args = tuple( + return tuple( self._resolve_term_arg(f"{group}.{name}.args[{index}]", value) for index, value in enumerate(term.args) ) - dispatcher = function - return PreparedInvocation( - f"{group}.{name}", - group, - dispatcher, - None, - args=args, - ) - - def _resolve_termination( - self, - name: str, - term: TerminationTerm, - term_index: int, - ) -> PreparedInvocation: - return self._resolve_dispatch_term("termination", name, term, term_index) - - # -- command hooks ------------------------------------------------------- - - def _resolve_receiver_term( - self, - group: str, - name: str, - term: Any, - function_name: str, - term_index: int, - output_size: int, - *, - function: Any | None = None, - receiver_key: str, - ) -> PreparedInvocation: - function = getattr(term, function_name) if function is None else function - has_receiver = receiver_key is not None - if not inspect.isfunction(function) or (not has_receiver and inspect.ismethod(function)): - raise TypeError(f"Manager term {group}.{name} {function_name} must be a static function.") - parameters = tuple(inspect.signature(function).parameters.values()) - if has_receiver: - if not parameters: - raise TypeError(f"Manager term {group}.{name} {function_name} must declare a receiver parameter.") - parameters = parameters[1:] - if not parameters: - raise TypeError(f"Manager term {group}.{name} {function_name} must declare a ManagerContext parameter.") - argument_parameters = parameters[1:] - annotations = get_type_hints(function, include_extras=True) - if annotations.get(parameters[0].name) is not ManagerContext: - raise TypeError( - f"Manager term {group}.{name} {function_name} must annotate its first parameter after the " - "receiver as ManagerContext (conventionally named ctx)." - ) - if argument_parameters: - raise TypeError( - f"Manager term {group}.{name} {function_name} has unsupported dependency parameters " - f"{[parameter.name for parameter in argument_parameters]}; read manager and simulator data " - "from the ManagerContext parameter." - ) - kind = f"command_{function_name}" if group == "command" else group - context = f"{group}.{name} ({function.__module__}.{function.__qualname__})" - dispatcher = function - return PreparedInvocation(context, kind, dispatcher, receiver_key, output_size=output_size) # -- argument resolution --------------------------------------------------- diff --git a/motrix_env_core/src/motrix_env_core/torch_manager/runtime.py b/motrix_env_core/src/motrix_env_core/torch_manager/runtime.py index b354ba4d..18c4a8ff 100644 --- a/motrix_env_core/src/motrix_env_core/torch_manager/runtime.py +++ b/motrix_env_core/src/motrix_env_core/torch_manager/runtime.py @@ -4,12 +4,12 @@ """Generic Torch executor for resolved Manager programs. The runtime consumes the metadata produced by the manager resolver (read -plan, layout, resolved context/term-argument trees, and prepared invocations) -and executes the same term schedule as the previous fused Numba kernels — -command ``update`` hooks, rewards, terminations, command ``advance`` hooks -for evaluate; per-group observation terms for observe; command ``reset_env`` -hooks plus simulator reset terms for reset — with one batched tensor call per -term instead of a per-lane loop. All writes go back into the persistent numpy +plan, layout, resolved context/term-argument trees, and per-term-kind typed +call records) and executes the manager term schedule — command ``update`` +hooks, rewards, terminations, command ``advance`` hooks for evaluate; +per-group observation terms for observe; command ``reset_env`` hooks plus +simulator reset terms for reset — with one batched tensor call per term +instead of a per-lane loop. All writes go back into the persistent numpy buffers the host lifecycle already owns, so ``ManagerEnv`` keeps its current API and state surfaces unchanged. @@ -35,12 +35,23 @@ from motrix_env_core.manager.kernel import KernelOutputs from motrix_env_core.manager.kernel_data import KernelDataScope -from motrix_env_core.manager.resolver import LeafArg, SimArg, TreeArg +from motrix_env_core.manager.resolver import ( + CommandAdvanceHook, + CommandResetHook, + CommandUpdateHook, + LeafArg, + ObservationTermCall, + RewardTermCall, + SimArg, + SimResetTermCall, + TerminationTermCall, + TreeArg, +) from motrix_env_core.torch_manager.rand import draw_uniform if TYPE_CHECKING: from motrix_env_core.manager.env import ManagerEnv - from motrix_env_core.manager.resolver import PreparedInvocation, ResolvedManagerProgram + from motrix_env_core.manager.resolver import ResolvedManagerProgram # Simulator-backed read-only arrays are wrapped as tensors but never written # by the executor; suppress PyTorch's non-writable-array warning once instead @@ -397,51 +408,30 @@ def __init__(self, env: ManagerEnv, program: ResolvedManagerProgram, compile: bo self._sim_paths = {path[1]: path for path in self._context_leaves if len(path) == 2 and path[0] == "sim"} self._rand_path = ("rand", "state") - # Command hooks are static functions on the CommandTerm class taking - # the command's kernel-data namespace as their receiver. - self.invocations = tuple(program.invocations) - # The resolver groups invocations by term kind; the evaluate schedule - # must instead follow the fused-kernel statement order: command - # update hooks -> rewards -> terminations -> command advance hooks - # (rewards and terminations read pre-advance command state). - self.evaluate_invocations = tuple( - sorted( - (i for i in self.invocations if i.kind in _EVALUATE_KINDS), - key=lambda invocation: _EVALUATE_ORDER[invocation.kind], - ) - ) - self.observe_invocations = tuple(i for i in self.invocations if i.kind == "observation") - self._reward_indices = { - i.context: index for index, i in enumerate(i for i in self.evaluate_invocations if i.kind == "reward") - } - self._termination_indices = { - i.context: index for index, i in enumerate(i for i in self.evaluate_invocations if i.kind == "termination") - } - self._observation_by_name = {i.context: i for i in self.observe_invocations} + # Typed call sequences encode the execution contract structurally: + # evaluate follows the fixed phase order command update hooks -> + # rewards -> terminations -> command advance hooks (rewards and + # terminations read pre-advance command state); each phase keeps its + # declaration order and reward/termination indices match + # ``layout.rewards``/``layout.terminations``. + self.command_updates: tuple[CommandUpdateHook, ...] = program.command_updates + self.rewards: tuple[RewardTermCall, ...] = program.rewards + self.terminations: tuple[TerminationTermCall, ...] = program.terminations + self.command_advances: tuple[CommandAdvanceHook, ...] = program.command_advances + self.command_reset_hooks: tuple[CommandResetHook, ...] = program.command_reset_hooks + self.observations: dict[str, tuple[ObservationTermCall, ...]] = program.observations reward_cfgs = env.cfg.reward_cfgs() self.reward_weights = np.asarray( [float(reward_cfgs[term.name].weight) for term in self.layout.rewards], dtype=np.float32 ) # Reset-term compact-write layout mirrors SimResetRuntime.buffers: # one entry per declared write, ordered by term then output name. - self._reset_terms: tuple[tuple[Any, int, tuple[str, ...]], ...] = () - offsets: dict[str, int] = {} - offset = 0 - for name, writes in env.sim_reset_writes.items(): - fields = tuple(writes) - offsets[name] = offset - offset += len(fields) - self._reset_terms = tuple( - ( - invocation, - offsets[name], - tuple(env.sim_reset_writes[name]), - ) - for invocation in program.invocations - if invocation.kind == "reset" - for name in (invocation.context[len("sim_reset.") :],) + # The write-field names come from the host (same insertion order the + # resolver used when assigning offsets). + self._reset_terms: tuple[tuple[SimResetTermCall, int, tuple[str, ...]], ...] = tuple( + (reset.call, reset.output_offset, tuple(env.sim_reset_writes[name])) + for name, reset in zip(env.sim_reset_terms, program.sim_reset_terms, strict=True) ) - self._reset_env_invocations = tuple(i for i in self.invocations if i.kind == "command_reset_env") # Per-anchor caches: the host reuses one persistent resolved-input # anchor (and at most two kernel-output buffer sets), so the batched @@ -488,7 +478,7 @@ def observe(self, inputs: tuple[Any, ...], outputs: KernelOutputs) -> None: ) def reset(self, inputs: tuple[Any, ...], env_ids: np.ndarray, reset_buffers: tuple[np.ndarray, ...]) -> None: - if not self._reset_env_invocations and not self._reset_terms: + if not self.command_reset_hooks and not self._reset_terms: return ids = torch.from_numpy(np.ascontiguousarray(env_ids, dtype=np.int64)) kernel = self._compiled_kernels.get(("reset", id(inputs), id(reset_buffers))) @@ -509,29 +499,30 @@ def _build_compiled_evaluate( weighted_terms = torch.from_numpy(buffers[1]) termination_masks = torch.from_numpy(buffers[2]) dt = ctx.dt - reward_indices = self._reward_indices - termination_indices = self._termination_indices - plan = tuple((invocation.kind, invocation, args) for invocation, args in self._evaluate_plan(inputs)) + update_hooks = self.command_updates + advance_hooks = self.command_advances + reward_plan = tuple((term, self._invocation_args_for(term, inputs)) for term in self.rewards) + termination_plan = tuple((term, self._invocation_args_for(term, inputs)) for term in self.terminations) def kernel(reward_view: torch.Tensor, terminated_view: torch.Tensor) -> None: + for hook in update_hooks: + receiver = ctx.commands[hook.receiver_key] + hook.function(receiver, ctx) total = None + for term, args in reward_plan: + value = term.function(ctx, *args).to(torch.float32) + weighted = value * weights[term.index] * dt + raw_terms[:, term.index] = value + weighted_terms[:, term.index] = weighted + total = weighted if total is None else total + weighted terminated = None - for kind, invocation, args in plan: - if kind in ("command_update", "command_advance"): - receiver = ctx.commands[invocation.receiver_key] - invocation.dispatcher(receiver, ctx) - elif kind == "reward": - value = invocation.dispatcher(ctx, *args).to(torch.float32) - index = reward_indices[invocation.context] - weighted = value * weights[index] * dt - raw_terms[:, index] = value - weighted_terms[:, index] = weighted - total = weighted if total is None else total + weighted - else: - mask = invocation.dispatcher(ctx, *args).to(torch.bool) - index = termination_indices[invocation.context] - termination_masks[:, index] = mask - terminated = mask if terminated is None else terminated | mask + for term, args in termination_plan: + mask = term.function(ctx, *args).to(torch.bool) + termination_masks[:, term.index] = mask + terminated = mask if terminated is None else terminated | mask + for hook in advance_hooks: + receiver = ctx.commands[hook.receiver_key] + hook.function(receiver, ctx) reward_view.copy_(total if total is not None else torch.zeros_like(reward_view)) terminated_view.copy_(terminated if terminated is not None else torch.zeros_like(terminated_view)) @@ -541,13 +532,6 @@ def kernel(reward_view: torch.Tensor, terminated_view: torch.Tensor) -> None: def _build_compiled_observe(self, inputs: tuple[Any, ...]) -> Callable[[torch.Tensor, torch.Tensor | None], None]: ctx = self._context_for(inputs) - observation_by_name = self._observation_by_name - plan = tuple( - (group_name, term, observation_by_name[f"{group_name}.{term.name}"]) - for group_name, group in self.layout.observations.items() - for term in group.terms - ) - invocation_args = {id(inv): self._invocation_args_for(inv, inputs) for _, _, inv in plan} # Each term writes a private contiguous buffer; a single fused # ``cat + copy`` assembles the group output. Writing through strided # slice views of the group buffer forces per-term unfusable store @@ -565,12 +549,8 @@ def _build_compiled_observe(self, inputs: tuple[Any, ...]) -> Callable[[torch.Te ( group_name, tuple( - (inv, invocation_args[id(inv)], buf) - for (_group_name, _term, inv), buf in zip( - tuple(entry for entry in plan if entry[0] == group_name), - bufs, - strict=True, - ) + (call, self._invocation_args_for(call, inputs), buf) + for call, buf in zip(self.observations[group_name], bufs, strict=True) ), ) for group_name, bufs in term_bufs.items() @@ -581,12 +561,12 @@ def _build_compiled_observe(self, inputs: tuple[Any, ...]) -> Callable[[torch.Te # which measurably beats a single fused graph on many-core hosts. def make_group_kernel( - entries: tuple[tuple[PreparedInvocation, tuple[Any, ...], torch.Tensor], ...], + entries: tuple[tuple[ObservationTermCall, tuple[Any, ...], torch.Tensor], ...], ) -> Callable[[torch.Tensor], None]: def group_kernel(target: torch.Tensor) -> None: - for invocation, args, buf in entries: - invocation.dispatcher(ctx, buf, *args) - target.copy_(torch.cat([buf for _inv, _args, buf in entries], dim=1)) + for call, args, buf in entries: + call.function(ctx, buf, *args) + target.copy_(torch.cat([buf for _call, _args, buf in entries], dim=1)) if self.compile_enabled: return torch.compile(group_kernel, dynamic=False, options=_compile_options()) @@ -623,21 +603,21 @@ def _build_compiled_reset( per_env = self._per_env_paths sim_paths = frozenset(self._sim_paths.values()) get = _leaf_lookup(self._context_tensors) - command_reset_env = self._reset_env_invocations + command_reset_hooks = self.command_reset_hooks num_envs = ctx.rand.state.shape[0] mask = torch.zeros(num_envs, dtype=torch.bool) all_ids = torch.arange(num_envs) write_plan = tuple( ( - invocation, + call, tuple(torch.from_numpy(reset_buffers[offset + index]) for index in range(len(fields))), tuple( torch.empty_like(torch.from_numpy(reset_buffers[offset + index])) for index in range(len(fields)) ), fields, - self._invocation_args_for(invocation, inputs), + self._invocation_args_for(call, inputs), ) - for invocation, offset, fields in self._reset_terms + for call, offset, fields in self._reset_terms ) def kernel() -> None: @@ -664,12 +644,12 @@ def lane_get(path: tuple[str, ...]) -> torch.Tensor: all_ids, actions=_ActionsMap(self._context_tensors, self.env.action_terms, lane_get), ) - for invocation in command_reset_env: - receiver = lane_ctx.commands[invocation.receiver_key] - invocation.dispatcher(receiver, lane_ctx) - for invocation, _buffers, staging, fields, args in write_plan: + for hook in command_reset_hooks: + receiver = lane_ctx.commands[hook.receiver_key] + hook.function(receiver, lane_ctx) + for call, _buffers, staging, fields, args in write_plan: sim_writes = {name: buffer for name, buffer in zip(fields, staging, strict=True)} - invocation.dispatcher(lane_ctx, _WritesMap(sim_writes), *args) + call.function(lane_ctx, _WritesMap(sim_writes), *args) for path, tensor in clones.items(): if tensor is not get(path): full = get(path) @@ -708,35 +688,30 @@ def _context_for(self, inputs: tuple[Any, ...]) -> _BatchContext: self._invocation_args_cache = {} return self._batch_ctx - def _evaluate_plan(self, inputs: tuple[Any, ...]) -> tuple[tuple[Any, tuple[Any, ...]], ...]: - plan = self._plan_cache(inputs).get("evaluate") - if plan is None: - plan = tuple( - (invocation, self._invocation_args_for(invocation, inputs)) for invocation in self.evaluate_invocations - ) - self._plan_cache(inputs)["evaluate"] = plan - return plan - - def _plan_cache(self, inputs: tuple[Any, ...]) -> dict[str, Any]: + def _plan_cache(self, inputs: tuple[Any, ...]) -> dict[int, tuple[Any, ...]]: cache = self._invocation_args_cache.get(id(inputs)) if cache is None: cache = {} self._invocation_args_cache[id(inputs)] = cache return cache - def _invocation_args_for(self, invocation: Any, inputs: tuple[Any, ...]) -> tuple[Any, ...]: + def _invocation_args_for(self, call: Any, inputs: tuple[Any, ...]) -> tuple[Any, ...]: cache = self._plan_cache(inputs) - key = id(invocation) + key = id(call) args = cache.get(key) if args is None: - args = self._invocation_args(invocation, inputs) + args = self._invocation_args(call, inputs) cache[key] = args return args - def _invocation_args(self, invocation: PreparedInvocation, inputs: tuple[Any, ...]) -> tuple[Any, ...]: + def _invocation_args( + self, + call: RewardTermCall | TerminationTermCall | ObservationTermCall | SimResetTermCall, + inputs: tuple[Any, ...], + ) -> tuple[Any, ...]: assert self._tree_tensors is not None args: list[Any] = [] - for arg in invocation.args: + for arg in call.args: if isinstance(arg, TreeArg): leaves = self._tree_tensors[arg.tree_index] host = self.program.prepared_trees[arg.tree_index].value @@ -798,8 +773,4 @@ def __init__( self.actions = actions if actions is not None else base.actions -_EVALUATE_KINDS = ("command_update", "reward", "termination", "command_advance") -_EVALUATE_ORDER = {"command_update": 0, "reward": 1, "termination": 2, "command_advance": 3} - - __all__ = ["TorchManagerRuntime", "TorchTaskProgram"] diff --git a/wiki/design/manager/kernel-abi.md b/wiki/design/manager/kernel-abi.md index 793b7fe8..817cbf9e 100644 --- a/wiki/design/manager/kernel-abi.md +++ b/wiki/design/manager/kernel-abi.md @@ -96,7 +96,7 @@ Term 可以读取其他 term 的数据,但不应调用其他 term 的行为方 - 收集 read program 的 sim input map; - 把 context 与 term 参数解析为 path-keyed resolved trees 和参数描述符(常量/叶子/树/sim key); - 校验方法签名、注解契约和 layout; -- 产出 `ManagerLayout`、`ReadPlan`、resolved trees 与 `PreparedInvocation`,供 Torch 执行器调度。 +- 产出 `ManagerLayout`、`ReadPlan`、resolved trees 与 per-term-kind typed call records,供 Torch 执行器调度。 构造期 warm-up 以克隆输入执行一次 evaluate/observe,确认输入/输出契约(shape、dtype、batch 维、可写性); schema rename、term type 变化、scope 变化和 query layout 变化都会改变 resolved layout。 diff --git a/wiki/design/manager/runtime.md b/wiki/design/manager/runtime.md index 15378050..a5f3318a 100644 --- a/wiki/design/manager/runtime.md +++ b/wiki/design/manager/runtime.md @@ -60,7 +60,7 @@ Manager group 的顺序是行为契约的一部分:action slice、observation 3. 创建 action、command、observation、reward 和 termination terms; 4. 为每个 term 执行一次 `canonicalize_kernel_data`; 5. resolver(`ManagerProgramResolver`)构造 `ManagerContext` layout、read plan 和 - prepared invocations:term dispatch 以签名校验(`inspect` + 类型标注),不做任何 + per-term-kind typed call records:term dispatch 以签名校验(`inspect` + 类型标注),不做任何 代码生成或 dispatcher 编译; 6. Torch 执行器按元数据调度;首次真实调用完成惰性 specialization,`warmup()` 用克隆 输入执行一次 evaluate/observe 预热。 diff --git a/wiki/design/manager/torch-executor.md b/wiki/design/manager/torch-executor.md index 5515813b..fe5dc642 100644 --- a/wiki/design/manager/torch-executor.md +++ b/wiki/design/manager/torch-executor.md @@ -5,7 +5,7 @@ 本文定义 `ManagerEnv` 的数值执行器:通用 Torch 批量执行器 (`motrix_env_core.torch_manager.runtime.TorchManagerRuntime`)。resolver (`motrix_env_core.manager.resolver.ManagerProgramResolver`)把 manager 配置 -解析为 read plan / layout / prepared invocations 等执行元数据,Torch 执行器 +解析为 read plan / layout / typed term calls 等执行元数据,Torch 执行器 以整批张量按固定语句顺序调度 term;host 生命周期与 term 配置 API 不变。 Torch 是唯一执行器——历史上并存过的 Numba fused-kernel 路径与 `executor="numba"|"torch"` 双执行器切换已经移除。相关背景见 @@ -24,8 +24,9 @@ RNG draw 顺序与 eager 完全一致;Inductor 重结合仅带来 ≤1e-8 量 - `ManagerProgramResolver`(backend-neutral)负责输入准备:把 `@kernel_data` term 参数与 `ManagerContext` 解析为 path-keyed 的 resolved tree(叶子持有 live value 引用与 `PER_ENV`/`SHARED` scope),构建 `ManagerLayout`、 - `ReadPlan`、`PreparedInvocation`(每个 term 的 dispatcher 与 resolved - argument 描述符),并校验签名注解契约与签名契约。它不做任何代码生成、 + `ReadPlan`、per-term-kind typed call records(command hooks 携带 receiver 与 + static function,reward/termination 携带 buffer column index 与 resolved + argument 描述符等),并校验签名注解契约与签名契约。它不做任何代码生成、 缓存或 dispatcher 编译。 - `TorchManagerRuntime` 消费这些元数据并执行 term 调度;所有写入落回 host 持有的 numpy 缓冲(reward terms/masks/obs/command state/metrics), From 3dfd50b15efb7fb6936868d29d70cb64c40d4cfe Mon Sep 17 00:00:00 2001 From: zilch <147668916@qq.com> Date: Fri, 6 Nov 2026 02:13:10 +0800 Subject: [PATCH 23/41] refactor(torch): derive the reset clone surface from the write contract MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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). --- .../motrix_env_core/torch_manager/runtime.py | 29 +++++++++++++------ 1 file changed, 20 insertions(+), 9 deletions(-) diff --git a/motrix_env_core/src/motrix_env_core/torch_manager/runtime.py b/motrix_env_core/src/motrix_env_core/torch_manager/runtime.py index 18c4a8ff..6c6c9b17 100644 --- a/motrix_env_core/src/motrix_env_core/torch_manager/runtime.py +++ b/motrix_env_core/src/motrix_env_core/torch_manager/runtime.py @@ -34,7 +34,6 @@ import torch from motrix_env_core.manager.kernel import KernelOutputs -from motrix_env_core.manager.kernel_data import KernelDataScope from motrix_env_core.manager.resolver import ( CommandAdvanceHook, CommandResetHook, @@ -402,8 +401,14 @@ def __init__(self, env: ManagerEnv, program: ResolvedManagerProgram, compile: bo self._context_tree = program.prepared_trees[program.context.tree_index] self._context_leaves = self._context_tree.leaves - self._per_env_paths = frozenset( - path for path, leaf in self._context_leaves.items() if leaf.scope is KernelDataScope.PER_ENV + # Reset terms write through three surfaces only: each reset hook's + # receiver command tree, the RNG state, and metrics. Everything + # else (sim views, action state, clip arrays, shared params) is + # read-only during reset and passes through uncloned. + self._reset_clone_prefixes: tuple[tuple[str, ...], ...] = ( + *(("commands", hook.receiver_key) for hook in program.command_reset_hooks), + ("rand", "state"), + ("metrics",), ) self._sim_paths = {path[1]: path for path in self._context_leaves if len(path) == 2 and path[0] == "sim"} self._rand_path = ("rand", "state") @@ -600,8 +605,7 @@ def _build_compiled_reset( """ ctx = self._context_for(inputs) assert self._context_tensors is not None - per_env = self._per_env_paths - sim_paths = frozenset(self._sim_paths.values()) + clone_prefixes = self._reset_clone_prefixes get = _leaf_lookup(self._context_tensors) command_reset_hooks = self.command_reset_hooks num_envs = ctx.rand.state.shape[0] @@ -627,10 +631,17 @@ def lane_get(path: tuple[str, ...]) -> torch.Tensor: cached = clones.get(path) if cached is None: full = get(path) - # Simulator reads are never written by reset terms (writes - # go through the sim_writes map), so they pass through - # uncloned; every other per-env leaf may be mutated. - cached = full.clone() if path in per_env and path not in sim_paths else full + # Reset terms only write the receiver command trees, + # RNG state and metrics (sim writes go through the + # sim_writes map); every other leaf is read-only and + # passes through uncloned. + cached = ( + full.clone() + if isinstance(full, torch.Tensor) + and full.shape[:1] == (num_envs,) + and any(path[: len(prefix)] == prefix for prefix in clone_prefixes) + else full + ) clones[path] = cached return cached From 7db922c71e47f53530b3ea613c8484a41cc22ad3 Mon Sep 17 00:00:00 2001 From: zilch <147668916@qq.com> Date: Fri, 6 Nov 2026 11:06:35 +0800 Subject: [PATCH 24/41] refactor(manager)!: delete SharedArray and the kernel-data scope machinery 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. --- .../tutorial/building_envs/manager_env.md | 6 +- .../tutorial/building_envs/manager_env.md | 4 +- .../src/motrix_env_core/manager/__init__.py | 3 +- .../src/motrix_env_core/manager/actions.py | 4 +- .../src/motrix_env_core/manager/kernel.py | 15 +-- .../manager/kernel_data/__init__.py | 4 - .../manager/kernel_data/lowering.py | 116 ++++-------------- .../manager/kernel_data/resolved.py | 8 +- .../src/motrix_env_core/manager/resolver.py | 18 +-- .../src/motrix_env_core/mdp/action.py | 10 +- .../src/motrix_env_core/mdp/terrain.py | 18 +-- motrix_env_core/tests/test_kernel_data.py | 10 +- motrix_env_core/tests/test_manager_env.py | 8 +- .../locomotion/ball_balance/mdp/reset.py | 8 +- .../locomotion/ball_balance/mdp/rewards.py | 4 +- .../humanoid/walk_manager_mdp/command.py | 9 +- .../humanoid/walk_manager_mdp/reset.py | 15 ++- .../motrix_envs/locomotion/wbt/mdp/command.py | 7 +- .../locomotion/wbt/mdp/observations.py | 3 +- .../motrix_envs/locomotion/wbt/mdp/rewards.py | 6 +- motrix_envs/src/motrix_envs/motion/tracked.py | 32 ++--- wiki/design/manager/kernel-abi.md | 37 +++--- wiki/design/manager/runtime.md | 6 +- wiki/design/manager/task-authoring.md | 3 +- wiki/design/manager/torch-executor.md | 9 +- 25 files changed, 129 insertions(+), 234 deletions(-) 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 index 2fd8ac85..39fad3d9 100644 --- a/docs/source/en/user_guide/tutorial/building_envs/manager_env.md +++ b/docs/source/en/user_guide/tutorial/building_envs/manager_env.md @@ -84,8 +84,8 @@ distributed to the terms. 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 + (`np.ndarray` fields, one row per environment) and shared model data (scalar or + batch-free `np.ndarray` 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. @@ -103,7 +103,7 @@ class MyJointActionCfg(ActionCfg): @kernel_data class MyJointAction(ActionTerm): current: np.ndarray # persistent state: previous action (num_envs rows) - scale: SharedArray # static parameter: passed in by __call__ + scale: np.float32 # static parameter: passed in by __call__ def action_space(self, env, actuators) -> gym.spaces.Box: ... # derived from the actuators' ctrl_range 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 index 091b10c9..095f8207 100644 --- 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 @@ -73,7 +73,7 @@ step 循环内不再编译。 `actuator_names` 为空元组 `()` 时控制全部执行器,写显式名称时只控制列出的执行器; 一个执行器只能归属一项。设为 `None` 表示本项不写执行器目标。 2. `@kernel_data` 装饰的 `ActionTerm` 子类:持有持久状态(`np.ndarray` 字段,按环境行)与 - 共享模型数据(`SharedArray` 字段),实现三个方法:`action_space()` 给出本项的动作空间、 + 共享模型数据(无 batch 维的 `np.ndarray` / 标量字段),实现三个方法:`action_space()` 给出本项的动作空间、 `process(actions)` 把动作批次变换为路由执行器的控制量、`reset(env_ids)` 清理持久状态。 ```python @@ -90,7 +90,7 @@ class MyJointActionCfg(ActionCfg): @kernel_data class MyJointAction(ActionTerm): current: np.ndarray # 持久状态:上一动作缓存(num_envs 行) - scale: SharedArray # 静态参数:由 __call__ 传入 + scale: np.float32 # 静态参数:由 __call__ 传入 def action_space(self, env, actuators) -> gym.spaces.Box: ... # 由执行器 ctrl_range 推导 diff --git a/motrix_env_core/src/motrix_env_core/manager/__init__.py b/motrix_env_core/src/motrix_env_core/manager/__init__.py index a78f7061..a19d1a09 100644 --- a/motrix_env_core/src/motrix_env_core/manager/__init__.py +++ b/motrix_env_core/src/motrix_env_core/manager/__init__.py @@ -15,7 +15,7 @@ from motrix_env_core.manager.dispatch import dispatch from motrix_env_core.manager.env import ManagerBasedEnvCfg, ManagerEnv from motrix_env_core.manager.kernel import ManagerWarmupResult -from motrix_env_core.manager.kernel_data import SharedArray, kernel_data +from motrix_env_core.manager.kernel_data import kernel_data from motrix_env_core.manager.metrics import metric from motrix_env_core.manager.observations import ( ManagerObservationGroupCfg, @@ -61,7 +61,6 @@ "ManagerResetCfg", "ResetTerm", "ResetTermCfg", - "SharedArray", "TerminationManager", "TerminationTermCfg", "TerminationTerm", diff --git a/motrix_env_core/src/motrix_env_core/manager/actions.py b/motrix_env_core/src/motrix_env_core/manager/actions.py index f116e807..447e7024 100644 --- a/motrix_env_core/src/motrix_env_core/manager/actions.py +++ b/motrix_env_core/src/motrix_env_core/manager/actions.py @@ -12,7 +12,7 @@ from motrix_env_core.config import configclass from motrix_env_core.manager.dispatch import dispatch -from motrix_env_core.manager.kernel_data import SharedArray, canonicalize_kernel_data, kernel_data +from motrix_env_core.manager.kernel_data import canonicalize_kernel_data, kernel_data if TYPE_CHECKING: from motrix_env_core.manager.env import ManagerEnv @@ -31,7 +31,7 @@ class ActionState: """ action_queue: np.ndarray - action_ptr: SharedArray + action_ptr: np.ndarray @dispatch def current(self) -> np.ndarray: diff --git a/motrix_env_core/src/motrix_env_core/manager/kernel.py b/motrix_env_core/src/motrix_env_core/manager/kernel.py index 471fdf92..9c8d09d5 100644 --- a/motrix_env_core/src/motrix_env_core/manager/kernel.py +++ b/motrix_env_core/src/motrix_env_core/manager/kernel.py @@ -6,7 +6,6 @@ import numpy as np -from motrix_env_core.manager.kernel_data.lowering import KernelDataScope from motrix_env_core.manager.kernel_data.resolved import ResolvedLeaf, ResolvedTree T = TypeVar("T") @@ -44,20 +43,14 @@ def validate_kernel_context( """Validate the numeric kernel contract once when a context is bound. ``trees`` is the resolved program input: one :class:`ResolvedTree` per - prepared kernel-data source. Each leaf is validated against its declared - scope (``PER_ENV`` leaves must lead with the batch dimension). + prepared kernel-data source. Leaf values are checked for dtype and type + sanity without any batch-dimension assumptions. """ if not isinstance(trees, tuple): raise TypeError(f"inputs must be a tuple of resolved trees, got {type(trees).__name__}.") for tree in trees: for path, leaf in tree.leaves.items(): - _validate_value( - leaf.value, - f"inputs.{'.'.join(path)}", - num_envs=num_envs, - batched=leaf.scope is KernelDataScope.PER_ENV, - writable=False, - ) + _validate_value(leaf.value, f"inputs.{'.'.join(path)}", num_envs=num_envs, batched=False, writable=False) _validate_value(reward_weights, "reward_weights", num_envs=num_envs, batched=False, writable=False) _validate_value(buffers, "buffers", num_envs=num_envs, batched=True, writable=True) _validate_outputs(outputs, num_envs) @@ -74,7 +67,7 @@ def clone_resolved_trees(trees: tuple[ResolvedTree, ...]) -> tuple[ResolvedTree, return tuple( ResolvedTree( tree.value, - {path: ResolvedLeaf(clone_kernel_value(leaf.value), leaf.scope) for path, leaf in tree.leaves.items()}, + {path: ResolvedLeaf(clone_kernel_value(leaf.value)) for path, leaf in tree.leaves.items()}, ) for tree in trees ) diff --git a/motrix_env_core/src/motrix_env_core/manager/kernel_data/__init__.py b/motrix_env_core/src/motrix_env_core/manager/kernel_data/__init__.py index ae6780b7..d3adafc2 100644 --- a/motrix_env_core/src/motrix_env_core/manager/kernel_data/__init__.py +++ b/motrix_env_core/src/motrix_env_core/manager/kernel_data/__init__.py @@ -4,11 +4,9 @@ from motrix_env_core.manager.kernel_data.lowering import ( KernelDataLayout, KernelDataLowering, - KernelDataScope, KernelLeafLayout, KernelMapLayout, KernelRecordLayout, - SharedArray, iter_layout_leaves, ) from motrix_env_core.manager.kernel_data.map import Map @@ -33,7 +31,6 @@ "KernelDataLayout", "Map", "KernelDataLowering", - "KernelDataScope", "KernelLeafLayout", "KernelMapDef", "FieldDef", @@ -42,7 +39,6 @@ "KernelMapLayout", "KernelRecordLayout", "TreeClassDef", - "SharedArray", "ResolvedLeaf", "ResolvedTree", "canonicalize_kernel_data", diff --git a/motrix_env_core/src/motrix_env_core/manager/kernel_data/lowering.py b/motrix_env_core/src/motrix_env_core/manager/kernel_data/lowering.py index bfa2ade6..f8f90252 100644 --- a/motrix_env_core/src/motrix_env_core/manager/kernel_data/lowering.py +++ b/motrix_env_core/src/motrix_env_core/manager/kernel_data/lowering.py @@ -5,8 +5,7 @@ import hashlib from dataclasses import dataclass -from enum import Enum -from typing import Annotated, Any, TypeAlias, get_args, get_origin +from typing import Any import numpy as np @@ -21,23 +20,6 @@ ] = {} -class KernelDataScope(str, Enum): - """How one resolved KernelData leaf is exposed to execution. - - Attributes: - PER_ENV: Leaf is an ``(N, ...)`` array covering the whole batch; the - reset schedule clones it for selected lanes. - SHARED: Pass the complete leaf unchanged to every environment lane. - """ - - PER_ENV = "per_env" - SHARED = "shared" - - -SharedArray: TypeAlias = Annotated[np.ndarray, KernelDataScope.SHARED] -"""A writable ndarray passed unchanged to every environment lane.""" - - @dataclass(frozen=True) class KernelLeafLayout: """Lowered ABI description of one logical KernelData leaf. @@ -45,13 +27,11 @@ class KernelLeafLayout: Attributes: tree_def: Logical leaf definition that owns the path and annotation. value_type: Concrete Python/NumPy type accepted for the leaf value. - scope: Whether the leaf is per-environment batched or shared. slot_index: Zero-based declaration index relative to the root layout. """ tree_def: LeafDef value_type: type - scope: KernelDataScope slot_index: int @property @@ -108,8 +88,8 @@ class KernelRecordLayout: Attributes: tree_def: Logical record definition being lowered. fields: Lowered children in logical dataclass declaration order. - fingerprint: Deterministic identity including child types, scopes, - and dispatch-method code fingerprints. + fingerprint: Deterministic identity including child types and + dispatch-method code fingerprints. """ tree_def: TreeClassDef @@ -154,9 +134,9 @@ class KernelDataLowering: """Validate one logical ``TreeClassDef`` and derive its leaf layout. The lowering checks that every leaf annotation is a supported - ndarray/scalar type, assigns each leaf its data scope, and produces a - deterministic structural fingerprint (field names, types, scopes, and - dispatch-method code) used for resolved-program invalidation. + ndarray/scalar type and produces a deterministic structural fingerprint + (field names, types, and dispatch-method code) used for + resolved-program invalidation. """ def lower( @@ -164,14 +144,13 @@ def lower( tree_def: TreeClassDef, *, context: str, - force_shared: bool = False, ) -> KernelRecordLayout: method_key = hashlib.sha256(repr(_tree_method_fingerprint(tree_def)).encode()).hexdigest() - cache_key = (f"{tree_def.fingerprint}:{method_key}", force_shared) + cache_key = f"{tree_def.fingerprint}:{method_key}" cached = _LAYOUT_CACHE.get(cache_key) if cached is not None: return cached - layout, _ = self._lower_record(tree_def, 0, force_shared, context) + layout, _ = self._lower_record(tree_def, 0, context) _LAYOUT_CACHE[cache_key] = layout return layout @@ -179,7 +158,6 @@ def _lower_record( self, tree_def: TreeClassDef, slot_index: int, - force_shared: bool, context: str, ) -> tuple[KernelRecordLayout, int]: field_layouts: list[KernelFieldLayout] = [] @@ -187,33 +165,18 @@ def _lower_record( for field in tree_def.fields: child: KernelDataLayout if isinstance(field.tree_def, TreeClassDef): - child, slot_index = self._lower_record( - field.tree_def, - slot_index, - force_shared, - context, - ) + child, slot_index = self._lower_record(field.tree_def, slot_index, context) elif isinstance(field.tree_def, KernelMapDef): - child, slot_index = self._lower_map( - field.tree_def, - slot_index, - force_shared, - context, - ) + child, slot_index = self._lower_map(field.tree_def, slot_index, context) else: assert isinstance(field.tree_def, LeafDef) - value_type, scope = self._leaf_type( - field.tree_def.annotation, - context=f"{context} KernelData {tree_def.logical_type.__name__}.{'.'.join(field.tree_def.path)}", - force_shared=force_shared, - ) + value_type = self._leaf_type(field.tree_def.annotation) if value_type is None: raise TypeError( f"{context} KernelData {tree_def.logical_type.__name__}." - f"{'.'.join(field.tree_def.path)} expected " - "np.ndarray, SharedArray, or numeric scalar." + f"{'.'.join(field.tree_def.path)} expected np.ndarray or numeric scalar." ) - child = KernelLeafLayout(field.tree_def, value_type, scope, slot_index) + child = KernelLeafLayout(field.tree_def, value_type, slot_index) slot_index += 1 field_layouts.append(KernelFieldLayout(field.name, child)) fingerprint_fields.append((field.name, self._fingerprint_part(child))) @@ -232,39 +195,23 @@ def _lower_map( self, tree_def: KernelMapDef, slot_index: int, - force_shared: bool, context: str, ) -> tuple[KernelMapLayout, int]: entry_layouts: list[KernelMapEntryLayout] = [] fingerprint_entries: list[Any] = [] for entry in tree_def.entries: if isinstance(entry.tree_def, TreeClassDef): - child, slot_index = self._lower_record( - entry.tree_def, - slot_index, - force_shared, - context, - ) + child, slot_index = self._lower_record(entry.tree_def, slot_index, context) elif isinstance(entry.tree_def, KernelMapDef): - child, slot_index = self._lower_map( - entry.tree_def, - slot_index, - force_shared, - context, - ) + child, slot_index = self._lower_map(entry.tree_def, slot_index, context) else: assert isinstance(entry.tree_def, LeafDef) - value_type, scope = self._leaf_type( - entry.tree_def.annotation, - context=f"{context} Map {'.'.join(entry.tree_def.path)}", - force_shared=force_shared, - ) + value_type = self._leaf_type(entry.tree_def.annotation) if value_type is None: raise TypeError( - f"{context} Map {'.'.join(entry.tree_def.path)} expected np.ndarray, " - "SharedArray, or numeric scalar." + f"{context} Map {'.'.join(entry.tree_def.path)} expected np.ndarray or numeric scalar." ) - child = KernelLeafLayout(entry.tree_def, value_type, scope, slot_index) + child = KernelLeafLayout(entry.tree_def, value_type, slot_index) slot_index += 1 entry_layouts.append(KernelMapEntryLayout(entry.key, child)) fingerprint_entries.append((entry.key, self._fingerprint_part(child))) @@ -279,7 +226,6 @@ def _fingerprint_part(layout: KernelDataLayout) -> Any: "leaf", layout.path, f"{layout.value_type.__module__}.{layout.value_type.__qualname__}", - layout.scope.value, layout.slot_index, ) if isinstance(layout, KernelMapLayout): @@ -298,30 +244,16 @@ def _fingerprint_part(layout: KernelDataLayout) -> Any: ) @staticmethod - def _leaf_type( - annotation: Any, - *, - context: str, - force_shared: bool, - ) -> tuple[type[Any] | None, KernelDataScope]: - metadata: tuple[Any, ...] = () - if get_origin(annotation) is Annotated: - annotation, *metadata_values = get_args(annotation) - metadata = tuple(metadata_values) - if metadata and metadata != (KernelDataScope.SHARED,): - raise TypeError(f"{context} has unsupported or conflicting metadata {metadata!r}.") + def _leaf_type(annotation: Any) -> type[Any] | None: if annotation is np.ndarray: - scope = KernelDataScope.SHARED if metadata or force_shared else KernelDataScope.PER_ENV - return annotation, scope - if metadata: - raise TypeError(f"{context} scope metadata is only valid on np.ndarray leaves.") + return annotation if annotation in {bool, int, float}: - return annotation, KernelDataScope.SHARED + return annotation if isinstance(annotation, type) and issubclass(annotation, np.generic): try: dtype = np.dtype(annotation) except TypeError: - return None, KernelDataScope.SHARED + return None if dtype.type is annotation and dtype.kind in "biuf": - return annotation, KernelDataScope.SHARED - return None, KernelDataScope.SHARED + return annotation + return None diff --git a/motrix_env_core/src/motrix_env_core/manager/kernel_data/resolved.py b/motrix_env_core/src/motrix_env_core/manager/kernel_data/resolved.py index e3b29535..d59f1b75 100644 --- a/motrix_env_core/src/motrix_env_core/manager/kernel_data/resolved.py +++ b/motrix_env_core/src/motrix_env_core/manager/kernel_data/resolved.py @@ -9,8 +9,6 @@ from dataclasses import dataclass from typing import Any -from motrix_env_core.manager.kernel_data.lowering import KernelDataScope - @dataclass(frozen=True) class ResolvedLeaf: @@ -19,11 +17,9 @@ class ResolvedLeaf: Attributes: value: Live leaf value reference — an ndarray held and refreshed in place by its owning runtime object, or an immutable scalar. - scope: Whether the leaf is an ``(N, ...)`` batched array or shared. """ value: Any - scope: KernelDataScope @dataclass(frozen=True) @@ -32,8 +28,8 @@ class ResolvedTree: ``value`` keeps the host root object (used as the attribute fallback for executor namespace views); ``leaves`` maps every leaf path to its live - value reference and scope. Array leaves are rewritten in place by their - owning runtime objects, so the tree itself never needs rebuilding. + value reference. Array leaves are rewritten in place by their owning + runtime objects, so the tree itself never needs rebuilding. """ value: Any diff --git a/motrix_env_core/src/motrix_env_core/manager/resolver.py b/motrix_env_core/src/motrix_env_core/manager/resolver.py index 3562a46c..51dafd57 100644 --- a/motrix_env_core/src/motrix_env_core/manager/resolver.py +++ b/motrix_env_core/src/motrix_env_core/manager/resolver.py @@ -5,7 +5,7 @@ Lowers one constructed manager environment into its execution metadata: the simulator read plan, the resolved context/term-argument trees (path-keyed -leaf descriptors with live value references and data scopes), prepared term +leaf descriptors with live value references), prepared term arguments, and one typed call record per manager term (command hooks, observations, rewards, terminations, simulator resets). Numerical executors — currently the Torch executor in :mod:`motrix_env_core.torch_manager` — @@ -32,7 +32,6 @@ from motrix_env_core.manager.kernel_data import ( KernelDataLayout, KernelDataLowering, - KernelDataScope, Map, ResolvedLeaf, ResolvedTree, @@ -314,11 +313,10 @@ def _manager_sim_inputs(sim_data: PhysicsReadProgram) -> tuple[Mapping[str, Mana def _tree_leaf_fingerprint(tree: ResolvedTree) -> tuple[Any, ...]: - """Per-leaf content fingerprint: scope plus array shape/dtype identity.""" + """Per-leaf content fingerprint: array shape/dtype identity.""" return tuple( ( path, - leaf.scope.value, None if not isinstance(leaf.value, np.ndarray) else (leaf.value.shape, leaf.value.dtype.str, leaf.value.strides), @@ -702,16 +700,12 @@ def _resolve_term_arg(self, source_name: str, value: Any) -> PreparedArg: """ if is_kernel_data(value): _, tree_def = flatten_kernel_data(value) - layout = self._kernel_data_lowering.lower( - tree_def, - context=source_name, - force_shared=True, - ) + layout = self._kernel_data_lowering.lower(tree_def, context=source_name) return TreeArg(self._register_tree(source_name, value, layout)) if isinstance(value, (float, np.floating)): - return LeafArg(ResolvedLeaf(np.float32(value), KernelDataScope.SHARED)) + return LeafArg(ResolvedLeaf(np.float32(value))) if isinstance(value, np.ndarray): - return LeafArg(ResolvedLeaf(value, KernelDataScope.SHARED)) + return LeafArg(ResolvedLeaf(value)) if isinstance(value, SimDataQuery): # A query passed as a term argument resolves to its registered # read-program batched view; the dispatch receives the array in @@ -743,7 +737,7 @@ def _register_tree( f"{source_name} field {'.'.join(prepared_field.path)!r} must return " f"{prepared_field.value_type.__name__}, got {type(field_value).__name__}." ) - leaves[prepared_field.path] = ResolvedLeaf(field_value, prepared_field.scope) + leaves[prepared_field.path] = ResolvedLeaf(field_value) self._prepared_trees.append(ResolvedTree(value, leaves)) return len(self._prepared_trees) - 1 diff --git a/motrix_env_core/src/motrix_env_core/mdp/action.py b/motrix_env_core/src/motrix_env_core/mdp/action.py index 90ed3ed3..66e5105a 100644 --- a/motrix_env_core/src/motrix_env_core/mdp/action.py +++ b/motrix_env_core/src/motrix_env_core/mdp/action.py @@ -8,7 +8,7 @@ from motrix_env_core.config import configclass from motrix_env_core.config.scene import RobotCfg -from motrix_env_core.manager import ActionCfg, ActionState, ActionTerm, ManagerEnv, SharedArray, kernel_data +from motrix_env_core.manager import ActionCfg, ActionState, ActionTerm, ManagerEnv, kernel_data from motrix_env_core.mdp.action_space import joint_position_action_space_from_ctrl_ranges from motrix_env_core.sim.model import ActuatorSpec, ActuatorType @@ -41,10 +41,10 @@ class JointPositionActionState(ActionState): path entirely. """ - default_angles: SharedArray - joint_lower: SharedArray - joint_upper: SharedArray - action_scales: SharedArray + default_angles: np.ndarray + joint_lower: np.ndarray + joint_upper: np.ndarray + action_scales: np.ndarray delay_steps: np.ndarray delay_lo: int delay_hi: int diff --git a/motrix_env_core/src/motrix_env_core/mdp/terrain.py b/motrix_env_core/src/motrix_env_core/mdp/terrain.py index 0caff9c1..b3d80b28 100644 --- a/motrix_env_core/src/motrix_env_core/mdp/terrain.py +++ b/motrix_env_core/src/motrix_env_core/mdp/terrain.py @@ -6,7 +6,7 @@ import numpy as np import torch -from motrix_env_core.manager import SharedArray, kernel_data +from motrix_env_core.manager import kernel_data @kernel_data @@ -20,24 +20,24 @@ class HeightFieldGrid: world ``xy`` to fractional grid coordinates. Attributes: - heights: ``SharedArray`` of shape ``(nrow, ncol)`` holding the + heights: ``np.ndarray`` of shape ``(nrow, ncol)`` holding the height-field sample values (row = y 方向, column = x 方向), 单位与引擎导出的 hfield 数据一致。 - origin: 长度为 2 的 ``SharedArray``,网格首个采样点 ``(x, y)`` 的 + origin: 长度为 2 的 ``np.ndarray``,网格首个采样点 ``(x, y)`` 的 世界坐标,是 bilinear 插值前的平移原点。 - spacing: 长度为 2 的 ``SharedArray``,相邻采样点在世界坐标系下的 + spacing: 长度为 2 的 ``np.ndarray``,相邻采样点在世界坐标系下的 间距 ``(dx, dy)``,用于把世界坐标映射到分数网格坐标。 - z0: 长度为 1 的 ``SharedArray``,叠加在插值结果之上的基准高度偏移。 + z0: 长度为 1 的 ``np.ndarray``,叠加在插值结果之上的基准高度偏移。 constant: 平地退化模式下的常数高度;仅在 ``enabled=False`` 时生效, 对应平地 ground geom 的世界 z 坐标。 enabled: 是否启用网格插值;``False`` 时高度查找直接返回 ``constant``,忽略所有网格字段。 """ - heights: SharedArray - origin: SharedArray - spacing: SharedArray - z0: SharedArray + heights: np.ndarray + origin: np.ndarray + spacing: np.ndarray + z0: np.ndarray constant: np.float32 enabled: bool diff --git a/motrix_env_core/tests/test_kernel_data.py b/motrix_env_core/tests/test_kernel_data.py index ecb47b9b..a417aeb3 100644 --- a/motrix_env_core/tests/test_kernel_data.py +++ b/motrix_env_core/tests/test_kernel_data.py @@ -8,9 +8,7 @@ from motrix_env_core.manager.kernel_data import ( KernelDataLowering, - KernelDataScope, Map, - SharedArray, flatten_kernel_data, iter_layout_leaves, kernel_data, @@ -21,7 +19,7 @@ @kernel_data class _Pose: position: np.ndarray - orientation: SharedArray + orientation: np.ndarray @kernel_data @@ -67,12 +65,6 @@ def test_kernel_data_is_frozen_and_compiles_nested_inherited_leaves() -> None: ("robot", "pose", "orientation"), ("phase",), ] - assert [leaf.scope for leaf in leaves] == [ - KernelDataScope.PER_ENV, - KernelDataScope.PER_ENV, - KernelDataScope.SHARED, - KernelDataScope.SHARED, - ] def test_kernel_data_flatten_is_zero_copy_and_rebuilds_host_value() -> None: diff --git a/motrix_env_core/tests/test_manager_env.py b/motrix_env_core/tests/test_manager_env.py index 43144459..61445065 100644 --- a/motrix_env_core/tests/test_manager_env.py +++ b/motrix_env_core/tests/test_manager_env.py @@ -38,7 +38,6 @@ ObsTerm, RewardTerm, RewardTermCfg, - SharedArray, TerminationTerm, TerminationTermCfg, kernel_data, @@ -97,7 +96,7 @@ def _second_action_space(_: ManagerEnv, __: np.ndarray | None) -> gym.spaces.Box @kernel_data class _ObservationParams: - offset: SharedArray + offset: np.ndarray scale: np.float32 @@ -529,9 +528,8 @@ def test_manager_context_is_injected_once_and_reused_across_all_term_kinds() -> command = _counter_command(env) assert env._compiled_manager_program is not None context_tree = env._compiled_manager_program.prepared_trees[env._compiled_manager_program.context.tree_index] - scopes = {path: leaf.scope for path, leaf in context_tree.leaves.items()} - assert scopes[("actions", "test", "action_queue")].value == "per_env" - assert scopes[("actions", "test", "action_ptr")].value == "shared" + assert context_tree.leaves[("actions", "test", "action_queue")].value is action.state.action_queue + assert context_tree.leaves[("actions", "test", "action_ptr")].value is action.state.action_ptr assert [term.output_slice for term in env.manager_layout.observations["policy"].terms] == [ slice(0, 2), slice(2, 3), diff --git a/motrix_envs/src/motrix_envs/locomotion/ball_balance/mdp/reset.py b/motrix_envs/src/motrix_envs/locomotion/ball_balance/mdp/reset.py index c50ccb23..d83901b8 100644 --- a/motrix_envs/src/motrix_envs/locomotion/ball_balance/mdp/reset.py +++ b/motrix_envs/src/motrix_envs/locomotion/ball_balance/mdp/reset.py @@ -14,7 +14,7 @@ ResetTerm, ResetTermCfg, ) -from motrix_env_core.manager.kernel_data import Map, SharedArray, kernel_data +from motrix_env_core.manager.kernel_data import Map, kernel_data from motrix_env_core.math.quaternion_array import from_euler from motrix_env_core.sim import ( BodyAngularVelocityWrite, @@ -127,9 +127,9 @@ def __call__(self, ctx) -> ResetTerm: @kernel_data class DofResetParams: - default_pos: SharedArray - joint_lower: SharedArray - joint_upper: SharedArray + default_pos: np.ndarray + joint_lower: np.ndarray + joint_upper: np.ndarray noise_scale: np.float32 diff --git a/motrix_envs/src/motrix_envs/locomotion/ball_balance/mdp/rewards.py b/motrix_envs/src/motrix_envs/locomotion/ball_balance/mdp/rewards.py index 97737ef2..3ea69cc2 100644 --- a/motrix_envs/src/motrix_envs/locomotion/ball_balance/mdp/rewards.py +++ b/motrix_envs/src/motrix_envs/locomotion/ball_balance/mdp/rewards.py @@ -8,7 +8,7 @@ from motrix_env_core.config import configclass from motrix_env_core.manager import ManagerContext, RewardTerm, RewardTermCfg -from motrix_env_core.manager.kernel_data import SharedArray, kernel_data +from motrix_env_core.manager.kernel_data import kernel_data from motrix_env_core.math.quaternion_array import rotate_inverse as rotate_inverse_batch @@ -72,7 +72,7 @@ def __call__(self, ctx) -> RewardTerm: @kernel_data class DofDefaultParams: - reference: SharedArray + reference: np.ndarray sigma: np.float32 diff --git a/motrix_envs/src/motrix_envs/locomotion/humanoid/walk_manager_mdp/command.py b/motrix_envs/src/motrix_envs/locomotion/humanoid/walk_manager_mdp/command.py index 1f4060f9..45e88412 100644 --- a/motrix_envs/src/motrix_envs/locomotion/humanoid/walk_manager_mdp/command.py +++ b/motrix_envs/src/motrix_envs/locomotion/humanoid/walk_manager_mdp/command.py @@ -14,7 +14,6 @@ CommandTerm, ManagerContext, ManagerEnv, - SharedArray, kernel_data, metric, ) @@ -126,15 +125,15 @@ class WalkCommand(CommandTerm): the ``command_steps`` metric). """ - vel_limit_low: SharedArray - vel_limit_high: SharedArray + vel_limit_low: np.ndarray + vel_limit_high: np.ndarray stand_prob: np.float32 resample_steps: np.float32 gait_period_base: np.float32 gait_period_width: np.float32 curriculum_enabled: bool - penalty_scale: SharedArray - avg_ep_len: SharedArray + penalty_scale: np.ndarray + avg_ep_len: np.ndarray level_down_threshold: np.float32 level_up_threshold: np.float32 degree: np.float32 diff --git a/motrix_envs/src/motrix_envs/locomotion/humanoid/walk_manager_mdp/reset.py b/motrix_envs/src/motrix_envs/locomotion/humanoid/walk_manager_mdp/reset.py index a4dd0de8..bee2cbcb 100644 --- a/motrix_envs/src/motrix_envs/locomotion/humanoid/walk_manager_mdp/reset.py +++ b/motrix_envs/src/motrix_envs/locomotion/humanoid/walk_manager_mdp/reset.py @@ -11,7 +11,6 @@ ManagerContext, ResetTerm, ResetTermCfg, - SharedArray, kernel_data, ) from motrix_env_core.manager.context import BuildContext @@ -72,18 +71,18 @@ class WalkResetParams: base center of mass. """ - default_joint_angles: SharedArray - init_pose: SharedArray + default_joint_angles: np.ndarray + init_pose: np.ndarray heightfield: HeightFieldGrid spawn_range: np.float32 randomization_enabled: bool - kp_default: SharedArray - damping_default: SharedArray - link_mass_default: SharedArray + kp_default: np.ndarray + damping_default: np.ndarray + link_mass_default: np.ndarray base_link_index: np.int64 base_mass_default: np.float32 - base_com_default: SharedArray - friction_default: SharedArray + base_com_default: np.ndarray + friction_default: np.ndarray kp_range: np.ndarray damping_range: np.ndarray friction_range: np.ndarray diff --git a/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/command.py b/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/command.py index ef72ac33..2b9fc191 100644 --- a/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/command.py +++ b/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/command.py @@ -14,7 +14,6 @@ CommandTerm, ManagerContext, ManagerEnv, - SharedArray, kernel_data, metric, ) @@ -101,9 +100,9 @@ class WbtMotionCommand(CommandTerm): target_body_orientation_relative: np.ndarray # Host-side adaptive frame-sampling state and parameters. - adaptive_bin_failed_count: SharedArray - adaptive_current_bin_failed_count: SharedArray - sampling_cdf: SharedArray + adaptive_bin_failed_count: np.ndarray + adaptive_current_bin_failed_count: np.ndarray + sampling_cdf: np.ndarray start_at_timestep_zero_prob: np.float32 hold_at_clip_end: bool sequential_clips: bool diff --git a/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/observations.py b/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/observations.py index 3a5cfaaf..1017e8bf 100644 --- a/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/observations.py +++ b/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/observations.py @@ -14,7 +14,6 @@ ObsTerm, kernel_data, ) -from motrix_env_core.manager.kernel_data import SharedArray from motrix_env_core.math import quaternion_array from motrix_env_core.math.quaternion_array import relative_orientation, rotate_inverse from motrix_env_core.mdp.observations import ( @@ -121,7 +120,7 @@ def __call__(self, ctx) -> ObsTerm: @kernel_data class RelativePositionParams: - reference: SharedArray + reference: np.ndarray noise_amplitude: np.float32 diff --git a/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/rewards.py b/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/rewards.py index ff40192b..4cba9ac0 100644 --- a/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/rewards.py +++ b/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/rewards.py @@ -10,7 +10,7 @@ from motrix_env_core.config import configclass from motrix_env_core.manager import ManagerContext, RewardTerm, RewardTermCfg -from motrix_env_core.manager.kernel_data import SharedArray, kernel_data +from motrix_env_core.manager.kernel_data import kernel_data from motrix_env_core.math.quaternion_array import rotation_distance from motrix_env_core.mdp.action import JointPositionActionState @@ -148,8 +148,8 @@ def __call__(self, ctx) -> RewardTerm: @kernel_data class DofLimitParams: - midpoint: SharedArray - half_range: SharedArray + midpoint: np.ndarray + half_range: np.ndarray soft_limit: np.float32 cap: np.float32 diff --git a/motrix_envs/src/motrix_envs/motion/tracked.py b/motrix_envs/src/motrix_envs/motion/tracked.py index 74b1c481..d3b2398d 100644 --- a/motrix_envs/src/motrix_envs/motion/tracked.py +++ b/motrix_envs/src/motrix_envs/motion/tracked.py @@ -7,7 +7,7 @@ import numpy as np -from motrix_env_core.manager.kernel_data import Map, SharedArray, kernel_data +from motrix_env_core.manager.kernel_data import Map, kernel_data from motrix_envs.motion.loader import MotionFormatError, MotrixMotion @@ -42,11 +42,11 @@ def load_extension_channel(motion: MotrixMotion, name: str) -> np.ndarray: class MotionChannel: """One declared ``ext_`` channel, aligned with the clip's frame axis. - Map leaves cannot carry ``SharedArray`` annotations directly, so each + Map leaves are flattened by value type, so each channel is wrapped in this record whose field is shared across lanes. """ - data: SharedArray + data: np.ndarray @kernel_data @@ -80,19 +80,19 @@ class WbtMotionClip: axis; empty when the task declares none. """ - joint_pos: SharedArray - joint_vel: SharedArray - tracked_bodies_pos_w: SharedArray - tracked_bodies_quat_w: SharedArray - tracked_bodies_lin_vel_w: SharedArray - tracked_bodies_ang_vel_w: SharedArray - root_body_pos_w: SharedArray - root_body_quat_w: SharedArray - root_body_lin_vel_w: SharedArray - root_body_ang_vel_w: SharedArray - reference_body_pos_w: SharedArray - reference_body_quat_w: SharedArray - frame_clip_end: SharedArray + joint_pos: np.ndarray + joint_vel: np.ndarray + tracked_bodies_pos_w: np.ndarray + tracked_bodies_quat_w: np.ndarray + tracked_bodies_lin_vel_w: np.ndarray + tracked_bodies_ang_vel_w: np.ndarray + root_body_pos_w: np.ndarray + root_body_quat_w: np.ndarray + root_body_lin_vel_w: np.ndarray + root_body_ang_vel_w: np.ndarray + reference_body_pos_w: np.ndarray + reference_body_quat_w: np.ndarray + frame_clip_end: np.ndarray extensions: Map @staticmethod diff --git a/wiki/design/manager/kernel-abi.md b/wiki/design/manager/kernel-abi.md index 817cbf9e..9d028fdd 100644 --- a/wiki/design/manager/kernel-abi.md +++ b/wiki/design/manager/kernel-abi.md @@ -4,7 +4,7 @@ 本文定义 Manager term 的数据 ABI:`@kernel_data` logical tree 如何解析为 path-keyed 的 resolved tree,`ManagerContext` 如何以整批张量 -提供固定 namespace,以及 `PER_ENV`、`SHARED` 和 metric backing 的 ownership 语义。任务开发者只需遵守 +提供固定 namespace,以及 metric backing 的 ownership 语义。任务开发者只需遵守 [Manager Task API](./task-authoring.md) 中的 term 协议;数值执行器见 [Torch 执行器设计](./torch-executor.md), backend 读侧和 simulator 边界见 [SimBackend](./sim-backend.md)。 @@ -26,7 +26,7 @@ class ManagerContext: Context 不包含 config、provider、backend handle、dynamic Python container 或 generic `states` / `resources` store。 Persistent runtime data 直接归属 concrete term;framework-owned random state 通过 `rand` 暴露,`dt` 提供当前 -`ctrl_dt` 的固定 float32 标量。执行器把该结构解析为固定 layout(path → leaf 引用 + scope),运行时以 namespace 视图(`ctx.sim[key]`、 +`ctrl_dt` 的固定 float32 标量。执行器把该结构解析为固定 layout(path → leaf 引用),运行时以 namespace 视图(`ctx.sim[key]`、 `ctx.commands[name].field` 等)暴露整批张量。 ## 2. 批量语义 @@ -37,11 +37,12 @@ Host 侧 `env_id` 是: np.arange(num_envs, dtype=np.int64) ``` -所有 `PER_ENV` leaves 是 `(N, ...)` 整批数组;term 一次调用处理整批。partial reset 采用静态全批 masked 设计: -reset 数学对全部 lane 计算,选中的 lane 由布尔掩码选择,`PER_ENV` leaves 按 clone-on-first-access 修改后以 +term 一次调用处理整批:lane 形状的 leaves 是 `(N, ...)` 整批数组。partial reset 采用静态全批 masked 设计: +reset 数学对全部 lane 计算,选中的 lane 由布尔掩码选择,reset term 可写的 leaves(reset hook 的 receiver +command tree、RNG state 与 lane 形状的 metric)按 clone-on-first-access 修改后以 `where(mask, clone, original)` 合并回常驻缓冲,未选中 lane 保持不变,不会把 reset batch row 当作 environment id。 -所有 Map 的 key 顺序、concrete value schema、field tree、leaf dtype、scope 和 resolved path 在 resolve 前固定。Map lookup 使用 literal key: +所有 Map 的 key 顺序、concrete value schema、field tree、leaf dtype 和 resolved path 在 resolve 前固定。Map lookup 使用 literal key: ```python motion = ctx.commands["motion"] @@ -50,17 +51,16 @@ dof_pos = ctx.sim["robot_dof_pos"] 动态字符串 key、运行期插入/删除和 fallback lookup 不属于 ABI;错误 key 在 resolve 或首次执行阶段失败。 -## 3. Leaf scope +## 3. Leaf 类型 -| 类型 | 默认 scope | term 访问语义 | -| --- | --- | --- | -| `np.ndarray` | `PER_ENV` | 整批 `(N, ...)` 数组/张量;reset 场景为 masked clone 行 | -| `SharedArray` | `SHARED` | 完整可写 ndarray,所有环境共享同一份 | -| Python scalar | `SHARED` | 所有环境使用同一个标量 | -| 固定 dtype NumPy scalar | `SHARED` | 所有环境使用同一个定 dtype 标量 | +| 类型 | term 访问语义 | +| --- | --- | +| `np.ndarray` | 完整 ndarray 引用,term 按自己的契约决定第 0 维是否为 batch 维 | +| Python scalar | 所有环境使用同一个标量 | +| 固定 dtype NumPy scalar | 所有环境使用同一个定 dtype 标量 | -`SharedArray` 只描述 scope,不代表 immutable。数组是否可以更新由 concrete term 的 owner/lifecycle 契约决定。 -Term receiver 的参数通常按 shared lowering 处理,以免把静态参数误切片。 +数组是否可以更新由 concrete term 的 owner/lifecycle 契约决定;reset term 可写的 leaves 由 reset write contract +(reset hook 的 receiver command tree、RNG state 与 lane 形状的 metric)界定。 ## 4. Logical tree 与 resolved tree @@ -74,8 +74,7 @@ rebuilt = unflatten_kernel_data(tree_def, leaves) Record 可以嵌套另一个 KernelData record;不支持 multiple inheritance、覆盖 inherited field、普通 dataclass、递归结构、dict、list、字符串、object array 或 backend handle。 Logical tree 与 execution lowering 分离:同一 concrete value 在 host term 和 context 中保留相同的 tree schema, -但 resolver 根据使用位置决定 scope 并产出 path-keyed 的 resolved tree(leaf 引用 + scope),运行时据此构建 -namespace 视图。只有 ndarray/scalar leaves 可以成为 resolved leaves;term 内不应产生稳态临时大数组分配。 +resolver 产出 path-keyed 的 resolved tree(leaf 引用),运行时据此构建 namespace 视图。只有 ndarray/scalar leaves 可以成为 resolved leaves;term 内不应产生稳态临时大数组分配。 ## 5. Ownership 与 metric @@ -92,20 +91,20 @@ Term 可以读取其他 term 的数据,但不应调用其他 term 的行为方 `ManagerProgramResolver`(backend-neutral)负责: -- 校验 concrete runtime terms 和 `ManagerContext` 的 KernelData schema(leaf 类型与 scope); +- 校验 concrete runtime terms 和 `ManagerContext` 的 KernelData schema(leaf 类型); - 收集 read program 的 sim input map; - 把 context 与 term 参数解析为 path-keyed resolved trees 和参数描述符(常量/叶子/树/sim key); - 校验方法签名、注解契约和 layout; - 产出 `ManagerLayout`、`ReadPlan`、resolved trees 与 per-term-kind typed call records,供 Torch 执行器调度。 构造期 warm-up 以克隆输入执行一次 evaluate/observe,确认输入/输出契约(shape、dtype、batch 维、可写性); -schema rename、term type 变化、scope 变化和 query layout 变化都会改变 resolved layout。 +schema rename、term type 变化和 query layout 变化都会改变 resolved layout。 ## 7. 核心不变量 - Context schema 在一个 environment 生命周期内固定; - 每个 key 对应唯一 concrete type 和 stable resolved path; -- 所有 `PER_ENV` leaves 共享同一个第 0 维(整批为 N;reset 数学以 masked 全批形式运行); +- lane 形状的 leaves 共享同一个第 0 维(整批为 N;reset 数学以 masked 全批形式运行); - `ctx.env_id` 在 step/reset 中与其他 per-env data 对齐; - simulator snapshots 只来自编译后的 read program; - 不建立 generic `states`、`resources` 或 `ManagerValue` registry; diff --git a/wiki/design/manager/runtime.md b/wiki/design/manager/runtime.md index a5f3318a..a2b85b13 100644 --- a/wiki/design/manager/runtime.md +++ b/wiki/design/manager/runtime.md @@ -82,7 +82,7 @@ reset 只处理指定的原始 `env_ids`: - action/command/reset terms 只修改目标 rows; - backend reset program 只写目标 rows; - read program 只刷新目标 rows; -- reset 路径按原始 environment id gather `PER_ENV` leaves,不把 reset batch row 当作环境 id; +- reset 路径按原始 environment id gather lane 状态,不把 reset batch row 当作环境 id; - `ArrayEnv` 继续负责 auto-reset、episode steps 和 truncation 语义。 ## 4. Term 协议 @@ -93,7 +93,7 @@ Term 不要求继承统一实现基类,但必须遵循对应 protocol: @kernel_data class BaseActionState: action_queue: np.ndarray # host (N, W, A), lane (W, A); W >= 2 - action_ptr: SharedArray # one shared cursor, advanced once per control step + action_ptr: np.ndarray # one shared cursor, advanced once per control step @dataclass(frozen=True) @@ -183,5 +183,5 @@ Observation terms use the host-side ``ObservationTerm(size, dispatch, *args)`` f - config 不持有 runtime state,concrete term 是其 persistent state 的唯一 owner; - action policy layout 与 actuator route 分离,route 非空且不重叠; - sim snapshot 只来自编译后的 read program; -- `ctx.env_id` 和所有 `PER_ENV` leaves 使用同一原始 environment id; +- `ctx.env_id` 和所有 lane 形状 leaves 使用同一原始 environment id; - Manager 不引入 generic `ManagerValue` registry 或独立 Hydra backend 维度。 diff --git a/wiki/design/manager/task-authoring.md b/wiki/design/manager/task-authoring.md index b79181f2..043c5746 100644 --- a/wiki/design/manager/task-authoring.md +++ b/wiki/design/manager/task-authoring.md @@ -92,8 +92,7 @@ class TrackingErrorTermination: `@kernel_data` 的字段必须有类型标注,结构在构造期固定。Term record 不可替换字段,但 owner 可以按契约原地更新 ndarray backing。 支持的 leaf 类型为: -- 普通 `np.ndarray`:默认按 environment 切片,第一维为 `num_envs`; -- `SharedArray`:所有 lane 访问完整数组,适合模型参数、查表数据和 motion clip; +- 普通 `np.ndarray`:term 拿到完整数组引用;lane 形状数组第一维为 `num_envs`,模型参数、查表数据和 motion clip 等共享数组则没有 batch 维; - Python 标量和固定 dtype NumPy scalar; - 另一个 `@kernel_data` 类型。 diff --git a/wiki/design/manager/torch-executor.md b/wiki/design/manager/torch-executor.md index fe5dc642..ba099dd5 100644 --- a/wiki/design/manager/torch-executor.md +++ b/wiki/design/manager/torch-executor.md @@ -23,7 +23,7 @@ RNG draw 顺序与 eager 完全一致;Inductor 重结合仅带来 ≤1e-8 量 - `ManagerProgramResolver`(backend-neutral)负责输入准备:把 `@kernel_data` term 参数与 `ManagerContext` 解析为 path-keyed 的 resolved tree(叶子持有 - live value 引用与 `PER_ENV`/`SHARED` scope),构建 `ManagerLayout`、 + live value 引用),构建 `ManagerLayout`、 `ReadPlan`、per-term-kind typed call records(command hooks 携带 receiver 与 static function,reward/termination 携带 buffer column index 与 resolved argument 描述符等),并校验签名注解契约与签名契约。它不做任何代码生成、 @@ -46,9 +46,10 @@ RNG draw 顺序与 eager 完全一致;Inductor 重结合仅带来 ≤1e-8 量 `_TensorCache` 在同一 anchor 上按叶子值身份缓存 tensor 视图,0 维 int/float/bool 叶子转为 Python 标量(graph-break-free 编译的关键)。 - reset 路径是静态全批 masked 设计:选中的 lane 由布尔掩码选择, - `PER_ENV` 叶子在 reset term 首次访问时 clone(sim 叶子除外),term 写 - clone 行后以 `where(mask, clone, original)` 合并回常驻缓冲,未选中 lane - (含 RNG 状态)保持不变;compact sim-reset write 缓冲布局与 + reset write contract 覆盖的可写叶子(reset hook 的 receiver command + tree、RNG state 与 lane 形状的 metric)在 reset term 首次访问时 clone, + term 写 clone 行后以 `where(mask, clone, original)` 合并回常驻缓冲, + 未选中 lane(含 RNG 状态)保持不变;compact sim-reset write 缓冲布局与 `SimResetRuntime.buffers` 一致。 ## 批量 term 约定 From ed79068a21b94acebfeb503b226d13470412e525 Mon Sep 17 00:00:00 2001 From: zilch <147668916@qq.com> Date: Fri, 6 Nov 2026 12:15:04 +0800 Subject: [PATCH 25/41] refactor(manager)!: drop @kernel_data and own carrier fields as torch tensors MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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). --- .../tutorial/building_envs/manager_env.md | 12 +-- .../tutorial/building_envs/manager_env.md | 10 +- .../src/motrix_env_core/manager/__init__.py | 2 - .../src/motrix_env_core/manager/actions.py | 27 +++--- .../src/motrix_env_core/manager/commands.py | 10 +- .../src/motrix_env_core/manager/context.py | 10 +- .../src/motrix_env_core/manager/kernel.py | 10 +- .../manager/kernel_data/__init__.py | 2 - .../manager/kernel_data/canonical.py | 12 ++- .../manager/kernel_data/lowering.py | 14 +-- .../manager/kernel_data/tree.py | 64 +++++-------- .../src/motrix_env_core/manager/metrics.py | 17 +++- .../src/motrix_env_core/manager/rand.py | 17 ++-- .../src/motrix_env_core/manager/resolver.py | 24 +++-- .../src/motrix_env_core/manager/terms.py | 4 +- .../src/motrix_env_core/mdp/action.py | 46 +++++---- .../src/motrix_env_core/mdp/state.py | 8 +- .../src/motrix_env_core/mdp/terrain.py | 14 +-- .../motrix_env_core/torch_manager/runtime.py | 27 +++++- motrix_env_core/tests/test_action_state.py | 84 +++++++++-------- motrix_env_core/tests/test_kernel_data.py | 54 +++++------ motrix_env_core/tests/test_manager_env.py | 89 +++++++++--------- .../tests/test_manager_sim_backend.py | 14 +-- .../tests/test_torch_manager_actions.py | 56 +++++------ .../locomotion/ball_balance/mdp/reset.py | 18 ++-- .../locomotion/ball_balance/mdp/rewards.py | 9 +- .../humanoid/walk_manager_mdp/command.py | 80 ++++++++-------- .../humanoid/walk_manager_mdp/reset.py | 93 ++++++++++--------- .../humanoid/walk_manager_mdp/terrain.py | 17 ++-- .../motrix_envs/locomotion/wbt/mdp/command.py | 76 ++++++++------- .../locomotion/wbt/mdp/observations.py | 9 +- .../motrix_envs/locomotion/wbt/mdp/rewards.py | 8 +- motrix_envs/src/motrix_envs/motion/library.py | 9 +- motrix_envs/src/motrix_envs/motion/tracked.py | 70 +++++++------- motrix_envs/tests/test_motion_library.py | 5 +- motrix_envs/tests/test_motion_loader.py | 4 +- motrix_envs/tests/test_walk_randomization.py | 47 +++++----- .../manager/isaaclab-manager-cfg-adapter.md | 16 ++-- wiki/design/manager/kernel-abi.md | 20 ++-- wiki/design/manager/runtime.md | 8 +- wiki/design/manager/task-authoring.md | 20 ++-- wiki/design/manager/term-required-queries.md | 4 +- wiki/design/manager/torch-executor.md | 6 +- wiki/design/motion-multi-clip-infra.md | 2 +- 44 files changed, 624 insertions(+), 524 deletions(-) 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 index 39fad3d9..72306540 100644 --- a/docs/source/en/user_guide/tutorial/building_envs/manager_env.md +++ b/docs/source/en/user_guide/tutorial/building_envs/manager_env.md @@ -83,9 +83,9 @@ distributed to the terms. 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 (scalar or - batch-free `np.ndarray` fields), and implements three methods — `action_space()` returns this term's action +2. An `ActionTerm` subclass — a plain frozen dataclass: it holds persistent state + (`torch.Tensor` fields, one row per environment) and shared model data (scalar or + batch-free tensor 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. @@ -100,9 +100,9 @@ class MyJointActionCfg(ActionCfg): actuators=actuators) -@kernel_data +@dataclass(frozen=True, slots=True) class MyJointAction(ActionTerm): - current: np.ndarray # persistent state: previous action (num_envs rows) + current: torch.Tensor # persistent state: previous action (num_envs rows) scale: np.float32 # static parameter: passed in by __call__ def action_space(self, env, actuators) -> gym.spaces.Box: @@ -248,7 +248,7 @@ through `env.state.metrics` / `env.metrics`. Most tasks can reuse built-in comma 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 +returns a `CommandTerm` subclass — a plain frozen dataclass — implementing four callbacks: | Callback | Runs in | Responsibility | 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 index 095f8207..5a4d44f4 100644 --- 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 @@ -72,8 +72,8 @@ step 循环内不再编译。 1. `ActionCfg` 子类:声明静态参数与执行器路由,`__call__(env, actuators)` 返回运行时 term。 `actuator_names` 为空元组 `()` 时控制全部执行器,写显式名称时只控制列出的执行器; 一个执行器只能归属一项。设为 `None` 表示本项不写执行器目标。 -2. `@kernel_data` 装饰的 `ActionTerm` 子类:持有持久状态(`np.ndarray` 字段,按环境行)与 - 共享模型数据(无 batch 维的 `np.ndarray` / 标量字段),实现三个方法:`action_space()` 给出本项的动作空间、 +2. plain frozen dataclass 的 `ActionTerm` 子类:持有持久状态(`torch.Tensor` 字段,按环境行)与 + 共享模型数据(无 batch 维的 tensor / 标量字段),实现三个方法:`action_space()` 给出本项的动作空间、 `process(actions)` 把动作批次变换为路由执行器的控制量、`reset(env_ids)` 清理持久状态。 ```python @@ -87,9 +87,9 @@ class MyJointActionCfg(ActionCfg): actuators=actuators) -@kernel_data +@dataclass(frozen=True, slots=True) class MyJointAction(ActionTerm): - current: np.ndarray # 持久状态:上一动作缓存(num_envs 行) + current: torch.Tensor # 持久状态:上一动作缓存(num_envs 行) scale: np.float32 # 静态参数:由 __call__ 传入 def action_space(self, env, actuators) -> gym.spaces.Box: @@ -225,7 +225,7 @@ class BodyPosResetCfg(ResetTermCfg): 不需要自定义。 **如何定义**:一个命令 term 是 `CommandCfg` 子类,`__call__(env)` 返回 -`@kernel_data` 的 `CommandTerm` 子类,实现四个回调: +plain frozen dataclass 的 `CommandTerm` 子类,实现四个回调: | 回调 | 执行位置 | 职责 | | ------------------------ | ----------------- | ------------------------------------------ | diff --git a/motrix_env_core/src/motrix_env_core/manager/__init__.py b/motrix_env_core/src/motrix_env_core/manager/__init__.py index a19d1a09..0bf870fa 100644 --- a/motrix_env_core/src/motrix_env_core/manager/__init__.py +++ b/motrix_env_core/src/motrix_env_core/manager/__init__.py @@ -15,7 +15,6 @@ from motrix_env_core.manager.dispatch import dispatch from motrix_env_core.manager.env import ManagerBasedEnvCfg, ManagerEnv from motrix_env_core.manager.kernel import ManagerWarmupResult -from motrix_env_core.manager.kernel_data import kernel_data from motrix_env_core.manager.metrics import metric from motrix_env_core.manager.observations import ( ManagerObservationGroupCfg, @@ -64,6 +63,5 @@ "TerminationManager", "TerminationTermCfg", "TerminationTerm", - "kernel_data", "metric", ] diff --git a/motrix_env_core/src/motrix_env_core/manager/actions.py b/motrix_env_core/src/motrix_env_core/manager/actions.py index 447e7024..f0353089 100644 --- a/motrix_env_core/src/motrix_env_core/manager/actions.py +++ b/motrix_env_core/src/motrix_env_core/manager/actions.py @@ -4,25 +4,30 @@ from __future__ import annotations import abc -from dataclasses import fields +from dataclasses import dataclass, fields from typing import TYPE_CHECKING import gymnasium as gym import numpy as np +import torch from motrix_env_core.config import configclass from motrix_env_core.manager.dispatch import dispatch -from motrix_env_core.manager.kernel_data import canonicalize_kernel_data, kernel_data +from motrix_env_core.manager.kernel_data import canonicalize_kernel_data if TYPE_CHECKING: from motrix_env_core.manager.env import ManagerEnv from motrix_env_core.sim.model import ActuatorSpec -@kernel_data +@dataclass(frozen=True, slots=True) class ActionState: """Manager-owned raw-action history shared by every action type. + Array fields are zero-copy torch tensors over host-owned numpy arenas, so + host code may write through ``tensor.numpy()`` views while kernels write + through the tensors — both see one backing buffer. + Attributes: action_queue: Raw policy actions, host ``(N, W, A)`` and kernel lane ``(W, A)``; ``W >= 2`` retains current and previous actions. @@ -30,25 +35,25 @@ class ActionState: slot. Only the manager advances it; partial reset leaves it intact. """ - action_queue: np.ndarray - action_ptr: np.ndarray + action_queue: torch.Tensor + action_ptr: torch.Tensor @dispatch - def current(self) -> np.ndarray: + def current(self) -> torch.Tensor: """Return current raw actions: host ``(N, A)``, kernel lane ``(A,)``. The returned view shares history storage and does not advance the cursor. """ - return self.action_queue[..., self.action_ptr[0], :] + return self.action_queue[:, self.action_ptr[0], :] @dispatch - def previous(self) -> np.ndarray: + def previous(self) -> torch.Tensor: """Return previous raw actions: host ``(N, A)``, kernel lane ``(A,)``. The returned view wraps around the ring without copying data. """ ptr = (self.action_ptr[0] - 1) % self.action_queue.shape[-2] - return self.action_queue[..., ptr, :] + return self.action_queue[:, ptr, :] class ActionTerm(abc.ABC): @@ -67,7 +72,7 @@ def process(self, actions: np.ndarray) -> np.ndarray | None: state = self.state ptr = (int(state.action_ptr[0]) + 1) % state.action_queue.shape[1] state.action_ptr[0] = ptr - state.action_queue[:, ptr] = actions + state.action_queue.numpy()[:, ptr] = actions return self._process(actions) @abc.abstractmethod @@ -76,7 +81,7 @@ def _process(self, actions: np.ndarray) -> np.ndarray | None: def reset(self, env_ids: np.ndarray) -> None: """Clear raw history and reset action-specific state for selected rows.""" - self.state.action_queue[env_ids] = 0.0 + self.state.action_queue.numpy()[env_ids] = 0.0 self._reset(env_ids) @abc.abstractmethod diff --git a/motrix_env_core/src/motrix_env_core/manager/commands.py b/motrix_env_core/src/motrix_env_core/manager/commands.py index 77b23b3d..ebbea044 100644 --- a/motrix_env_core/src/motrix_env_core/manager/commands.py +++ b/motrix_env_core/src/motrix_env_core/manager/commands.py @@ -8,10 +8,10 @@ from typing import TYPE_CHECKING, Any import numpy as np +import torch from motrix_env_core.config import configclass from motrix_env_core.manager.dispatch import dispatch -from motrix_env_core.manager.kernel_data import kernel_data if TYPE_CHECKING: from motrix_env_core.manager.context import ManagerContext @@ -27,18 +27,18 @@ class ResetContext: metrics: dict[str, Any] -@kernel_data +@dataclass(frozen=True, slots=True) class CommandTerm(abc.ABC): - """Environment-local KernelData command pipeline with persistent runtime state. + """Environment-local kernel-data command pipeline with persistent runtime state. Contract: every command term carries a ``command`` field — one - ``(num_envs, command_dim)`` float array holding the lane's goal vector. + ``(num_envs, command_dim)`` float tensor holding the lane's goal vector. The kernel lowering hands each lane a writable row view, and generic consumers (for example ``mdp.observations.CommandObsCfg``) read it by field name. """ - command: np.ndarray + command: torch.Tensor @dispatch @abc.abstractmethod diff --git a/motrix_env_core/src/motrix_env_core/manager/context.py b/motrix_env_core/src/motrix_env_core/manager/context.py index 75871f42..c093fe74 100644 --- a/motrix_env_core/src/motrix_env_core/manager/context.py +++ b/motrix_env_core/src/motrix_env_core/manager/context.py @@ -5,11 +5,13 @@ from __future__ import annotations +from dataclasses import dataclass from typing import TYPE_CHECKING, Any, TypeVar import numpy as np +import torch -from motrix_env_core.manager.kernel_data import Map, kernel_data +from motrix_env_core.manager.kernel_data import Map from motrix_env_core.manager.rand import RandValue if TYPE_CHECKING: @@ -20,11 +22,11 @@ T = TypeVar("T") -@kernel_data +@dataclass(frozen=True, slots=True) class ManagerContext: """Framework-owned manager data available to every compiled term.""" - env_id: np.ndarray + env_id: torch.Tensor actions: Map commands: Map metrics: Map[np.ndarray] @@ -35,7 +37,7 @@ class ManagerContext: # request simulator-state rematerialization for that lane after the # transition. It is not an episode reset: episode bookkeeping and # action-term state are untouched (see CommandTerm.advance). - sim_reset_requested: np.ndarray + sim_reset_requested: torch.Tensor class BuildContext: diff --git a/motrix_env_core/src/motrix_env_core/manager/kernel.py b/motrix_env_core/src/motrix_env_core/manager/kernel.py index 9c8d09d5..15df1362 100644 --- a/motrix_env_core/src/motrix_env_core/manager/kernel.py +++ b/motrix_env_core/src/motrix_env_core/manager/kernel.py @@ -5,6 +5,7 @@ from typing import Any, NamedTuple, TypeVar, cast import numpy as np +import torch from motrix_env_core.manager.kernel_data.resolved import ResolvedLeaf, ResolvedTree @@ -74,7 +75,9 @@ def clone_resolved_trees(trees: tuple[ResolvedTree, ...]) -> tuple[ResolvedTree, def clone_kernel_value(value: T) -> T: - """Recursively clone arrays while preserving task-specific tuple types.""" + """Recursively clone arrays and tensors while preserving task-specific tuple types.""" + if isinstance(value, torch.Tensor): + return cast(T, value.clone()) if isinstance(value, np.ndarray): # Warmup must use the same Numba array type as real execution: copy() # normalizes strided views to C layout and drops the readonly flag. @@ -98,6 +101,11 @@ def clone_kernel_value(value: T) -> T: def _validate_value(value: Any, path: str, *, num_envs: int, batched: bool, writable: bool) -> None: + if isinstance(value, torch.Tensor): + if batched and (value.dim() == 0 or value.shape[0] != num_envs): + raise ValueError(f"{path} must have batch dimension {num_envs}, got shape {tuple(value.shape)}.") + return + if isinstance(value, np.ndarray): if value.dtype.hasobject: raise TypeError(f"{path} must not use object dtype, got {value.dtype}.") diff --git a/motrix_env_core/src/motrix_env_core/manager/kernel_data/__init__.py b/motrix_env_core/src/motrix_env_core/manager/kernel_data/__init__.py index d3adafc2..d6ef1772 100644 --- a/motrix_env_core/src/motrix_env_core/manager/kernel_data/__init__.py +++ b/motrix_env_core/src/motrix_env_core/manager/kernel_data/__init__.py @@ -22,7 +22,6 @@ is_kernel_data, is_kernel_data_type, iter_tree_leaves, - kernel_data, unflatten_kernel_data, ) @@ -47,6 +46,5 @@ "is_kernel_data_type", "iter_layout_leaves", "iter_tree_leaves", - "kernel_data", "unflatten_kernel_data", ] diff --git a/motrix_env_core/src/motrix_env_core/manager/kernel_data/canonical.py b/motrix_env_core/src/motrix_env_core/manager/kernel_data/canonical.py index 1918df79..541be3f7 100644 --- a/motrix_env_core/src/motrix_env_core/manager/kernel_data/canonical.py +++ b/motrix_env_core/src/motrix_env_core/manager/kernel_data/canonical.py @@ -1,9 +1,10 @@ -# Copyright Motphys Technology Co., Ltd. 2025, 2026 +# Copyright Motrix Technology Co., Ltd. 2025, 2026 # SPDX-License-Identifier: Apache-2.0 from typing import TypeVar, cast import numpy as np +import torch from motrix_env_core.manager.kernel_data.lowering import KernelDataLowering from motrix_env_core.manager.kernel_data.tree import ( @@ -16,12 +17,15 @@ def canonicalize_kernel_data(value: KernelDataType, *, context: str) -> KernelDataType: - """Copy one KernelData value into environment-owned backing arrays.""" + """Copy one kernel-data value into environment-owned backing arrays.""" if not is_kernel_data(value): - raise TypeError(f"{context} must return a concrete @kernel_data, got {type(value).__name__}.") + raise TypeError(f"{context} must return a concrete kernel-data dataclass, got {type(value).__name__}.") leaves, tree_def = flatten_kernel_data(value) KernelDataLowering().lower(tree_def, context=context) - copied = tuple(leaf.copy() if isinstance(leaf, np.ndarray) else leaf for leaf in leaves) + copied = tuple( + leaf.clone() if isinstance(leaf, torch.Tensor) else leaf.copy() if isinstance(leaf, np.ndarray) else leaf + for leaf in leaves + ) return cast(KernelDataType, unflatten_kernel_data(tree_def, copied)) diff --git a/motrix_env_core/src/motrix_env_core/manager/kernel_data/lowering.py b/motrix_env_core/src/motrix_env_core/manager/kernel_data/lowering.py index f8f90252..9d83e7ea 100644 --- a/motrix_env_core/src/motrix_env_core/manager/kernel_data/lowering.py +++ b/motrix_env_core/src/motrix_env_core/manager/kernel_data/lowering.py @@ -8,6 +8,7 @@ from typing import Any import numpy as np +import torch from motrix_env_core.manager.fingerprint import function_fingerprint from motrix_env_core.manager.kernel_data.methods import dispatch_methods @@ -22,7 +23,7 @@ @dataclass(frozen=True) class KernelLeafLayout: - """Lowered ABI description of one logical KernelData leaf. + """Lowered ABI description of one logical kernel-data leaf. Attributes: tree_def: Logical leaf definition that owns the path and annotation. @@ -83,7 +84,7 @@ class KernelMapLayout: @dataclass(frozen=True) class KernelRecordLayout: - """Lowered schema for one concrete ``@kernel_data`` class tree. + """Lowered schema for one concrete kernel-data class tree. Attributes: tree_def: Logical record definition being lowered. @@ -98,7 +99,7 @@ class KernelRecordLayout: @property def logical_type(self) -> type: - """Return the host ``@kernel_data`` type represented by this layout.""" + """Return the host kernel-data type represented by this layout.""" return self.tree_def.logical_type @@ -174,7 +175,7 @@ def _lower_record( if value_type is None: raise TypeError( f"{context} KernelData {tree_def.logical_type.__name__}." - f"{'.'.join(field.tree_def.path)} expected np.ndarray or numeric scalar." + f"{'.'.join(field.tree_def.path)} expected a tensor, an ndarray, or a numeric scalar." ) child = KernelLeafLayout(field.tree_def, value_type, slot_index) slot_index += 1 @@ -209,7 +210,8 @@ def _lower_map( value_type = self._leaf_type(entry.tree_def.annotation) if value_type is None: raise TypeError( - f"{context} Map {'.'.join(entry.tree_def.path)} expected np.ndarray or numeric scalar." + f"{context} Map {'.'.join(entry.tree_def.path)} expected a tensor, an ndarray, or a numeric " + "scalar." ) child = KernelLeafLayout(entry.tree_def, value_type, slot_index) slot_index += 1 @@ -245,7 +247,7 @@ def _fingerprint_part(layout: KernelDataLayout) -> Any: @staticmethod def _leaf_type(annotation: Any) -> type[Any] | None: - if annotation is np.ndarray: + if annotation in (np.ndarray, torch.Tensor): return annotation if annotation in {bool, int, float}: return annotation diff --git a/motrix_env_core/src/motrix_env_core/manager/kernel_data/tree.py b/motrix_env_core/src/motrix_env_core/manager/kernel_data/tree.py index 00ce8778..3f782194 100644 --- a/motrix_env_core/src/motrix_env_core/manager/kernel_data/tree.py +++ b/motrix_env_core/src/motrix_env_core/manager/kernel_data/tree.py @@ -4,36 +4,35 @@ from __future__ import annotations import hashlib -from dataclasses import dataclass, fields -from typing import Annotated, Any, TypeVar, cast, get_args, get_origin, get_type_hints +from dataclasses import dataclass, fields, is_dataclass +from typing import Annotated, Any, get_args, get_origin, get_type_hints -from typing_extensions import dataclass_transform +import numpy as np +import torch from motrix_env_core.manager.kernel_data.map import Map +from motrix_env_core.sim.read import SimDataQuery -_KernelDataType = TypeVar("_KernelDataType", bound=type) -_MARKER = "__motrix_kernel_data__" _TREE_SCHEMA_VERSION = 2 _TREE_DEF_CACHE: dict[type, TreeClassDef] = {} -@dataclass_transform(frozen_default=True, field_specifiers=()) -def kernel_data(cls: _KernelDataType) -> _KernelDataType: - """Declare one immutable logical class value that can be flattened for Numba kernels.""" - _validate_declaration(cls) - decorated: type[Any] = dataclass(frozen=True, slots=True)(cls) - setattr(decorated, _MARKER, True) - return cast(_KernelDataType, decorated) - - def is_kernel_data_type(value_type: object) -> bool: - """Return whether a type is a concrete ``@kernel_data`` logical class.""" - return isinstance(value_type, type) and value_type.__dict__.get(_MARKER) is True + """Return whether a type is a concrete kernel-data dataclass.""" + return isinstance(value_type, type) and is_dataclass(value_type) def is_kernel_data(value: object) -> bool: - """Return whether a value is an instance of a concrete KernelData type.""" - return is_kernel_data_type(type(value)) + """Return whether a value is an instance of a kernel-data dataclass. + + Simulator read queries are dataclasses too; they resolve through the read + program instead of the kernel-data tree and are excluded here. + """ + return ( + is_dataclass(value) + and not isinstance(value, (np.ndarray, torch.Tensor)) + and not isinstance(value, SimDataQuery) + ) class BaseDef: @@ -58,7 +57,7 @@ class FieldDef: """Named child of a ``TreeClassDef``. Attributes: - name: Python dataclass field name on the logical KernelData class. + name: Python dataclass field name on the logical kernel-data class. tree_def: Logical tree definition for the field value. """ @@ -68,9 +67,9 @@ class FieldDef: @dataclass(frozen=True) class TreeClassDef(BaseDef): - """Root definition of one concrete ``@kernel_data`` class tree. + """Root definition of one concrete kernel-data class tree. - Every flattened KernelData value is represented by a ``TreeClassDef`` root. + Every flattened kernel-data value is represented by a ``TreeClassDef`` root. Its fields recursively contain nested class, map, and leaf definitions. Attributes: @@ -121,9 +120,9 @@ class KernelMapDef(BaseDef): def flatten_kernel_data(value: Any) -> tuple[tuple[Any, ...], TreeClassDef]: - """Return zero-copy leaves and the concrete class tree definition for one KernelData value.""" + """Return zero-copy leaves and the concrete class tree definition for one kernel-data value.""" if not is_kernel_data(value): - raise TypeError(f"KernelData tree root must be a concrete @kernel_data value, got {type(value).__name__}.") + raise TypeError(f"KernelData tree root must be a dataclass value, got {type(value).__name__}.") tree_def, dynamic = _build_class_def(type(value), value, (), ()) if not dynamic: cached = _TREE_DEF_CACHE.get(type(value)) @@ -303,22 +302,3 @@ def _metadata_identity(value: Any) -> Any: def _node_annotation(annotation: Any) -> Any: return get_args(annotation)[0] if get_origin(annotation) is Annotated else annotation - - -def _validate_declaration(logical_type: type) -> None: - kernel_bases = [base for base in logical_type.__bases__ if is_kernel_data_type(base)] - if len(kernel_bases) > 1: - raise TypeError(f"KernelData {logical_type.__name__} must not inherit multiple KernelData bases.") - inherited_fields: set[str] = set() - for base in logical_type.__mro__[1:]: - if is_kernel_data_type(base): - inherited_fields.update(field.name for field in fields(base)) - elif hasattr(base, "__dataclass_fields__") and base is not object: - raise TypeError( - f"KernelData {logical_type.__name__} must not inherit fields from ordinary dataclass {base.__name__}." - ) - own_annotations = logical_type.__dict__.get("__annotations__", {}) - overridden = inherited_fields.intersection((*logical_type.__dict__, *own_annotations)) - if overridden: - names = ", ".join(sorted(overridden)) - raise TypeError(f"KernelData {logical_type.__name__} must not override inherited fields: {names}.") diff --git a/motrix_env_core/src/motrix_env_core/manager/metrics.py b/motrix_env_core/src/motrix_env_core/manager/metrics.py index c4cad8ca..803e491f 100644 --- a/motrix_env_core/src/motrix_env_core/manager/metrics.py +++ b/motrix_env_core/src/motrix_env_core/manager/metrics.py @@ -7,6 +7,7 @@ from typing import Any import numpy as np +import torch _METRIC_METADATA_KEY = "motrix_metric" _UNBOUND_METRIC = object() @@ -43,7 +44,7 @@ def _field_metric(data_field: Any) -> Metric | None: def materialize_metrics(value: Any, num_envs: int, *, context: str) -> Any: - """Replace metric() placeholders with writable ``(num_envs,)`` arrays.""" + """Replace metric() placeholders with writable ``(num_envs,)`` tensor views.""" if not is_dataclass(value): return value @@ -55,8 +56,8 @@ def materialize_metrics(value: Any, num_envs: int, *, context: str) -> Any: if current is _UNBOUND_METRIC: if spec is None: raise TypeError(f"{context} field {data_field.name!r} has an unbound non-metric value.") - dtype = np.dtype(np.float32 if spec.dtype is None else spec.dtype) - current = np.empty(num_envs, dtype=dtype) + dtype = np.float32 if spec.dtype is None else spec.dtype + current = torch.from_numpy(np.empty(num_envs, dtype=dtype)) changed = True values.append(current) return type(value)(*values) if changed else value @@ -69,8 +70,12 @@ def collect_metrics( context: str, context_path: tuple[str, ...], ) -> tuple[tuple[str, np.ndarray], ...]: - """Validate direct metric fields on one canonical manager term.""" + """Validate direct metric fields on one canonical manager term. + Tensor fields register their zero-copy numpy views, so kernel writes + through the tensor and host reads through ``ArrayEnvState.metrics`` share + one backing buffer. + """ del context_path bindings = [] for data_field in fields(value): @@ -78,8 +83,10 @@ def collect_metrics( if spec is None: continue array = getattr(value, data_field.name) + if isinstance(array, torch.Tensor): + array = array.numpy() if not isinstance(array, np.ndarray): - raise TypeError(f"{context} metric field {data_field.name!r} must be an ndarray.") + raise TypeError(f"{context} metric field {data_field.name!r} must hold a tensor or an ndarray.") if array.shape not in {(num_envs,), (num_envs, 1)}: raise ValueError( f"{context} metric field {data_field.name!r} must have shape " diff --git a/motrix_env_core/src/motrix_env_core/manager/rand.py b/motrix_env_core/src/motrix_env_core/manager/rand.py index 68a190e0..a8b67125 100644 --- a/motrix_env_core/src/motrix_env_core/manager/rand.py +++ b/motrix_env_core/src/motrix_env_core/manager/rand.py @@ -3,9 +3,10 @@ """SplitMix64 PRNG primitives for manager environments.""" -import numpy as np +from dataclasses import dataclass -from motrix_env_core.manager.kernel_data import kernel_data +import numpy as np +import torch _GOLDEN_RATIO_64 = np.uint64(0x9E3779B97F4A7C15) _ENV_SALT_64 = np.uint64(0xD2B74407B1CE6E93) @@ -13,15 +14,19 @@ _MIX_MULTIPLIER_2 = np.uint64(0x94D049BB133111EB) -@kernel_data +@dataclass(frozen=True, slots=True) class RandValue: - """Mutable per-environment PRNG state shared by manager terms.""" + """Mutable per-environment PRNG state shared by manager terms. + + ``state`` is a signed-int64 torch view over the per-lane uint64 bit + buffer; the bit pattern is the SplitMix64 state. + """ - state: np.ndarray + state: torch.Tensor def next_uniform(self) -> np.float32: """Advance this environment's PRNG and return a value in ``[-1, 1)``.""" - return next_uniform(self.state) + return next_uniform(self.state.numpy().view(np.uint64)) def initialize_rand_states(num_envs: int, rand_seed: int) -> np.ndarray: diff --git a/motrix_env_core/src/motrix_env_core/manager/resolver.py b/motrix_env_core/src/motrix_env_core/manager/resolver.py index 51dafd57..2aee0ea9 100644 --- a/motrix_env_core/src/motrix_env_core/manager/resolver.py +++ b/motrix_env_core/src/motrix_env_core/manager/resolver.py @@ -313,13 +313,13 @@ def _manager_sim_inputs(sim_data: PhysicsReadProgram) -> tuple[Mapping[str, Mana def _tree_leaf_fingerprint(tree: ResolvedTree) -> tuple[Any, ...]: - """Per-leaf content fingerprint: array shape/dtype identity.""" + """Per-leaf content fingerprint: array/tensor shape/dtype identity.""" return tuple( ( path, None - if not isinstance(leaf.value, np.ndarray) - else (leaf.value.shape, leaf.value.dtype.str, leaf.value.strides), + if not isinstance(leaf.value, (np.ndarray, torch.Tensor)) + else (tuple(leaf.value.shape), str(leaf.value.dtype), getattr(leaf.value, "strides", None)), ) for path, leaf in tree.leaves.items() ) @@ -516,14 +516,14 @@ def _resolve_manager_context( sim = {key: binding.value for key, binding in self._sim_inputs.items()} commands = dict(self._env.command_terms) value = ManagerContext( - env_id=np.arange(self._env.num_envs, dtype=np.int64), + env_id=torch.from_numpy(np.arange(self._env.num_envs, dtype=np.int64)), actions=Map(actions), commands=Map(commands), metrics=Map(self._env.metrics), rand=self._env._rand, sim=Map(sim), dt=np.float32(self._env.cfg.ctrl_dt), - sim_reset_requested=self._env._sim_reset_requested, + sim_reset_requested=torch.from_numpy(self._env._sim_reset_requested), ) _, tree_def = flatten_kernel_data(value) layout = self._kernel_data_lowering.lower( @@ -693,8 +693,8 @@ def _resolve_dispatch_args(self, group: str, name: str, term: Any) -> tuple[Prep def _resolve_term_arg(self, source_name: str, value: Any) -> PreparedArg: """Resolve one manager term argument into a prepared descriptor. - KernelData arguments become structured trees; float-like scalars and - ndarrays become shared leaf slots read by reference; simulator queries + Kernel-data dataclass arguments become structured trees; float-like + scalars and arrays/tensors become shared leaf slots read by reference; simulator queries resolve to their read-program batched array; every other value is a literal constant. """ @@ -730,9 +730,13 @@ def _register_tree( raise TypeError(f"{source_name} value has a different KernelData schema.") leaves: dict[tuple[str, ...], ResolvedLeaf] = {} for prepared_field, field_value in zip(iter_layout_leaves(layout), values, strict=True): - if prepared_field.value_type is np.ndarray and not isinstance(field_value, np.ndarray): - raise TypeError(f"{source_name} field {'.'.join(prepared_field.path)!r} must be np.ndarray.") - if prepared_field.value_type is not np.ndarray and type(field_value) is not prepared_field.value_type: + if prepared_field.value_type in (np.ndarray, torch.Tensor): + if not isinstance(field_value, prepared_field.value_type): + raise TypeError( + f"{source_name} field {'.'.join(prepared_field.path)!r} must be " + f"{prepared_field.value_type.__name__}." + ) + elif type(field_value) is not prepared_field.value_type: raise TypeError( f"{source_name} field {'.'.join(prepared_field.path)!r} must return " f"{prepared_field.value_type.__name__}, got {type(field_value).__name__}." diff --git a/motrix_env_core/src/motrix_env_core/manager/terms.py b/motrix_env_core/src/motrix_env_core/manager/terms.py index e5bd55f8..8fb28b25 100644 --- a/motrix_env_core/src/motrix_env_core/manager/terms.py +++ b/motrix_env_core/src/motrix_env_core/manager/terms.py @@ -24,7 +24,7 @@ def canonicalize_term_args(args: tuple[Any, ...], *, context: str) -> tuple[Any, Accepted leaves are scalars (including strings for compile-time Map-name lookups), scalar tuples, ndarrays, tensors, simulator data queries, and - ``@kernel_data`` carriers (which are normalized in place). + kernel-data dataclass carriers (which are normalized in place). """ values: list[Any] = [] for index, value in enumerate(args): @@ -37,7 +37,7 @@ def canonicalize_term_args(args: tuple[Any, ...], *, context: str) -> tuple[Any, else: raise TypeError( f"{context} args[{index}] must be a scalar, a scalar tuple, an ndarray, a tensor, " - f"a @kernel_data value, or a simulator data query; got {type(value).__name__}." + f"a kernel-data dataclass, or a simulator data query; got {type(value).__name__}." ) return tuple(values) diff --git a/motrix_env_core/src/motrix_env_core/mdp/action.py b/motrix_env_core/src/motrix_env_core/mdp/action.py index 66e5105a..e1e7c35b 100644 --- a/motrix_env_core/src/motrix_env_core/mdp/action.py +++ b/motrix_env_core/src/motrix_env_core/mdp/action.py @@ -3,17 +3,20 @@ """Generic joint-position action terms for manager-based environments.""" +from dataclasses import dataclass + import gymnasium as gym import numpy as np +import torch from motrix_env_core.config import configclass from motrix_env_core.config.scene import RobotCfg -from motrix_env_core.manager import ActionCfg, ActionState, ActionTerm, ManagerEnv, kernel_data +from motrix_env_core.manager import ActionCfg, ActionState, ActionTerm, ManagerEnv from motrix_env_core.mdp.action_space import joint_position_action_space_from_ctrl_ranges from motrix_env_core.sim.model import ActuatorSpec, ActuatorType -@kernel_data +@dataclass(frozen=True, slots=True) class JointPositionActionState(ActionState): """Persistent joint-position action pipeline, history, and shared model data. @@ -41,11 +44,11 @@ class JointPositionActionState(ActionState): path entirely. """ - default_angles: np.ndarray - joint_lower: np.ndarray - joint_upper: np.ndarray - action_scales: np.ndarray - delay_steps: np.ndarray + default_angles: torch.Tensor + joint_lower: torch.Tensor + joint_upper: torch.Tensor + action_scales: torch.Tensor + delay_steps: torch.Tensor delay_lo: int delay_hi: int @@ -64,20 +67,23 @@ def _process(self, actions: np.ndarray) -> np.ndarray: ptr = int(state.action_ptr[0]) target_raw = actions if state.delay_hi > 0: - width = state.action_queue.shape[1] - delay_indices = (ptr - state.delay_steps) % width + queue = state.action_queue.numpy() + width = queue.shape[1] + delay_indices = (ptr - state.delay_steps.numpy()) % width target_raw = np.take_along_axis( - state.action_queue, + queue, delay_indices[:, None, None], axis=1, )[:, 0, :] - return target_raw * state.action_scales + state.default_angles + return target_raw * state.action_scales.numpy() + state.default_angles.numpy() def _reset(self, env_ids: np.ndarray) -> None: """Resample per-episode delay after the base term clears raw history.""" state = self.state if state.delay_hi > 0: - state.delay_steps[env_ids] = np.random.randint(state.delay_lo, state.delay_hi + 1, size=env_ids.size) + state.delay_steps.numpy()[env_ids] = np.random.randint( + state.delay_lo, state.delay_hi + 1, size=env_ids.size + ) @configclass(kw_only=True) @@ -137,13 +143,13 @@ def __call__(self, env: ManagerEnv, actuators: tuple[ActuatorSpec, ...] | None) delay_steps = np.zeros(env.num_envs, dtype=np.int64) action_ptr = np.zeros(1, dtype=np.int64) state = JointPositionActionState( - action_queue=action_queue, - default_angles=default_angles[indices], - joint_lower=joint_lower, - joint_upper=joint_upper, - action_scales=action_scales[indices], - delay_steps=delay_steps, - action_ptr=action_ptr, + action_queue=torch.from_numpy(action_queue), + default_angles=torch.from_numpy(default_angles[indices]), + joint_lower=torch.from_numpy(joint_lower), + joint_upper=torch.from_numpy(joint_upper), + action_scales=torch.from_numpy(action_scales[indices]), + delay_steps=torch.from_numpy(delay_steps), + action_ptr=torch.from_numpy(action_ptr), delay_lo=int(lo), delay_hi=int(hi), ) @@ -155,7 +161,7 @@ def __call__(self, env: ManagerEnv, actuators: tuple[ActuatorSpec, ...] | None) raise ValueError(f"position actuator {spec.name!r} must define or inherit ctrl_range") ctrl_ranges.append(spec.ctrl_range) space = joint_position_action_space_from_ctrl_ranges( - np.asarray(ctrl_ranges, dtype=np.float32), state.default_angles, state.action_scales + np.asarray(ctrl_ranges, dtype=np.float32), state.default_angles.numpy(), state.action_scales.numpy() ) return JointPositionActionTerm(space, state) diff --git a/motrix_env_core/src/motrix_env_core/mdp/state.py b/motrix_env_core/src/motrix_env_core/mdp/state.py index 3cbf892a..92328e74 100644 --- a/motrix_env_core/src/motrix_env_core/mdp/state.py +++ b/motrix_env_core/src/motrix_env_core/mdp/state.py @@ -7,6 +7,9 @@ from typing import TYPE_CHECKING +import numpy as np +import torch + from motrix_env_core.manager.rand import RandValue, initialize_rand_states if TYPE_CHECKING: @@ -15,7 +18,10 @@ def _create_rand_value(env: ManagerEnv) -> RandValue: """Create the framework-owned per-environment PRNG state.""" - return RandValue(initialize_rand_states(env.num_envs, env._rand_seed)) + # View the uint64 bit buffer as signed int64 so Torch kernels can consume + # it directly; the SplitMix64 bit pattern is preserved exactly. + states = initialize_rand_states(env.num_envs, env._rand_seed) + return RandValue(torch.from_numpy(states.view(np.int64))) __all__ = [ diff --git a/motrix_env_core/src/motrix_env_core/mdp/terrain.py b/motrix_env_core/src/motrix_env_core/mdp/terrain.py index b3d80b28..178b1481 100644 --- a/motrix_env_core/src/motrix_env_core/mdp/terrain.py +++ b/motrix_env_core/src/motrix_env_core/mdp/terrain.py @@ -3,13 +3,13 @@ """Static terrain-height grid shared by manager reset and reward terms.""" +from dataclasses import dataclass + import numpy as np import torch -from motrix_env_core.manager import kernel_data - -@kernel_data +@dataclass(frozen=True, slots=True) class HeightFieldGrid: """Static terrain grid for batched bilinear height lookup. @@ -34,10 +34,10 @@ class HeightFieldGrid: ``constant``,忽略所有网格字段。 """ - heights: np.ndarray - origin: np.ndarray - spacing: np.ndarray - z0: np.ndarray + heights: torch.Tensor + origin: torch.Tensor + spacing: torch.Tensor + z0: torch.Tensor constant: np.float32 enabled: bool diff --git a/motrix_env_core/src/motrix_env_core/torch_manager/runtime.py b/motrix_env_core/src/motrix_env_core/torch_manager/runtime.py index 6c6c9b17..a9ddf5ae 100644 --- a/motrix_env_core/src/motrix_env_core/torch_manager/runtime.py +++ b/motrix_env_core/src/motrix_env_core/torch_manager/runtime.py @@ -113,14 +113,29 @@ def tensor(self, anchor: Any, value: Any) -> Any: key = id(value) tensor = tensors.get(key) if tensor is None: - if isinstance(value, np.ndarray): - # 0-dim scalar leaves are configuration constants (body + if isinstance(value, torch.Tensor): + # 0-dim tensor leaves are configuration constants (body # indices, sizes, noise amplitudes, mode flags); as tensors # they force tensor-dependent control flow that Dynamo cannot # trace (``if amplitude != 0`` / ``if sequential_clips``) and # a scalar extraction Inductor cannot compile. Python scalars # carry no dtype-precision concern and specialize branches at - # trace time. + # trace time. Non-0-dim tensors are already zero-copy views + # over their owning arenas and pass through unchanged. + if value.dim() == 0: + if value.dtype in (torch.int64, torch.int32): + tensors[key] = int(value) + return tensors[key] + if value.dtype in (torch.float32, torch.float64): + tensors[key] = float(value) + return tensors[key] + if value.dtype is torch.bool: + tensors[key] = bool(value) + return tensors[key] + tensors[key] = value + return value + if isinstance(value, np.ndarray): + # 0-dim numpy scalar leaves follow the same rule. if value.shape == (): if np.issubdtype(value.dtype, np.integer): tensors[key] = int(value) @@ -224,6 +239,12 @@ def _resolve(self, name: str) -> Any: if self._host is None: raise AttributeError(name) value = getattr(self._host, name) + if isinstance(value, torch.Tensor): + # Host records keep tensor fields; expose the same 0-dim + # integer-constant rule as _TensorCache.tensor. + if value.dim() == 0 and value.dtype in (torch.int64, torch.int32): + return int(value) + return value if isinstance(value, np.ndarray): # Same 0-dim integer-constant rule as _TensorCache.tensor. if value.shape == () and np.issubdtype(value.dtype, np.integer): diff --git a/motrix_env_core/tests/test_action_state.py b/motrix_env_core/tests/test_action_state.py index e5c3d8bd..55556477 100644 --- a/motrix_env_core/tests/test_action_state.py +++ b/motrix_env_core/tests/test_action_state.py @@ -3,16 +3,19 @@ """Host accessors for manager-owned raw-action history.""" +from dataclasses import dataclass + import gymnasium as gym import numpy as np import pytest +import torch from motrix_env_core.manager.actions import ActionState -from motrix_env_core.manager.kernel_data import canonicalize_kernel_data, kernel_data +from motrix_env_core.manager.kernel_data import canonicalize_kernel_data from motrix_env_core.mdp.action import JointPositionActionState, JointPositionActionTerm -@kernel_data +@dataclass(frozen=True, slots=True) class _ActionState(ActionState): scale: float @@ -21,22 +24,25 @@ class _ActionState(ActionState): def test_action_state_accessors_return_history_views(ptr: int) -> None: queue = np.arange(18, dtype=np.float32).reshape(2, 3, 3) state = canonicalize_kernel_data( - _ActionState(action_queue=queue, action_ptr=np.asarray([ptr], dtype=np.int64), scale=1.0), + _ActionState( + action_queue=torch.from_numpy(queue), + action_ptr=torch.from_numpy(np.asarray([ptr], dtype=np.int64)), + scale=1.0, + ), context="action state test", ) - history = state.action_queue.copy() - pointer = state.action_ptr.copy() + history = state.action_queue.numpy().copy() + pointer = state.action_ptr.numpy().copy() current = state.current() previous = state.previous() - np.testing.assert_array_equal(current, history[:, ptr]) - np.testing.assert_array_equal(previous, history[:, (ptr - 1) % 3]) - assert current.shape == previous.shape == (2, 3) - assert np.shares_memory(current, state.action_queue) - assert np.shares_memory(previous, state.action_queue) - np.testing.assert_array_equal(state.action_queue, history) - np.testing.assert_array_equal(state.action_ptr, pointer) + np.testing.assert_array_equal(current.numpy(), history[:, ptr]) + np.testing.assert_array_equal(previous.numpy(), history[:, (ptr - 1) % 3]) + assert tuple(current.shape) == tuple(previous.shape) == (2, 3) + assert current.data_ptr() == state.action_queue[:, ptr, :].data_ptr() + np.testing.assert_array_equal(state.action_queue.numpy(), history) + np.testing.assert_array_equal(state.action_ptr.numpy(), pointer) current[0, 0] = -1.0 assert state.action_queue[0, ptr, 0] == -1.0 @@ -44,19 +50,21 @@ def test_action_state_accessors_return_history_views(ptr: int) -> None: def test_joint_position_action_applies_per_lane_delay() -> None: state = JointPositionActionState( - action_queue=np.asarray( - [ - [[10.0], [20.0], [30.0], [40.0]], - [[1.0], [2.0], [3.0], [4.0]], - ], - dtype=np.float32, + action_queue=torch.from_numpy( + np.asarray( + [ + [[10.0], [20.0], [30.0], [40.0]], + [[1.0], [2.0], [3.0], [4.0]], + ], + dtype=np.float32, + ) ), - action_ptr=np.asarray([2], dtype=np.int64), - default_angles=np.asarray([100.0], dtype=np.float32), - joint_lower=np.asarray([-100.0], dtype=np.float32), - joint_upper=np.asarray([100.0], dtype=np.float32), - action_scales=np.asarray([2.0], dtype=np.float32), - delay_steps=np.asarray([1, 2], dtype=np.int64), + action_ptr=torch.from_numpy(np.asarray([2], dtype=np.int64)), + default_angles=torch.from_numpy(np.asarray([100.0], dtype=np.float32)), + joint_lower=torch.from_numpy(np.asarray([-100.0], dtype=np.float32)), + joint_upper=torch.from_numpy(np.asarray([100.0], dtype=np.float32)), + action_scales=torch.from_numpy(np.asarray([2.0], dtype=np.float32)), + delay_steps=torch.from_numpy(np.asarray([1, 2], dtype=np.int64)), delay_lo=0, delay_hi=2, ) @@ -67,13 +75,13 @@ def test_joint_position_action_applies_per_lane_delay() -> None: def test_joint_position_action_resamples_delay_on_reset(monkeypatch) -> None: state = JointPositionActionState( - action_queue=np.zeros((3, 2, 1), dtype=np.float32), - action_ptr=np.asarray([0], dtype=np.int64), - default_angles=np.zeros(1, dtype=np.float32), - joint_lower=np.zeros(1, dtype=np.float32), - joint_upper=np.ones(1, dtype=np.float32), - action_scales=np.ones(1, dtype=np.float32), - delay_steps=np.zeros(3, dtype=np.int64), + action_queue=torch.zeros((3, 2, 1), dtype=torch.float32), + action_ptr=torch.zeros(1, dtype=torch.int64), + default_angles=torch.zeros(1, dtype=torch.float32), + joint_lower=torch.zeros(1, dtype=torch.float32), + joint_upper=torch.ones(1, dtype=torch.float32), + action_scales=torch.ones(1, dtype=torch.float32), + delay_steps=torch.zeros(3, dtype=torch.int64), delay_lo=1, delay_hi=3, ) @@ -82,17 +90,17 @@ def test_joint_position_action_resamples_delay_on_reset(monkeypatch) -> None: term._reset(np.asarray([0, 2], dtype=np.int64)) - np.testing.assert_array_equal(term.state.delay_steps, [2, 0, 2]) + np.testing.assert_array_equal(term.state.delay_steps.numpy(), [2, 0, 2]) def test_action_state_accessors_follow_cursor_changes() -> None: state = ActionState( - action_queue=np.arange(12, dtype=np.float32).reshape(2, 2, 3), - action_ptr=np.zeros(1, dtype=np.int64), + action_queue=torch.from_numpy(np.arange(12, dtype=np.float32).reshape(2, 2, 3)), + action_ptr=torch.zeros(1, dtype=torch.int64), ) - initial_current = state.current().copy() - initial_previous = state.previous().copy() + initial_current = state.current().clone() + initial_previous = state.previous().clone() state.action_ptr[0] = 1 - np.testing.assert_array_equal(state.current(), initial_previous) - np.testing.assert_array_equal(state.previous(), initial_current) + np.testing.assert_array_equal(state.current().numpy(), initial_previous.numpy()) + np.testing.assert_array_equal(state.previous().numpy(), initial_current.numpy()) diff --git a/motrix_env_core/tests/test_kernel_data.py b/motrix_env_core/tests/test_kernel_data.py index a417aeb3..452a4435 100644 --- a/motrix_env_core/tests/test_kernel_data.py +++ b/motrix_env_core/tests/test_kernel_data.py @@ -1,49 +1,50 @@ # Copyright Motphys Technology Co., Ltd. 2025, 2026 # SPDX-License-Identifier: Apache-2.0 -from dataclasses import FrozenInstanceError +from dataclasses import FrozenInstanceError, dataclass import numpy as np import pytest +import torch from motrix_env_core.manager.kernel_data import ( KernelDataLowering, Map, flatten_kernel_data, + is_kernel_data, iter_layout_leaves, - kernel_data, unflatten_kernel_data, ) -@kernel_data +@dataclass(frozen=True, slots=True) class _Pose: position: np.ndarray orientation: np.ndarray -@kernel_data +@dataclass(frozen=True, slots=True) class _RobotBase: velocity: np.ndarray -@kernel_data +@dataclass(frozen=True, slots=True) class _Robot(_RobotBase): pose: _Pose -@kernel_data +@dataclass(frozen=True, slots=True) class _Tracking: robot: _Robot phase: np.float32 -@kernel_data +@dataclass(frozen=True, slots=True) class _MappedValues: values: Map -@kernel_data +@dataclass(frozen=True, slots=True) class _ArrayMapValues: values: Map[np.ndarray] @@ -81,6 +82,19 @@ def test_kernel_data_flatten_is_zero_copy_and_rebuilds_host_value() -> None: assert rebuilt.phase == np.float32(0.25) +def test_kernel_data_accepts_tensor_leaves() -> None: + @dataclass(frozen=True, slots=True) + class _Tensors: + buffer: torch.Tensor + + buffer = torch.zeros((2, 3), dtype=torch.float32) + value = _Tensors(buffer) + leaves, tree_def = flatten_kernel_data(value) + assert leaves == (buffer,) + layout = KernelDataLowering().lower(tree_def, context="test") + assert iter_layout_leaves(layout)[0].value_type is torch.Tensor + + def test_kernel_data_schema_fingerprint_ignores_values() -> None: first_value = _Tracking( _Robot(np.asarray([[1.0]], dtype=np.float32), _Pose(np.asarray([[2.0]]), np.asarray([3.0]))), @@ -113,28 +127,16 @@ def test_map_schema_changes_tree_and_layout_identity() -> None: def test_kernel_data_rejects_unsupported_leaf() -> None: - @kernel_data + @dataclass(frozen=True, slots=True) class _Invalid: label: str _, tree_def = flatten_kernel_data(_Invalid("invalid")) - with pytest.raises(TypeError, match="expected np.ndarray"): + with pytest.raises(TypeError, match="expected a tensor"): KernelDataLowering().lower(tree_def, context="invalid") -def test_kernel_data_rejects_inherited_field_override_before_slots_transform() -> None: - @kernel_data - class _Base: - value: int = 1 - - with pytest.raises(TypeError, match="must not override inherited fields: value"): - - @kernel_data - class _AnnotationOverride(_Base): - value: int = 2 - - with pytest.raises(TypeError, match="must not override inherited fields: value"): - - @kernel_data - class _DefaultOverride(_Base): - value = 2 +def test_ndarrays_and_tensors_are_not_kernel_data_values() -> None: + assert not is_kernel_data(np.zeros(2)) + assert not is_kernel_data(torch.zeros(2)) + assert is_kernel_data(_Pose(np.zeros(2), np.zeros(2))) diff --git a/motrix_env_core/tests/test_manager_env.py b/motrix_env_core/tests/test_manager_env.py index 61445065..c089a11c 100644 --- a/motrix_env_core/tests/test_manager_env.py +++ b/motrix_env_core/tests/test_manager_env.py @@ -9,7 +9,7 @@ without a simulator backend. """ -from dataclasses import fields, replace +from dataclasses import dataclass, fields, replace import gymnasium as gym import numpy as np @@ -40,23 +40,22 @@ RewardTermCfg, TerminationTerm, TerminationTermCfg, - kernel_data, metric, ) from motrix_env_core.manager.dispatch import dispatch # noqa: E402 from motrix_env_core.manager.observations import create_observation_groups -@kernel_data +@dataclass(frozen=True, slots=True) class _TestActionState(ActionState): - reset_flags: np.ndarray - source: np.ndarray + reset_flags: torch.Tensor + source: torch.Tensor -@kernel_data +@dataclass(frozen=True, slots=True) class _SecondActionState(ActionState): - reset_flags: np.ndarray - source: np.ndarray + reset_flags: torch.Tensor + source: torch.Tensor class _TestActionTerm(ActionTerm): @@ -64,10 +63,10 @@ def __init__(self, action_space: gym.spaces.Box, state: _TestActionState) -> Non super().__init__(action_space, state) def _process(self, actions: np.ndarray) -> None: - self.state.source[...] = actions + self.state.source.copy_(torch.from_numpy(actions)) def _reset(self, env_ids: np.ndarray) -> None: - self.state.reset_flags.fill(False) + self.state.reset_flags.fill_(False) self.state.reset_flags[env_ids] = True @@ -76,10 +75,10 @@ def __init__(self, action_space: gym.spaces.Box, state: _SecondActionState) -> N super().__init__(action_space, state) def _process(self, actions: np.ndarray) -> None: - self.state.source[...] = actions + self.state.source.copy_(torch.from_numpy(actions)) def _reset(self, env_ids: np.ndarray) -> None: - self.state.reset_flags.fill(False) + self.state.reset_flags.fill_(False) self.state.reset_flags[env_ids] = True @@ -94,9 +93,9 @@ def _second_action_space(_: ManagerEnv, __: np.ndarray | None) -> gym.spaces.Box ) -@kernel_data +@dataclass(frozen=True, slots=True) class _ObservationParams: - offset: np.ndarray + offset: torch.Tensor scale: np.float32 @@ -115,7 +114,7 @@ def __call__(self, env: ManagerEnv) -> ObsTerm: return ObsTerm( 2, _injected_observation, - _ObservationParams(np.asarray([1.0], dtype=np.float32), np.float32(self.scale)), + _ObservationParams(torch.from_numpy(np.asarray([1.0], dtype=np.float32)), np.float32(self.scale)), ) @@ -197,10 +196,10 @@ def __call__(self, env: ManagerEnv, actuator_indices: np.ndarray | None) -> Acti return _TestActionTerm( _test_action_space(env, actuator_indices), _TestActionState( - np.zeros((env.num_envs, 2, 1), dtype=np.float32), - np.zeros(1, dtype=np.int64), - np.zeros((env.num_envs, 1), dtype=bool), - np.zeros((env.num_envs, 1), dtype=np.float32), + torch.from_numpy(np.zeros((env.num_envs, 2, 1), dtype=np.float32)), + torch.from_numpy(np.zeros(1, dtype=np.int64)), + torch.from_numpy(np.zeros((env.num_envs, 1), dtype=bool)), + torch.from_numpy(np.zeros((env.num_envs, 1), dtype=np.float32)), ), ) @@ -213,10 +212,10 @@ def __call__(self, env: ManagerEnv, actuator_indices: np.ndarray | None) -> Acti return _SecondActionTerm( _second_action_space(env, actuator_indices), _SecondActionState( - np.zeros((env.num_envs, 2, 2), dtype=np.float32), - np.zeros(1, dtype=np.int64), - np.zeros((env.num_envs, 1), dtype=bool), - np.zeros((env.num_envs, 2), dtype=np.float32), + torch.from_numpy(np.zeros((env.num_envs, 2, 2), dtype=np.float32)), + torch.from_numpy(np.zeros(1, dtype=np.int64)), + torch.from_numpy(np.zeros((env.num_envs, 1), dtype=bool)), + torch.from_numpy(np.zeros((env.num_envs, 2), dtype=np.float32)), ), ) @@ -226,14 +225,14 @@ class _CounterCommandCfg(CommandCfg): def __call__(self, env: ManagerEnv) -> "_CounterCommand": shape = (env.num_envs, 1) return _CounterCommand( - double=np.zeros(shape, dtype=np.float32), - command=np.zeros(shape, dtype=np.float32), + double=torch.zeros(shape, dtype=torch.float32), + command=torch.zeros(shape, dtype=torch.float32), ) -@kernel_data +@dataclass(frozen=True, slots=True) class _CounterCommand(CommandTerm): - double: np.ndarray = metric() + double: torch.Tensor = metric() @staticmethod @dispatch @@ -462,10 +461,11 @@ def test_rand_value_owns_seeded_per_environment_states() -> None: reproduced = _ManagerEnv(num_envs=3, rand_seed=7)._rand.state different = _ManagerEnv(num_envs=3, rand_seed=8)._rand.state - assert first.shape == (3, 1) - assert first.dtype == np.uint64 - np.testing.assert_array_equal(first, reproduced) - assert not np.array_equal(first, different) + # The state is a signed-int64 tensor view over the per-lane uint64 bits. + assert tuple(first.shape) == (3, 1) + assert first.dtype == torch.int64 + np.testing.assert_array_equal(first.numpy().view(np.uint64), reproduced.numpy().view(np.uint64)) + assert not np.array_equal(first.numpy(), different.numpy()) def test_step_perf_records_standard_and_manager_phases() -> None: @@ -536,7 +536,7 @@ def test_manager_context_is_injected_once_and_reused_across_all_term_kinds() -> slice(3, 4), ] - action.state.source[:, 0] = [0.25, 0.75] + action.state.source[:, 0] = torch.tensor([0.25, 0.75]) env.compute_transition(state) env.compute_observation(state) @@ -582,9 +582,9 @@ def test_command_evaluation_updates_state_only_in_evaluate_kernel() -> None: assert env._compiled_manager_program is not None action = env.action_terms["test"] assert isinstance(action, ActionTerm) - action.state.source[:, 0] = [1.0, 2.0] + action.state.source[:, 0] = torch.tensor([1.0, 2.0]) derived = _counter_command(env).double - derived.fill(-1.0) + derived.fill_(-1.0) env._refresh_sim_reads() inputs = env._resolved_inputs @@ -604,7 +604,7 @@ def test_command_term_updates_lifecycle_and_selected_reset_ids() -> None: command = _counter_command(env).command np.testing.assert_array_equal(command, -1.0) - command[:] = [[1.0], [2.0], [3.0]] + command[:] = torch.tensor([[1.0], [2.0], [3.0]]) action = env.action_terms["test"] state.terminated[:] = [False, True, False] env._reset_done_envs() @@ -674,19 +674,24 @@ def test_final_observations_capture_command_before_reset_and_survive_steps(with_ def test_final_observation_noise_uses_shared_observation_rng() -> None: env = _ManagerEnv() env.init_state() - rng_before = env._rand.state.copy() + rng_before = env._rand.state.clone() observe = env._task_program.observe_kernel def noisy_observe(inputs, outputs): observe(inputs, outputs) env._rand.state[:] += np.uint64(1) - outputs.policy_obs[:, 0] = env._rand.state[:, 0].astype(np.float32) + outputs.policy_obs[:, 0] = env._rand.state[:, 0].to(torch.float32) env._task_program = replace(env._task_program, observe_kernel=noisy_observe) final = env._capture_final_obs(np.array([0], dtype=np.int64)) - np.testing.assert_array_equal(env._rand.state, rng_before + np.uint64(1)) + np.testing.assert_array_equal( + env._rand.state.numpy().view(np.uint64), rng_before.numpy().view(np.uint64) + np.uint64(1) + ) saved_final = final.obs.policy.copy() - np.testing.assert_array_equal(final.obs.policy[:, 0], (rng_before[[0], 0] + np.uint64(1)).astype(np.float32)) + np.testing.assert_array_equal( + final.obs.policy[:, 0], + (rng_before[[0], 0] + 1).to(torch.float32).numpy(), + ) env._execute_observe_kernel(env._resolved_inputs) np.testing.assert_array_equal(env._rand.state, rng_before + np.uint64(2)) np.testing.assert_array_equal(final.obs.policy, saved_final) @@ -698,7 +703,7 @@ def test_post_reset_observation_does_not_run_command_evaluation() -> None: assert env._compiled_manager_program is not None action = env.action_terms["test"] assert isinstance(action, ActionTerm) - action.state.source[:, 0] = [0.25, 0.5, 0.75] + action.state.source[:, 0] = torch.tensor([0.25, 0.5, 0.75]) state.terminated[:] = [False, True, False] env._refresh_sim_reads() env._reset_done_envs() @@ -729,10 +734,10 @@ def test_shared_arrays_in_terms_are_writable() -> None: assert isinstance(injected, ObsTerm) assert injected.args - assert injected.args[0].offset.flags.writeable + assert torch.is_tensor(injected.args[0].offset) -def test_observation_term_must_be_kernel_data_with_instance_compute() -> None: +def test_observation_term_must_be_a_term_with_instance_compute() -> None: class _InvalidObservation: def compute(self, env_id: int, out: torch.Tensor) -> None: out[0] = env_id diff --git a/motrix_env_core/tests/test_manager_sim_backend.py b/motrix_env_core/tests/test_manager_sim_backend.py index 1980b5e7..a3e79133 100644 --- a/motrix_env_core/tests/test_manager_sim_backend.py +++ b/motrix_env_core/tests/test_manager_sim_backend.py @@ -6,6 +6,7 @@ import subprocess import sys from collections.abc import Mapping +from dataclasses import dataclass import gymnasium as gym import numpy as np @@ -16,7 +17,6 @@ from motrix_env_core.manager.actions import ActionCfg, ActionState, ActionTerm, ManagerActionsCfg from motrix_env_core.manager.context import ManagerContext from motrix_env_core.manager.env import ManagerBasedEnvCfg, ManagerEnv -from motrix_env_core.manager.kernel_data import kernel_data from motrix_env_core.manager.observations import ( ManagerObservationGroupCfg, ManagerObservationsCfg, @@ -199,9 +199,9 @@ def write_compiler(self): return self._write_compiler -@kernel_data +@dataclass(frozen=True, slots=True) class _CtrlActionState(ActionState): - source: np.ndarray + source: torch.Tensor class _CtrlActionTerm(ActionTerm): @@ -209,7 +209,7 @@ def __init__(self, action_space: gym.spaces.Box, state: _CtrlActionState) -> Non super().__init__(action_space, state) def _process(self, actions: np.ndarray) -> np.ndarray: - self.state.source[...] = actions + self.state.source.copy_(torch.from_numpy(actions)) return actions * np.float32(2.0) def _reset(self, env_ids: np.ndarray) -> None: @@ -229,9 +229,9 @@ def __call__(self, env, actuator_indices): return _CtrlActionTerm( _ctrl_action_space(env, actuator_indices), _CtrlActionState( - np.zeros((env.num_envs, 2, 1), dtype=np.float32), - np.zeros(1, dtype=np.int64), - np.zeros((env.num_envs, 1), dtype=np.float32), + torch.from_numpy(np.zeros((env.num_envs, 2, 1), dtype=np.float32)), + torch.from_numpy(np.zeros(1, dtype=np.int64)), + torch.from_numpy(np.zeros((env.num_envs, 1), dtype=np.float32)), ), ) diff --git a/motrix_env_core/tests/test_torch_manager_actions.py b/motrix_env_core/tests/test_torch_manager_actions.py index e63d19e6..f38a47e3 100644 --- a/motrix_env_core/tests/test_torch_manager_actions.py +++ b/motrix_env_core/tests/test_torch_manager_actions.py @@ -19,13 +19,15 @@ def _make_term(width: int, delayed: bool) -> JointPositionActionTerm: # Nonzero history catches off-by-one reads before fresh actions fill the ring. queue = np.arange(num_envs * width * num_actions, dtype=np.float32).reshape(num_envs, width, num_actions) / 8 state = JointPositionActionState( - action_queue=queue, - action_ptr=np.asarray([width - 1], dtype=np.int64), - default_angles=np.asarray([0.25, -2.0, 0.75], dtype=np.float32), - joint_lower=np.full(num_actions, -np.inf, dtype=np.float32), - joint_upper=np.full(num_actions, np.inf, dtype=np.float32), - action_scales=np.asarray([0.5, 2.0, 0.0], dtype=np.float32), - delay_steps=np.arange(num_envs, dtype=np.int64) % width if delayed else np.zeros(num_envs, dtype=np.int64), + action_queue=torch.from_numpy(queue), + action_ptr=torch.from_numpy(np.asarray([width - 1], dtype=np.int64)), + default_angles=torch.from_numpy(np.asarray([0.25, -2.0, 0.75], dtype=np.float32)), + joint_lower=torch.from_numpy(np.full(num_actions, -np.inf, dtype=np.float32)), + joint_upper=torch.from_numpy(np.full(num_actions, np.inf, dtype=np.float32)), + action_scales=torch.from_numpy(np.asarray([0.5, 2.0, 0.0], dtype=np.float32)), + delay_steps=torch.from_numpy( + np.arange(num_envs, dtype=np.int64) % width if delayed else np.zeros(num_envs, dtype=np.int64) + ), # Keep native reset resampling deterministic without mocking its path. delay_lo=width - 1 if delayed else 0, delay_hi=width - 1 if delayed else 0, @@ -36,23 +38,23 @@ def _make_term(width: int, delayed: bool) -> JointPositionActionTerm: def _exercise_process(process: Callable, *, width: int, delayed: bool) -> None: term = _make_term(width, delayed) state = term.state - queue = torch.from_numpy(state.action_queue.copy()) - ptr = torch.from_numpy(state.action_ptr.copy()) - delays = torch.from_numpy(state.delay_steps.copy()) - scales = torch.from_numpy(state.action_scales.copy()) - defaults = torch.from_numpy(state.default_angles.copy()) + queue = state.action_queue.clone() + ptr = state.action_ptr.clone() + delays = state.delay_steps.clone() + scales = state.action_scales.clone() + defaults = state.default_angles.clone() rng = np.random.default_rng(17) for _ in range(2 * width + 1): - raw = rng.uniform(-5, 5, size=state.current().shape).astype(np.float32) + raw = rng.uniform(-5, 5, size=tuple(state.current().shape)).astype(np.float32) actions = torch.from_numpy(raw.copy()) queue_before, ptr_before = queue.clone(), ptr.clone() targets, new_queue, new_ptr, current, previous = process(queue, ptr, actions, delays, scales, defaults) expected = term.process(raw) torch.testing.assert_close(targets, torch.from_numpy(expected)) - torch.testing.assert_close(new_queue, torch.from_numpy(state.action_queue), rtol=0, atol=0) - torch.testing.assert_close(new_ptr, torch.from_numpy(state.action_ptr), rtol=0, atol=0) - torch.testing.assert_close(current, torch.from_numpy(state.current()), rtol=0, atol=0) - torch.testing.assert_close(previous, torch.from_numpy(state.previous()), rtol=0, atol=0) + torch.testing.assert_close(new_queue, state.action_queue, rtol=0, atol=0) + torch.testing.assert_close(new_ptr, state.action_ptr, rtol=0, atol=0) + torch.testing.assert_close(current, state.current(), rtol=0, atol=0) + torch.testing.assert_close(previous, state.previous(), rtol=0, atol=0) torch.testing.assert_close(queue, queue_before, rtol=0, atol=0) torch.testing.assert_close(ptr, ptr_before, rtol=0, atol=0) torch.testing.assert_close(actions, torch.from_numpy(raw), rtol=0, atol=0) @@ -69,8 +71,8 @@ def test_joint_position_process_matches_native(width: int, delayed: bool) -> Non def _exercise_reset(reset: Callable, process: Callable) -> None: term = _make_term(4, True) state = term.state - queue = torch.from_numpy(state.action_queue.copy()) - ptr = torch.from_numpy(state.action_ptr.copy()) + queue = state.action_queue.clone() + ptr = state.action_ptr.clone() for mask in ( np.asarray([False, True, False, True, False]), np.zeros(5, dtype=bool), @@ -80,25 +82,25 @@ def _exercise_reset(reset: Callable, process: Callable) -> None: pointer_before = ptr.clone() queue_after = reset(queue, torch.from_numpy(mask)) term.reset(np.flatnonzero(mask)) - torch.testing.assert_close(queue_after, torch.from_numpy(state.action_queue), rtol=0, atol=0) + torch.testing.assert_close(queue_after, state.action_queue, rtol=0, atol=0) torch.testing.assert_close(queue, before, rtol=0, atol=0) torch.testing.assert_close(ptr, pointer_before, rtol=0, atol=0) - torch.testing.assert_close(ptr, torch.from_numpy(state.action_ptr), rtol=0, atol=0) + torch.testing.assert_close(ptr, state.action_ptr, rtol=0, atol=0) # Reset randomization belongs to the host, not the pure history helper. - delays = torch.from_numpy(state.delay_steps.copy()) + delays = state.delay_steps.clone() raw = np.full((5, 3), 3.5, dtype=np.float32) targets, queue, ptr, current, previous = process( queue_after, ptr, torch.from_numpy(raw), delays, - torch.from_numpy(state.action_scales), - torch.from_numpy(state.default_angles), + state.action_scales, + state.default_angles, ) torch.testing.assert_close(targets, torch.from_numpy(term.process(raw))) - torch.testing.assert_close(queue, torch.from_numpy(state.action_queue), rtol=0, atol=0) - torch.testing.assert_close(current, torch.from_numpy(state.current()), rtol=0, atol=0) - torch.testing.assert_close(previous, torch.from_numpy(state.previous()), rtol=0, atol=0) + torch.testing.assert_close(queue, state.action_queue, rtol=0, atol=0) + torch.testing.assert_close(current, state.current(), rtol=0, atol=0) + torch.testing.assert_close(previous, state.previous(), rtol=0, atol=0) def test_reset_action_history_matches_native_selected_lanes() -> None: diff --git a/motrix_envs/src/motrix_envs/locomotion/ball_balance/mdp/reset.py b/motrix_envs/src/motrix_envs/locomotion/ball_balance/mdp/reset.py index d83901b8..0bfbfbd6 100644 --- a/motrix_envs/src/motrix_envs/locomotion/ball_balance/mdp/reset.py +++ b/motrix_envs/src/motrix_envs/locomotion/ball_balance/mdp/reset.py @@ -5,6 +5,8 @@ from __future__ import annotations +from dataclasses import dataclass + import numpy as np import torch @@ -14,7 +16,7 @@ ResetTerm, ResetTermCfg, ) -from motrix_env_core.manager.kernel_data import Map, kernel_data +from motrix_env_core.manager.kernel_data import Map from motrix_env_core.math.quaternion_array import from_euler from motrix_env_core.sim import ( BodyAngularVelocityWrite, @@ -125,11 +127,11 @@ def __call__(self, ctx) -> ResetTerm: ) -@kernel_data +@dataclass(frozen=True, slots=True) class DofResetParams: - default_pos: np.ndarray - joint_lower: np.ndarray - joint_upper: np.ndarray + default_pos: torch.Tensor + joint_lower: torch.Tensor + joint_upper: torch.Tensor noise_scale: np.float32 @@ -168,9 +170,9 @@ def __call__(self, ctx) -> ResetTerm: return ResetTerm( _reset_body_dof_pos, DofResetParams( - default_pos=default_pos, - joint_lower=joint_lower, - joint_upper=joint_upper, + default_pos=torch.from_numpy(default_pos), + joint_lower=torch.from_numpy(joint_lower), + joint_upper=torch.from_numpy(joint_upper), noise_scale=np.float32(self.noise * self.noise_scale), ), writes={"position": JointPositionWrite(joints), "velocity": JointVelocityWrite(joints)}, diff --git a/motrix_envs/src/motrix_envs/locomotion/ball_balance/mdp/rewards.py b/motrix_envs/src/motrix_envs/locomotion/ball_balance/mdp/rewards.py index 3ea69cc2..e58dddc6 100644 --- a/motrix_envs/src/motrix_envs/locomotion/ball_balance/mdp/rewards.py +++ b/motrix_envs/src/motrix_envs/locomotion/ball_balance/mdp/rewards.py @@ -3,12 +3,13 @@ """Reward terms for the manager-based ball-balance task.""" +from dataclasses import dataclass + import numpy as np import torch from motrix_env_core.config import configclass from motrix_env_core.manager import ManagerContext, RewardTerm, RewardTermCfg -from motrix_env_core.manager.kernel_data import kernel_data from motrix_env_core.math.quaternion_array import rotate_inverse as rotate_inverse_batch @@ -70,9 +71,9 @@ def __call__(self, ctx) -> RewardTerm: return RewardTerm(ball_under_feet_reward, np.float32(self.sigma)) -@kernel_data +@dataclass(frozen=True, slots=True) class DofDefaultParams: - reference: np.ndarray + reference: torch.Tensor sigma: np.float32 @@ -102,5 +103,5 @@ def __call__(self, ctx) -> RewardTerm: if missing: raise KeyError(f"key pose 'default' is missing joints: {missing}") reference = np.asarray([default[name] for name in joint_names], dtype=np.float32) - params = DofDefaultParams(reference, np.float32(self.sigma)) + params = DofDefaultParams(torch.from_numpy(reference), np.float32(self.sigma)) return RewardTerm(dof_default_reward, params) diff --git a/motrix_envs/src/motrix_envs/locomotion/humanoid/walk_manager_mdp/command.py b/motrix_envs/src/motrix_envs/locomotion/humanoid/walk_manager_mdp/command.py index 45e88412..7da1fae6 100644 --- a/motrix_envs/src/motrix_envs/locomotion/humanoid/walk_manager_mdp/command.py +++ b/motrix_envs/src/motrix_envs/locomotion/humanoid/walk_manager_mdp/command.py @@ -4,6 +4,7 @@ """Velocity-command and gait-phase command term.""" import math +from dataclasses import dataclass import numpy as np import torch @@ -14,7 +15,6 @@ CommandTerm, ManagerContext, ManagerEnv, - kernel_data, metric, ) from motrix_env_core.manager.commands import ResetContext @@ -84,7 +84,7 @@ def _resample_command( commands.copy_(torch.where(resample[:, None], sampled, commands)) -@kernel_data +@dataclass(frozen=True, slots=True) class WalkCommand(CommandTerm): """Per-environment velocity command and gait-phase clock. @@ -125,15 +125,15 @@ class WalkCommand(CommandTerm): the ``command_steps`` metric). """ - vel_limit_low: np.ndarray - vel_limit_high: np.ndarray + vel_limit_low: torch.Tensor + vel_limit_high: torch.Tensor stand_prob: np.float32 resample_steps: np.float32 gait_period_base: np.float32 gait_period_width: np.float32 curriculum_enabled: bool - penalty_scale: np.ndarray - avg_ep_len: np.ndarray + penalty_scale: torch.Tensor + avg_ep_len: torch.Tensor level_down_threshold: np.float32 level_up_threshold: np.float32 degree: np.float32 @@ -143,12 +143,12 @@ class WalkCommand(CommandTerm): # Per-environment state. ``phase`` and ``steps`` are published as metrics; # the kernel lowering hands each lane a writable row view. # Contract goal vector (inherited CommandTerm.command row view). - # command: np.ndarray - phase_offset: np.ndarray - sin_cos: np.ndarray - phase: np.ndarray - phase_step: np.ndarray - steps: np.ndarray = metric(name="command_steps", dtype=np.float32) + # command: torch.Tensor + phase_offset: torch.Tensor + sin_cos: torch.Tensor + phase: torch.Tensor + phase_step: torch.Tensor + steps: torch.Tensor = metric(name="command_steps", dtype=np.float32) @staticmethod @dispatch @@ -216,18 +216,18 @@ def reset(self, ctx: ResetContext) -> None: truncated), so the EMA sample matches the direct env's ``done = terminated | truncated``. """ - ctx.metrics["penalty_scale"] = float(self.penalty_scale[0]) + penalty_scale = self.penalty_scale.numpy() + avg_ep_len = self.avg_ep_len.numpy() + ctx.metrics["penalty_scale"] = float(penalty_scale[0]) if not self.curriculum_enabled or ctx.env_ids.size == 0: return - ep_len = self.steps[ctx.env_ids, 0].astype(np.float64) + 1.0 - self.avg_ep_len[0] = np.float32(0.99 * self.avg_ep_len[0] + 0.01 * float(ep_len.mean())) - if self.avg_ep_len[0] < self.level_down_threshold: - self.penalty_scale[0] *= np.float32(1.0 - self.degree) - elif self.avg_ep_len[0] > self.level_up_threshold: - self.penalty_scale[0] *= np.float32(1.0 + self.degree) - self.penalty_scale[0] = np.float32( - min(max(float(self.penalty_scale[0]), float(self.min_scale)), float(self.max_scale)) - ) + ep_len = self.steps.numpy()[ctx.env_ids, 0].astype(np.float64) + 1.0 + avg_ep_len[0] = np.float32(0.99 * avg_ep_len[0] + 0.01 * float(ep_len.mean())) + if avg_ep_len[0] < self.level_down_threshold: + penalty_scale[0] *= np.float32(1.0 - self.degree) + elif avg_ep_len[0] > self.level_up_threshold: + penalty_scale[0] *= np.float32(1.0 + self.degree) + penalty_scale[0] = np.float32(min(max(float(penalty_scale[0]), float(self.min_scale)), float(self.max_scale))) @configclass(kw_only=True) @@ -255,32 +255,36 @@ class WalkCommandCfg(CommandCfg): def __call__(self, env: ManagerEnv) -> WalkCommand: num_envs = env.num_envs return WalkCommand( - vel_limit_low=np.asarray(self.vel_limit[0], dtype=np.float32), - vel_limit_high=np.asarray(self.vel_limit[1], dtype=np.float32), + vel_limit_low=torch.from_numpy(np.asarray(self.vel_limit[0], dtype=np.float32)), + vel_limit_high=torch.from_numpy(np.asarray(self.vel_limit[1], dtype=np.float32)), stand_prob=np.float32(self.stand_prob), resample_steps=np.float32(max(int(round(self.resampling_time / self.ctrl_dt)), 1)), - phase_step=np.full( - (num_envs, 1), - np.float32(2.0 * math.pi * self.ctrl_dt / self.gait_period), - dtype=np.float32, + phase_step=torch.from_numpy( + np.full( + (num_envs, 1), + np.float32(2.0 * math.pi * self.ctrl_dt / self.gait_period), + dtype=np.float32, + ) ), gait_period_base=np.float32(self.gait_period), gait_period_width=np.float32(self.gait_period_randomization_width), curriculum_enabled=self.curriculum_enabled, - penalty_scale=np.full( - (1,), - self.initial_scale if self.curriculum_enabled else 1.0, - dtype=np.float32, + penalty_scale=torch.from_numpy( + np.full( + (1,), + self.initial_scale if self.curriculum_enabled else 1.0, + dtype=np.float32, + ) ), - avg_ep_len=np.zeros((1,), dtype=np.float32), + avg_ep_len=torch.from_numpy(np.zeros((1,), dtype=np.float32)), level_down_threshold=np.float32(self.level_down_threshold), level_up_threshold=np.float32(self.level_up_threshold), degree=np.float32(self.degree), min_scale=np.float32(self.min_scale), max_scale=np.float32(self.max_scale), - command=np.zeros((num_envs, 3), dtype=np.float32), - phase_offset=np.zeros((num_envs, 2), dtype=np.float32), - sin_cos=np.zeros((num_envs, 4), dtype=np.float32), - phase=np.zeros((num_envs, 2), dtype=np.float32), - steps=np.zeros((num_envs, 1), dtype=np.float32), + command=torch.zeros((num_envs, 3), dtype=torch.float32), + phase_offset=torch.zeros((num_envs, 2), dtype=torch.float32), + sin_cos=torch.zeros((num_envs, 4), dtype=torch.float32), + phase=torch.zeros((num_envs, 2), dtype=torch.float32), + steps=torch.zeros((num_envs, 1), dtype=torch.float32), ) diff --git a/motrix_envs/src/motrix_envs/locomotion/humanoid/walk_manager_mdp/reset.py b/motrix_envs/src/motrix_envs/locomotion/humanoid/walk_manager_mdp/reset.py index bee2cbcb..9cb581db 100644 --- a/motrix_envs/src/motrix_envs/locomotion/humanoid/walk_manager_mdp/reset.py +++ b/motrix_envs/src/motrix_envs/locomotion/humanoid/walk_manager_mdp/reset.py @@ -3,6 +3,8 @@ """Simulator reset terms for the humanoid velocity-tracking task.""" +from dataclasses import dataclass + import numpy as np import torch @@ -11,7 +13,6 @@ ManagerContext, ResetTerm, ResetTermCfg, - kernel_data, ) from motrix_env_core.manager.context import BuildContext from motrix_env_core.manager.kernel_data.map import Map @@ -33,7 +34,7 @@ from motrix_envs.locomotion.humanoid.walk_manager_mdp.randomization import WalkRandomizationCfg -@kernel_data +@dataclass(frozen=True, slots=True) class WalkResetParams: """Reset parameters, including in-kernel rough-terrain spawn sampling. @@ -71,24 +72,24 @@ class WalkResetParams: base center of mass. """ - default_joint_angles: np.ndarray - init_pose: np.ndarray + default_joint_angles: torch.Tensor + init_pose: torch.Tensor heightfield: HeightFieldGrid spawn_range: np.float32 randomization_enabled: bool - kp_default: np.ndarray - damping_default: np.ndarray - link_mass_default: np.ndarray + kp_default: torch.Tensor + damping_default: torch.Tensor + link_mass_default: torch.Tensor base_link_index: np.int64 base_mass_default: np.float32 - base_com_default: np.ndarray - friction_default: np.ndarray - kp_range: np.ndarray - damping_range: np.ndarray - friction_range: np.ndarray - mass_scale_range: np.ndarray - base_mass_off_range: np.ndarray - com_noise: np.ndarray + base_com_default: torch.Tensor + friction_default: torch.Tensor + kp_range: torch.Tensor + damping_range: torch.Tensor + friction_range: torch.Tensor + mass_scale_range: torch.Tensor + base_mass_off_range: torch.Tensor + com_noise: torch.Tensor def _sample_spawn_xy(ctx: ManagerContext, params: WalkResetParams) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: @@ -226,8 +227,10 @@ def __call__(self, ctx: BuildContext) -> ResetTerm: "joints_velocity": JointVelocityWrite(joint_names), } params = dict( - default_joint_angles=body.init_joint_pos, - init_pose=np.concatenate([body.init_base_position, body.init_base_quat]).astype(np.float32), + default_joint_angles=torch.from_numpy(np.ascontiguousarray(body.init_joint_pos, dtype=np.float32)), + init_pose=torch.from_numpy( + np.concatenate([body.init_base_position, body.init_base_quat]).astype(np.float32) + ), heightfield=ground_height_grid(ctx, self.ground_geom), spawn_range=np.float32(self.spawn_xy_range), ) @@ -250,41 +253,45 @@ def __call__(self, ctx: BuildContext) -> ResetTerm: link_masses = np.asarray(others["randomize_link_masses"], dtype=np.float32) params.update( randomization_enabled=True, - kp_default=np.asarray(others["actuator_kp"], dtype=np.float32), - damping_default=np.asarray(others["randomize_actuator_kd"], dtype=np.float32), - link_mass_default=link_masses, + kp_default=torch.from_numpy(np.asarray(others["actuator_kp"], dtype=np.float32)), + damping_default=torch.from_numpy(np.asarray(others["randomize_actuator_kd"], dtype=np.float32)), + link_mass_default=torch.from_numpy(link_masses), base_link_index=np.int64(base_index), base_mass_default=np.float32(link_masses[base_index]), - base_com_default=np.asarray(others["randomize_base_com"], dtype=np.float32), - friction_default=np.asarray(others["randomize_friction_default"], dtype=np.float32), - kp_range=np.asarray(randomization.kp_scale_range, dtype=np.float32), - damping_range=np.asarray(randomization.damping_scale_range, dtype=np.float32), - friction_range=np.asarray( - (1.0, 1.0) - if randomization.sliding_friction_range is None - else randomization.sliding_friction_range, - dtype=np.float32, + base_com_default=torch.from_numpy(np.asarray(others["randomize_base_com"], dtype=np.float32)), + friction_default=torch.from_numpy(np.asarray(others["randomize_friction_default"], dtype=np.float32)), + kp_range=torch.from_numpy(np.asarray(randomization.kp_scale_range, dtype=np.float32)), + damping_range=torch.from_numpy(np.asarray(randomization.damping_scale_range, dtype=np.float32)), + friction_range=torch.from_numpy( + np.asarray( + (1.0, 1.0) + if randomization.sliding_friction_range is None + else randomization.sliding_friction_range, + dtype=np.float32, + ) + ), + mass_scale_range=torch.from_numpy(np.asarray(randomization.link_mass_scale_range, dtype=np.float32)), + base_mass_off_range=torch.from_numpy( + np.asarray(randomization.base_mass_offset_range, dtype=np.float32) ), - mass_scale_range=np.asarray(randomization.link_mass_scale_range, dtype=np.float32), - base_mass_off_range=np.asarray(randomization.base_mass_offset_range, dtype=np.float32), - com_noise=np.asarray(randomization.base_com_offset_noise, dtype=np.float32), + com_noise=torch.from_numpy(np.asarray(randomization.base_com_offset_noise, dtype=np.float32)), ) else: params.update( randomization_enabled=False, - kp_default=np.zeros(0, dtype=np.float32), - damping_default=np.zeros(0, dtype=np.float32), - link_mass_default=np.zeros(0, dtype=np.float32), + kp_default=torch.zeros(0, dtype=torch.float32), + damping_default=torch.zeros(0, dtype=torch.float32), + link_mass_default=torch.zeros(0, dtype=torch.float32), base_link_index=np.int64(0), base_mass_default=np.float32(0.0), - base_com_default=np.zeros(3, dtype=np.float32), - friction_default=np.zeros(3, dtype=np.float32), - kp_range=np.ones(2, dtype=np.float32), - damping_range=np.ones(2, dtype=np.float32), - friction_range=np.ones(2, dtype=np.float32), - mass_scale_range=np.ones(2, dtype=np.float32), - base_mass_off_range=np.zeros(2, dtype=np.float32), - com_noise=np.zeros(3, dtype=np.float32), + base_com_default=torch.zeros(3, dtype=torch.float32), + friction_default=torch.zeros(3, dtype=torch.float32), + kp_range=torch.ones(2, dtype=torch.float32), + damping_range=torch.ones(2, dtype=torch.float32), + friction_range=torch.ones(2, dtype=torch.float32), + mass_scale_range=torch.ones(2, dtype=torch.float32), + base_mass_off_range=torch.zeros(2, dtype=torch.float32), + com_noise=torch.zeros(3, dtype=torch.float32), ) kernel = reset_walk_state_randomized if randomization.enabled else reset_walk_state diff --git a/motrix_envs/src/motrix_envs/locomotion/humanoid/walk_manager_mdp/terrain.py b/motrix_envs/src/motrix_envs/locomotion/humanoid/walk_manager_mdp/terrain.py index c6081055..13aea2cb 100644 --- a/motrix_envs/src/motrix_envs/locomotion/humanoid/walk_manager_mdp/terrain.py +++ b/motrix_envs/src/motrix_envs/locomotion/humanoid/walk_manager_mdp/terrain.py @@ -4,6 +4,7 @@ """Walk-task builder for the shared in-kernel terrain grid.""" import numpy as np +import torch from motrix_env_core.mdp.terrain import HeightFieldGrid @@ -18,19 +19,19 @@ def ground_height_grid(ctx, ground_geom: str) -> HeightFieldGrid: if cfg.ground_heightfield_geom is not None: hf = ctx.model.others["ground_heightfield"] return HeightFieldGrid( - heights=hf["heights"], - origin=hf["origin"], - spacing=hf["spacing"], - z0=hf["z0"], + heights=torch.from_numpy(np.ascontiguousarray(hf["heights"], dtype=np.float32)), + origin=torch.from_numpy(np.ascontiguousarray(hf["origin"], dtype=np.float32)), + spacing=torch.from_numpy(np.ascontiguousarray(hf["spacing"], dtype=np.float32)), + z0=torch.from_numpy(np.ascontiguousarray(hf["z0"], dtype=np.float32)), constant=np.float32(0.0), enabled=True, ) spec = ctx.model.others["geoms"][ground_geom] return HeightFieldGrid( - heights=np.zeros((1, 1), dtype=np.float32), - origin=np.zeros((2,), dtype=np.float32), - spacing=np.ones((2,), dtype=np.float32), - z0=np.zeros((1,), dtype=np.float32), + heights=torch.zeros((1, 1), dtype=torch.float32), + origin=torch.zeros((2,), dtype=torch.float32), + spacing=torch.ones((2,), dtype=torch.float32), + z0=torch.zeros((1,), dtype=torch.float32), constant=np.float32(spec.local_pose[2]), enabled=False, ) diff --git a/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/command.py b/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/command.py index 2b9fc191..17a7daa2 100644 --- a/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/command.py +++ b/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/command.py @@ -3,6 +3,8 @@ """Command terms for manager-based whole-body tracking.""" +from dataclasses import dataclass + import numpy as np import torch from omegaconf import MISSING @@ -14,7 +16,6 @@ CommandTerm, ManagerContext, ManagerEnv, - kernel_data, metric, ) from motrix_env_core.manager.commands import ResetContext @@ -53,7 +54,7 @@ def _sample_motion_step( return torch.where(zero_unit[:, 0] < start_zero_prob, torch.zeros_like(steps), steps) -@kernel_data +@dataclass(frozen=True, slots=True) class WbtMotionCommand(CommandTerm): """Drive each environment through a shared whole-body reference-motion clip. @@ -96,13 +97,15 @@ class WbtMotionCommand(CommandTerm): # Runtime buffers and tracked-body alignment metadata. reference_index: np.int64 - target_body_position_relative: np.ndarray - target_body_orientation_relative: np.ndarray + target_body_position_relative: torch.Tensor + target_body_orientation_relative: torch.Tensor - # Host-side adaptive frame-sampling state and parameters. - adaptive_bin_failed_count: np.ndarray - adaptive_current_bin_failed_count: np.ndarray - sampling_cdf: np.ndarray + # Host-side adaptive frame-sampling state and parameters. Host hooks read + # and write them through zero-copy ``.numpy()`` views; kernels write the + # per-env members through the tensors themselves. + adaptive_bin_failed_count: torch.Tensor + adaptive_current_bin_failed_count: torch.Tensor + sampling_cdf: torch.Tensor start_at_timestep_zero_prob: np.float32 hold_at_clip_end: bool sequential_clips: bool @@ -115,22 +118,24 @@ class WbtMotionCommand(CommandTerm): # Kept as ``(num_envs, 1)`` per-env arrays: the kernel lowering hands each # lane a writable row view, so advance/reset_env can update the lane's # frame in place without a shared backing. - steps: np.ndarray = metric(name="motion_step", dtype=np.float32) - clip_ended: np.ndarray = metric() + steps: torch.Tensor = metric(name="motion_step", dtype=np.float32) + clip_ended: torch.Tensor = metric() def reset(self, ctx: ResetContext) -> None: """Update adaptive statistics and prepare sampling for selected environments.""" env_ids = ctx.env_ids - if self.sampling_cdf.size: + cdf = self.sampling_cdf.numpy() + if cdf.size: + failed = self.adaptive_bin_failed_count.numpy() episode_failed = ctx.terminated[env_ids] if np.any(episode_failed): - failed_steps = self.steps[env_ids, 0][episode_failed] - num_bins = self.adaptive_bin_failed_count.size + failed_steps = self.steps.numpy()[env_ids, 0][episode_failed] + num_bins = failed.size failed_bins = (failed_steps.astype(np.int64) * num_bins) // max(self.clip.joint_pos.shape[0], 1) failed_bins = np.clip(failed_bins, 0, num_bins - 1) - self.adaptive_current_bin_failed_count[:] += np.bincount(failed_bins, minlength=num_bins).astype( - np.float32 - ) + self.adaptive_current_bin_failed_count.numpy()[:] += np.bincount( + failed_bins, minlength=num_bins + ).astype(np.float32) probabilities = self._sampling_probabilities() entropy = -float(np.sum(probabilities * np.log(probabilities + 1e-12))) entropy_norm = entropy / float(np.log(max(probabilities.size, 2))) @@ -138,7 +143,7 @@ def reset(self, ctx: ResetContext) -> None: ctx.metrics["adaptive_sampling_entropy"] = entropy_norm ctx.metrics["adaptive_sampling_top1_prob"] = float(probabilities[top_bin]) ctx.metrics["adaptive_sampling_top1_bin"] = float(top_bin) / float(probabilities.size) - ctx.metrics["adaptive_failure_mass"] = float(np.sum(self.adaptive_bin_failed_count)) + ctx.metrics["adaptive_failure_mass"] = float(np.sum(failed)) num_frames = self.clip.joint_pos.shape[0] max_step = num_frames - 1 max_bin = int(np.clip(np.ceil(max_step * probabilities.size / max(num_frames, 1)), 1, probabilities.size)) @@ -149,22 +154,23 @@ def reset(self, ctx: ResetContext) -> None: probabilities[:max_bin] = 1.0 / max_bin else: probabilities /= total - self.sampling_cdf[:] = np.cumsum(probabilities, dtype=np.float32) - self.sampling_cdf[-1] = 1.0 + cdf[:] = np.cumsum(probabilities, dtype=np.float32) + cdf[-1] = 1.0 def on_transition(self) -> None: """Fold accumulated failures into the adaptive sampler once per step.""" - if self.adaptive_bin_failed_count.size: - self.adaptive_bin_failed_count[:] = ( - self.alpha * self.adaptive_current_bin_failed_count - + (np.float32(1.0) - self.alpha) * self.adaptive_bin_failed_count + failed = self.adaptive_bin_failed_count.numpy() + if failed.size: + failed[:] = self.alpha * self.adaptive_current_bin_failed_count.numpy() + (np.float32(1.0) - self.alpha) * ( + failed ) - self.adaptive_current_bin_failed_count.fill(0.0) + self.adaptive_current_bin_failed_count.numpy().fill(0.0) def _sampling_probabilities(self) -> np.ndarray: """Build normalized frame-bin probabilities from failure history.""" - num_bins = self.adaptive_bin_failed_count.size - probability = self.adaptive_bin_failed_count + self.uniform_ratio / num_bins + failed = self.adaptive_bin_failed_count.numpy() + num_bins = failed.size + probability = failed + self.uniform_ratio / num_bins kernel_size = max(self.kernel_size, 1) if kernel_size > 1: kernel = np.asarray([self.kernel_lambda**i for i in range(kernel_size)], dtype=np.float32) @@ -316,14 +322,16 @@ def __call__(self, env: ManagerEnv) -> CommandTerm: return WbtMotionCommand( clip=source, reference_index=np.int64(reference_index), - steps=np.zeros((env.num_envs, 1), dtype=np.int64), - clip_ended=np.zeros((env.num_envs, 1), dtype=bool), - command=np.empty((env.num_envs, 2 * env.num_actuators), dtype=np.float32), - target_body_position_relative=np.empty((env.num_envs, *tracked_shape), dtype=np.float32), - target_body_orientation_relative=np.empty((env.num_envs, tracked_shape[0], 4), dtype=np.float32), - adaptive_bin_failed_count=np.zeros((num_bins,), dtype=np.float32), - adaptive_current_bin_failed_count=np.zeros((num_bins,), dtype=np.float32), - sampling_cdf=np.ones((num_bins,), dtype=np.float32) if num_bins else np.empty((0,), dtype=np.float32), + steps=torch.zeros((env.num_envs, 1), dtype=torch.int64), + clip_ended=torch.zeros((env.num_envs, 1), dtype=torch.bool), + command=torch.empty((env.num_envs, 2 * env.num_actuators), dtype=torch.float32), + target_body_position_relative=torch.empty((env.num_envs, *tracked_shape), dtype=torch.float32), + target_body_orientation_relative=torch.empty((env.num_envs, tracked_shape[0], 4), dtype=torch.float32), + adaptive_bin_failed_count=torch.zeros((num_bins,), dtype=torch.float32), + adaptive_current_bin_failed_count=torch.zeros((num_bins,), dtype=torch.float32), + sampling_cdf=torch.ones((num_bins,), dtype=torch.float32) + if num_bins + else torch.empty((0,), dtype=torch.float32), start_at_timestep_zero_prob=np.float32(self.start_at_timestep_zero_prob), hold_at_clip_end=self.hold_at_clip_end, sequential_clips=self.sequential_clips, diff --git a/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/observations.py b/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/observations.py index 1017e8bf..5b0fc888 100644 --- a/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/observations.py +++ b/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/observations.py @@ -3,6 +3,8 @@ """WBT command-backed observation terms.""" +from dataclasses import dataclass + import numpy as np import torch @@ -12,7 +14,6 @@ ManagerContext, ObservationTermCfg, ObsTerm, - kernel_data, ) from motrix_env_core.math import quaternion_array from motrix_env_core.math.quaternion_array import relative_orientation, rotate_inverse @@ -118,9 +119,9 @@ def __call__(self, ctx) -> ObsTerm: return ObsTerm(6 * num_tracked, robot_body_orientation_obs) -@kernel_data +@dataclass(frozen=True, slots=True) class RelativePositionParams: - reference: np.ndarray + reference: torch.Tensor noise_amplitude: np.float32 @@ -162,7 +163,7 @@ def __call__(self, ctx) -> ObsTerm: if missing: raise ValueError(f"RobotCfg {robot_name!r} key pose is missing queried joints: {missing}.") reference = np.asarray([positions[name] for name in queried_joints], dtype=np.float32) - params = RelativePositionParams(reference, np.float32(self.noise.amplitude)) + params = RelativePositionParams(torch.from_numpy(reference), np.float32(self.noise.amplitude)) return ObsTerm(reference.shape[0], dof_pos_rel_obs, params) diff --git a/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/rewards.py b/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/rewards.py index 4cba9ac0..b7e06177 100644 --- a/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/rewards.py +++ b/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/rewards.py @@ -3,6 +3,7 @@ """WBT command-backed reward terms.""" +from dataclasses import dataclass from typing import cast import numpy as np @@ -10,7 +11,6 @@ from motrix_env_core.config import configclass from motrix_env_core.manager import ManagerContext, RewardTerm, RewardTermCfg -from motrix_env_core.manager.kernel_data import kernel_data from motrix_env_core.math.quaternion_array import rotation_distance from motrix_env_core.mdp.action import JointPositionActionState @@ -146,10 +146,10 @@ def __call__(self, ctx) -> RewardTerm: return RewardTerm(global_body_angular_velocity_reward, np.float32(self.sigma)) -@kernel_data +@dataclass(frozen=True, slots=True) class DofLimitParams: - midpoint: np.ndarray - half_range: np.ndarray + midpoint: torch.Tensor + half_range: torch.Tensor soft_limit: np.float32 cap: np.float32 diff --git a/motrix_envs/src/motrix_envs/motion/library.py b/motrix_envs/src/motrix_envs/motion/library.py index afa147c5..21db305b 100644 --- a/motrix_envs/src/motrix_envs/motion/library.py +++ b/motrix_envs/src/motrix_envs/motion/library.py @@ -15,6 +15,7 @@ from pathlib import Path import numpy as np +import torch from motrix_env_core.manager.kernel_data import Map from motrix_envs.motion.loader import MotrixMotion @@ -104,15 +105,13 @@ def assemble(self) -> WbtMotionClip: clip_offsets = np.zeros_like(clip_lengths) np.cumsum(clip_lengths[:-1], out=clip_offsets[1:]) return WbtMotionClip( - **{ - field: np.concatenate([getattr(segment, field) for segment in segments]) for field in _FRAME_AXIS_FIELDS - }, - frame_clip_end=np.concatenate( + **{field: torch.cat([getattr(segment, field) for segment in segments]) for field in _FRAME_AXIS_FIELDS}, + frame_clip_end=torch.cat( [segment.frame_clip_end + offset for segment, offset in zip(segments, clip_offsets)] ), extensions=Map( { - name: MotionChannel(np.concatenate([segment.extensions[name].data for segment in segments])) + name: MotionChannel(torch.cat([segment.extensions[name].data for segment in segments])) for name in self.extension_channels } ), diff --git a/motrix_envs/src/motrix_envs/motion/tracked.py b/motrix_envs/src/motrix_envs/motion/tracked.py index d3b2398d..f741f25b 100644 --- a/motrix_envs/src/motrix_envs/motion/tracked.py +++ b/motrix_envs/src/motrix_envs/motion/tracked.py @@ -5,9 +5,12 @@ from __future__ import annotations +from dataclasses import dataclass + import numpy as np +import torch -from motrix_env_core.manager.kernel_data import Map, kernel_data +from motrix_env_core.manager.kernel_data import Map from motrix_envs.motion.loader import MotionFormatError, MotrixMotion @@ -38,7 +41,7 @@ def load_extension_channel(motion: MotrixMotion, name: str) -> np.ndarray: return channel -@kernel_data +@dataclass(frozen=True, slots=True) class MotionChannel: """One declared ``ext_`` channel, aligned with the clip's frame axis. @@ -46,10 +49,10 @@ class MotionChannel: channel is wrapped in this record whose field is shared across lanes. """ - data: np.ndarray + data: torch.Tensor -@kernel_data +@dataclass(frozen=True, slots=True) class WbtMotionClip: """Numeric whole-body-tracking views over one clip or a concatenated corpus. @@ -80,19 +83,19 @@ class WbtMotionClip: axis; empty when the task declares none. """ - joint_pos: np.ndarray - joint_vel: np.ndarray - tracked_bodies_pos_w: np.ndarray - tracked_bodies_quat_w: np.ndarray - tracked_bodies_lin_vel_w: np.ndarray - tracked_bodies_ang_vel_w: np.ndarray - root_body_pos_w: np.ndarray - root_body_quat_w: np.ndarray - root_body_lin_vel_w: np.ndarray - root_body_ang_vel_w: np.ndarray - reference_body_pos_w: np.ndarray - reference_body_quat_w: np.ndarray - frame_clip_end: np.ndarray + joint_pos: torch.Tensor + joint_vel: torch.Tensor + tracked_bodies_pos_w: torch.Tensor + tracked_bodies_quat_w: torch.Tensor + tracked_bodies_lin_vel_w: torch.Tensor + tracked_bodies_ang_vel_w: torch.Tensor + root_body_pos_w: torch.Tensor + root_body_quat_w: torch.Tensor + root_body_lin_vel_w: torch.Tensor + root_body_ang_vel_w: torch.Tensor + reference_body_pos_w: torch.Tensor + reference_body_quat_w: torch.Tensor + frame_clip_end: torch.Tensor extensions: Map @staticmethod @@ -129,18 +132,23 @@ def create( reference_index = motion.body_index(reference_body_name) return WbtMotionClip( - joint_pos=np.ascontiguousarray(motion.joint_pos[:, joint_motion_idx]), - joint_vel=np.ascontiguousarray(motion.joint_vel[:, joint_motion_idx]), - tracked_bodies_pos_w=np.ascontiguousarray(motion.body_pos_w[:, tracked_idx]), - tracked_bodies_quat_w=np.ascontiguousarray(motion.body_quat_w[:, tracked_idx]), - tracked_bodies_lin_vel_w=np.ascontiguousarray(motion.body_lin_vel_w[:, tracked_idx]), - tracked_bodies_ang_vel_w=np.ascontiguousarray(motion.body_ang_vel_w[:, tracked_idx]), - root_body_pos_w=np.ascontiguousarray(motion.body_pos_w[:, root_index]), - root_body_quat_w=np.ascontiguousarray(motion.body_quat_w[:, root_index]), - root_body_lin_vel_w=np.ascontiguousarray(motion.body_lin_vel_w[:, root_index]), - root_body_ang_vel_w=np.ascontiguousarray(motion.body_ang_vel_w[:, root_index]), - reference_body_pos_w=np.ascontiguousarray(motion.body_pos_w[:, reference_index]), - reference_body_quat_w=np.ascontiguousarray(motion.body_quat_w[:, reference_index]), - frame_clip_end=np.full(motion.num_frames, motion.num_frames - 1, dtype=np.int64), - extensions=Map({name: MotionChannel(load_extension_channel(motion, name)) for name in extension_channels}), + joint_pos=torch.from_numpy(np.ascontiguousarray(motion.joint_pos[:, joint_motion_idx])), + joint_vel=torch.from_numpy(np.ascontiguousarray(motion.joint_vel[:, joint_motion_idx])), + tracked_bodies_pos_w=torch.from_numpy(np.ascontiguousarray(motion.body_pos_w[:, tracked_idx])), + tracked_bodies_quat_w=torch.from_numpy(np.ascontiguousarray(motion.body_quat_w[:, tracked_idx])), + tracked_bodies_lin_vel_w=torch.from_numpy(np.ascontiguousarray(motion.body_lin_vel_w[:, tracked_idx])), + tracked_bodies_ang_vel_w=torch.from_numpy(np.ascontiguousarray(motion.body_ang_vel_w[:, tracked_idx])), + root_body_pos_w=torch.from_numpy(np.ascontiguousarray(motion.body_pos_w[:, root_index])), + root_body_quat_w=torch.from_numpy(np.ascontiguousarray(motion.body_quat_w[:, root_index])), + root_body_lin_vel_w=torch.from_numpy(np.ascontiguousarray(motion.body_lin_vel_w[:, root_index])), + root_body_ang_vel_w=torch.from_numpy(np.ascontiguousarray(motion.body_ang_vel_w[:, root_index])), + reference_body_pos_w=torch.from_numpy(np.ascontiguousarray(motion.body_pos_w[:, reference_index])), + reference_body_quat_w=torch.from_numpy(np.ascontiguousarray(motion.body_quat_w[:, reference_index])), + frame_clip_end=torch.from_numpy(np.full(motion.num_frames, motion.num_frames - 1, dtype=np.int64)), + extensions=Map( + { + name: MotionChannel(torch.from_numpy(load_extension_channel(motion, name))) + for name in extension_channels + } + ), ) diff --git a/motrix_envs/tests/test_motion_library.py b/motrix_envs/tests/test_motion_library.py index e93c915b..8979675a 100644 --- a/motrix_envs/tests/test_motion_library.py +++ b/motrix_envs/tests/test_motion_library.py @@ -3,6 +3,7 @@ import numpy as np import pytest +import torch from motrix_envs.motion import MotionLibrary, MotrixMotion, WbtMotionClip from motrix_envs.motion.loader import MotionFormatError @@ -110,7 +111,7 @@ def test_library_emits_per_frame_clip_boundaries(tmp_path): clip = MotionLibrary([tmp_path / "a.npz", tmp_path / "b.npz", tmp_path / "c.npz"], **_TASK_ORDER).assemble() np.testing.assert_array_equal(clip.frame_clip_end, [5] * 6 + [9] * 4 + [14] * 5) - assert clip.frame_clip_end.dtype == np.int64 + assert clip.frame_clip_end.dtype == torch.int64 @pytest.mark.parametrize("directory_source", [False, True]) @@ -276,7 +277,7 @@ def test_wbt_motion_clip_loads_declared_extension_channel(tmp_path): extension_channels=("smpl_joints",), ) np.testing.assert_array_equal(clip.extensions["smpl_joints"].data, motion.extensions["smpl_joints"]) - assert clip.extensions["smpl_joints"].data.flags.c_contiguous + assert clip.extensions["smpl_joints"].data.is_contiguous() with pytest.raises(MotionFormatError, match="does not provide declared extension channel"): WbtMotionClip.create( diff --git a/motrix_envs/tests/test_motion_loader.py b/motrix_envs/tests/test_motion_loader.py index b3db183d..c3325475 100644 --- a/motrix_envs/tests/test_motion_loader.py +++ b/motrix_envs/tests/test_motion_loader.py @@ -147,7 +147,7 @@ def test_wbt_motion_clip_builds_name_ordered_numeric_views(tmp_path): joint_pos = wbt.joint_pos motion.joint_pos = np.zeros_like(motion.joint_pos) assert wbt.joint_pos is joint_pos - assert np.any(wbt.joint_pos) + assert bool(wbt.joint_pos.any()) array_fields = ( "joint_pos", @@ -163,4 +163,4 @@ def test_wbt_motion_clip_builds_name_ordered_numeric_views(tmp_path): "reference_body_pos_w", "reference_body_quat_w", ) - assert all(getattr(wbt, name).flags.c_contiguous for name in array_fields) + assert all(getattr(wbt, name).is_contiguous() for name in array_fields) diff --git a/motrix_envs/tests/test_walk_randomization.py b/motrix_envs/tests/test_walk_randomization.py index 3ca8d53a..d2e551aa 100644 --- a/motrix_envs/tests/test_walk_randomization.py +++ b/motrix_envs/tests/test_walk_randomization.py @@ -9,6 +9,7 @@ import gymnasium as gym import numpy as np import pytest +import torch from motrix_env_core.manager import ManagerEnv from motrix_env_core.mdp.action import JointPositionActionCfg, JointPositionActionState, JointPositionActionTerm @@ -29,13 +30,13 @@ def test_walk_randomization_cfg_rejects_invalid_ranges(): def _tiny_delay_manager(num_envs: int = 2, num_actuators: int = 3, lo: int = 0, hi: int = 2): action_state = JointPositionActionState( - action_queue=np.zeros((num_envs, max(hi + 1, 2), num_actuators), dtype=np.float32), - default_angles=np.zeros(num_actuators, dtype=np.float32), - joint_lower=np.zeros(num_actuators, dtype=np.float32), - joint_upper=np.ones(num_actuators, dtype=np.float32), - action_scales=np.ones(num_actuators, dtype=np.float32), - delay_steps=np.full(num_envs, hi, dtype=np.int64), - action_ptr=np.zeros(1, dtype=np.int64), + action_queue=torch.from_numpy(np.zeros((num_envs, max(hi + 1, 2), num_actuators), dtype=np.float32)), + default_angles=torch.from_numpy(np.zeros(num_actuators, dtype=np.float32)), + joint_lower=torch.from_numpy(np.zeros(num_actuators, dtype=np.float32)), + joint_upper=torch.from_numpy(np.ones(num_actuators, dtype=np.float32)), + action_scales=torch.from_numpy(np.ones(num_actuators, dtype=np.float32)), + delay_steps=torch.from_numpy(np.full(num_envs, hi, dtype=np.int64)), + action_ptr=torch.from_numpy(np.zeros(1, dtype=np.int64)), delay_lo=lo, delay_hi=hi, ) @@ -68,7 +69,7 @@ def test_action_delay_applies_delayed_action(): env.apply_action(actions, env._state) # Only the raw action exists so far; unwritten history remains zero. assert np.all(controls == 0.0) - assert np.all(action_state.action_queue[:, action_state.action_ptr[0]] == 1.0) + assert bool((action_state.action_queue[:, int(action_state.action_ptr[0])] == 1.0).all()) for value in (2.0, 3.0, 4.0): env.apply_action(np.full_like(actions, value), env._state) # Two-step delay: the applied action is the one from two steps ago. @@ -81,13 +82,13 @@ def test_action_delay_disabled_passes_actions_through_and_keeps_history(): env.apply_action(actions, env._state) assert np.all(controls == 0.5) ptr = int(action_state.action_ptr[0]) - assert np.all(action_state.action_queue[:, ptr] == 0.5) - assert np.all(action_state.action_queue[:, (ptr - 1) % action_state.action_queue.shape[1]] == 0.0) + assert bool((action_state.action_queue[:, ptr] == 0.5).all()) + assert bool((action_state.action_queue[:, (ptr - 1) % action_state.action_queue.shape[1]] == 0.0).all()) def test_action_queue_wraps_and_selects_delay_per_lane(): env, action_state, controls = _tiny_delay_manager(num_envs=2, lo=0, hi=2) - action_state.delay_steps[:] = [0, 2] + action_state.delay_steps[:] = torch.tensor([0, 2]) width = action_state.action_queue.shape[1] for value in range(1, width + 2): actions = np.full(controls.shape, value, dtype=np.float32) @@ -102,18 +103,18 @@ def test_action_reset_clears_only_selected_lane_history_and_keeps_pointer(): env, action_state, controls = _tiny_delay_manager(num_envs=2, lo=0, hi=2) for value in (1.0, 2.0, 3.0): env.apply_action(np.full(controls.shape, value, dtype=np.float32), env._state) - ptr_before_reset = action_state.action_ptr.copy() - other_lane_history = action_state.action_queue[1].copy() + ptr_before_reset = action_state.action_ptr.clone() + other_lane_history = action_state.action_queue[1].clone() env.reset(np.asarray([0], dtype=np.int64)) - assert np.all(action_state.action_queue[0] == 0.0) - np.testing.assert_array_equal(action_state.action_queue[1], other_lane_history) - np.testing.assert_array_equal(action_state.action_ptr, ptr_before_reset) + assert bool((action_state.action_queue[0] == 0.0).all()) + torch.testing.assert_close(action_state.action_queue[1], other_lane_history, rtol=0, atol=0) + torch.testing.assert_close(action_state.action_ptr, ptr_before_reset, rtol=0, atol=0) def test_action_delay_reset_resamples_within_range(): env, action_state, _ = _tiny_delay_manager(lo=0, hi=2) env.reset(np.arange(action_state.action_queue.shape[0])) - assert np.all((action_state.delay_steps >= 0) & (action_state.delay_steps <= 2)) + assert bool(((action_state.delay_steps >= 0) & (action_state.delay_steps <= 2)).all()) def test_joint_position_term_methods_manage_history_and_specific_state(): @@ -124,10 +125,10 @@ def test_joint_position_term_methods_manage_history_and_specific_state(): pointer = int(action_state.action_ptr[0]) term.process(np.full(controls.shape, 4.0, dtype=np.float32)) assert int(action_state.action_ptr[0]) == (pointer + 1) % action_state.action_queue.shape[1] - np.testing.assert_array_equal(action_state.current(), 4.0) + torch.testing.assert_close(action_state.current(), torch.full_like(action_state.current(), 4.0)) term.reset(np.asarray([0], dtype=np.int64)) - np.testing.assert_array_equal(action_state.action_queue[0], 0.0) - np.testing.assert_array_equal(action_state.action_queue[1, 0], 3.0) + assert bool((action_state.action_queue[0] == 0.0).all()) + assert bool((action_state.action_queue[1, 0] == 3.0).all()) def test_flat_cfg_randomization_settings_are_wired(): @@ -160,10 +161,10 @@ def test_walk_env_randomized_reset_samples_within_ranges(): max_period = max(command_cfg.gait_period + command_cfg.gait_period_randomization_width, 0.1) min_phase_step = 2.0 * np.pi * cfg.ctrl_dt / max_period max_phase_step = 2.0 * np.pi * cfg.ctrl_dt / min_period - assert np.all(walk.phase_step >= min_phase_step - 1e-6) - assert np.all(walk.phase_step <= max_phase_step + 1e-6) + assert bool((walk.phase_step.numpy() >= min_phase_step - 1e-6).all()) + assert bool((walk.phase_step.numpy() <= max_phase_step + 1e-6).all()) action_state = env.action_terms["joint_position"].state delay_lo, delay_hi = cfg.actions.joint_position.action_delay_steps - assert np.all((action_state.delay_steps >= delay_lo) & (action_state.delay_steps <= delay_hi)) + assert bool(((action_state.delay_steps >= delay_lo) & (action_state.delay_steps <= delay_hi)).all()) finally: del env diff --git a/wiki/design/manager/isaaclab-manager-cfg-adapter.md b/wiki/design/manager/isaaclab-manager-cfg-adapter.md index d2325c15..87a195b7 100644 --- a/wiki/design/manager/isaaclab-manager-cfg-adapter.md +++ b/wiki/design/manager/isaaclab-manager-cfg-adapter.md @@ -2,7 +2,7 @@ ## 摘要 -本设计定义一个面向 IsaacLab manager-based 任务配置的薄适配层,目标是尽量保留用户熟悉的 ManagerCfg 分组、`func + params`、term 命名、嵌套覆写和 selector 表达。Torch 批量自定义 MDP 函数在程序构建时被特化为私有 concrete `@kernel_data` Term,selector 被归一化为显式 `SimDataQuery`。适配层不直接执行任意 IsaacLab Torch callback,也不改变 [Manager Task API](./task-authoring.md) 和 [Manager Runtime](./runtime.md) 中的 Torch 执行器边界。 +本设计定义一个面向 IsaacLab manager-based 任务配置的薄适配层,目标是尽量保留用户熟悉的 ManagerCfg 分组、`func + params`、term 命名、嵌套覆写和 selector 表达。Torch 批量自定义 MDP 函数在程序构建时被特化为私有 concrete kernel-data dataclass Term,selector 被归一化为显式 `SimDataQuery`。适配层不直接执行任意 IsaacLab Torch callback,也不改变 [Manager Task API](./task-authoring.md) 和 [Manager Runtime](./runtime.md) 中的 Torch 执行器边界。 本文是目标设计,不表示下述 API 均已实现。兼容基线参考本地 IsaacLab 快照 `21bcb476b2`;该快照检查时落后 `origin/main` 202 个 commit,因此不声称覆盖所有上游版本。 @@ -14,7 +14,7 @@ IsaacLab-like ManagerCfg -> configuration normalization -> native TermCfg.__call__() / functional term generation - -> concrete @kernel_data Term + -> concrete kernel-data dataclass Term -> ManagerContext / SimDataQuery / ReadPlan -> Torch executor(torch.compile kernel,整批张量) ``` @@ -24,14 +24,14 @@ IsaacLab-like ManagerCfg - 用户可以继续按 action、command、observation、reward、termination 和 event 组织配置。 - 配置层可接受 IsaacLab 常见的嵌套属性修改和以 `None` 删除 term。 - 每个被接受的 term 最终必须成为 concrete runtime term;native typed cfg 直接创建,functional cfg 由框架在构建期生成。 -- [`@kernel_data`](./kernel-abi.md)、[`ManagerContext`](./kernel-abi.md) 和 [SimDataQuery 读侧编译](./sim-read-program.md) 的契约保持不变。 +- [kernel-data dataclass](./kernel-abi.md)、[`ManagerContext`](./kernel-abi.md) 和 [SimDataQuery 读侧编译](./sim-read-program.md) 的契约保持不变。 非目标: - 不直接执行 IsaacLab 中以 `env` 为参数的通用 Python/Torch `func`。 - 不允许 term 隐式读取整个 environment 或在 kernel 外隐藏 simulator 依赖。 - 不静默忽略 noise、clip、history 等未支持选项。 -- 不让 callable、`dict[str, Any]` 或完整 environment 进入 kernel_data ABI 或编译图。 +- 不让 callable、`dict[str, Any]` 或完整 environment 进入 kernel-data ABI 或编译图。 - 不为表面兼容牺牲 typed cfg、compiler-visible dependency 或批量 Torch kernel。 ## 2. 当前主要差异 @@ -106,7 +106,7 @@ def track_lin_vel_xy_exp( `func + params` normalizer 使用函数签名绑定 keyword params,再按参数值的实际类型生成 concrete Term。对上述配置,等价的内部结果是: ```python -@kernel_data +@dataclass(frozen=True, slots=True) class _GeneratedTrackLinVelXYExp: std: np.float32 @@ -126,9 +126,9 @@ class _GeneratedTrackLinVelXYExp: | --- | --- | | 数值标量 | 转换为类型确定的 Term field,例如 `np.float32` | | `str` 或 enum | 嵌入生成调用作为 compile-time literal;它在 MDP 函数中可以用作 command/map key 或普通常量 | -| `np.ndarray` | 转换为具有明确 shared/per-env scope 的 typed field | +| `np.ndarray` / `torch.Tensor` | 转换为具有明确 shared/per-env scope 的 typed tensor field | | `SceneEntityCfg` | 根据值的 selector 类型生成并去重 `SimDataQuery`,将解析后的 kernel 输入传给 MDP 计算 | -| nested `@kernel_data` | 作为 typed subtree flatten/lower | +| nested kernel-data record | 作为 typed subtree flatten/lower | | 任意 Python object、closure 或未支持容器 | 在构建期报出包含 group/term/param 路径的错误 | 生成的 Term class 是 compiler 实现细节,应按 `(term kind, func fingerprint, param schema, structural literals)` 缓存。数值参数只改变 Term 实例,不生成新类型;改变 string/enum/selector 属于结构配置变化,需要重建 program。调试输出和错误必须引用原始 `func` 与 `params` 路径,不向用户暴露 `_Generated*` 类名。 @@ -189,7 +189,7 @@ self.rewards.undesired_contacts = None - 提供 backend-neutral `SceneEntityCfg`,归一化为显式、可去重的 `SimDataQuery`。 - 保留 `.weight = ...` 与 `term = None` 等嵌套覆写方式。 - 对无状态 Reward/Termination 接受 Torch 批量 `func + params`,自动生成 concrete Term,不要求用户注册 factory。 -- `params` 字典只存在于 host config,不进入 kernel_data ABI 或编译图。 +- `params` 字典只存在于 host config,不进入 kernel-data ABI 或编译图。 ### P1:Observation 常用选项 diff --git a/wiki/design/manager/kernel-abi.md b/wiki/design/manager/kernel-abi.md index 9d028fdd..c22aa386 100644 --- a/wiki/design/manager/kernel-abi.md +++ b/wiki/design/manager/kernel-abi.md @@ -2,8 +2,8 @@ ## 摘要 -本文定义 Manager term 的数据 ABI:`@kernel_data` logical tree 如何解析为 -path-keyed 的 resolved tree,`ManagerContext` 如何以整批张量 +本文定义 Manager term 的数据 ABI:kernel-data dataclass(plain frozen dataclass) +logical tree 如何解析为 path-keyed 的 resolved tree,`ManagerContext` 如何以整批张量 提供固定 namespace,以及 metric backing 的 ownership 语义。任务开发者只需遵守 [Manager Task API](./task-authoring.md) 中的 term 协议;数值执行器见 [Torch 执行器设计](./torch-executor.md), backend 读侧和 simulator 边界见 [SimBackend](./sim-backend.md)。 @@ -13,15 +13,16 @@ backend 读侧和 simulator 边界见 [SimBackend](./sim-backend.md)。 `ManagerContext` 是 term dispatch 的唯一 manager/runtime 输入入口: ```python -@kernel_data +@dataclass(frozen=True, slots=True) class ManagerContext: - env_id: np.ndarray + env_id: torch.Tensor actions: Map commands: Map + metrics: Map[np.ndarray] rand: RandValue sim: Map dt: np.float32 - sim_reset_requested: np.ndarray + sim_reset_requested: torch.Tensor ``` Context 不包含 config、provider、backend handle、dynamic Python container 或 generic `states` / `resources` store。 @@ -55,7 +56,8 @@ dof_pos = ctx.sim["robot_dof_pos"] | 类型 | term 访问语义 | | --- | --- | -| `np.ndarray` | 完整 ndarray 引用,term 按自己的契约决定第 0 维是否为 batch 维 | +| `torch.Tensor` | term 拥有字段的默认类型。张量是宿主 numpy arena 之上的零拷贝视图(`torch.from_numpy`);host 侧代码可通过 `tensor.numpy()` 读写同一 backing | +| `np.ndarray` | 完整 ndarray 引用,term 按自己的契约决定第 0 维是否为 batch 维;SimBackend 读侧 arena(`ctx.sim`)与 metric backing 保持 ndarray | | Python scalar | 所有环境使用同一个标量 | | 固定 dtype NumPy scalar | 所有环境使用同一个定 dtype 标量 | @@ -64,14 +66,16 @@ dof_pos = ctx.sim["robot_dof_pos"] ## 4. Logical tree 与 resolved tree -`@kernel_data` 支持有序、可逆的 record tree: +Kernel-data record 是 plain frozen dataclass(`@dataclass(frozen=True, slots=True)`, +不再有专用装饰器),支持有序、可逆的 record tree: ```python leaves, tree_def = flatten_kernel_data(value) rebuilt = unflatten_kernel_data(tree_def, leaves) ``` -Record 可以嵌套另一个 KernelData record;不支持 multiple inheritance、覆盖 inherited field、普通 dataclass、递归结构、dict、list、字符串、object array 或 backend handle。 +Record 可以嵌套另一个 kernel-data record;不支持递归结构、dict、list、字符串、object array 或 backend handle +(SimDataQuery 是 dataclass 但通过 read program 解析,不属于 kernel-data tree)。 Logical tree 与 execution lowering 分离:同一 concrete value 在 host term 和 context 中保留相同的 tree schema, resolver 产出 path-keyed 的 resolved tree(leaf 引用),运行时据此构建 namespace 视图。只有 ndarray/scalar leaves 可以成为 resolved leaves;term 内不应产生稳态临时大数组分配。 diff --git a/wiki/design/manager/runtime.md b/wiki/design/manager/runtime.md index a2b85b13..124a474c 100644 --- a/wiki/design/manager/runtime.md +++ b/wiki/design/manager/runtime.md @@ -39,7 +39,7 @@ Manager 不是完整的 `step()` 替代品,也不编译物理推进。物理 ManagerBasedEnvCfg -> normalize manager groups -> create concrete terms - -> canonicalize @kernel_data terms + -> canonicalize kernel-data dataclass terms -> compile model/read/write/reset programs -> construct ManagerContext layout -> warm up generated kernels @@ -90,10 +90,10 @@ reset 只处理指定的原始 `env_ids`: Term 不要求继承统一实现基类,但必须遵循对应 protocol: ```python -@kernel_data +@dataclass(frozen=True, slots=True) class BaseActionState: - action_queue: np.ndarray # host (N, W, A), lane (W, A); W >= 2 - action_ptr: np.ndarray # one shared cursor, advanced once per control step + action_queue: torch.Tensor # host (N, W, A), lane (W, A); W >= 2 + action_ptr: torch.Tensor # one shared cursor, advanced once per control step @dataclass(frozen=True) diff --git a/wiki/design/manager/task-authoring.md b/wiki/design/manager/task-authoring.md index 043c5746..1e1d5200 100644 --- a/wiki/design/manager/task-authoring.md +++ b/wiki/design/manager/task-authoring.md @@ -3,7 +3,7 @@ ## 摘要 本文面向任务开发者,说明如何使用 `ManagerBasedEnvCfg` 和 Manager term 创建一个 config-driven `ManagerEnv`。 -内容覆盖任务文件组织、manager groups、term 创建、`SimDataQuery` 参数、`@kernel_data` 字段和常用 term 接口。 +内容覆盖任务文件组织、manager groups、term 创建、`SimDataQuery` 参数、kernel-data dataclass 字段和常用 term 接口。 运行时边界见 [Manager Runtime](./runtime.md),kernel ABI 见 [Manager Kernel ABI](./kernel-abi.md)。 ## 1. 任务文件组织 @@ -75,13 +75,13 @@ class MyTaskEnvCfg(ManagerBasedEnvCfg): ## 3. Term 创建与字段 -各 term cfg 通过 `__call__()` 创建 concrete runtime term。Concrete term 使用 `@kernel_data` 声明固定字段: +各 term cfg 通过 `__call__()` 创建 concrete runtime term。Concrete term 使用 plain frozen dataclass(kernel-data record)声明固定字段: ```python -@kernel_data +@dataclass(frozen=True, slots=True) class TrackingErrorTermination: threshold: np.float32 - error: np.ndarray = metric() + error: torch.Tensor = metric() def compute(self, ctx: ManagerContext) -> bool: value = ... @@ -89,12 +89,12 @@ class TrackingErrorTermination: return value > self.threshold ``` -`@kernel_data` 的字段必须有类型标注,结构在构造期固定。Term record 不可替换字段,但 owner 可以按契约原地更新 ndarray backing。 +kernel-data record 的字段必须有类型标注,结构在构造期固定。Term record 不可替换字段,但 owner 可以按契约原地更新 backing(tensor 写入或经 `tensor.numpy()` 视图写入)。 支持的 leaf 类型为: -- 普通 `np.ndarray`:term 拿到完整数组引用;lane 形状数组第一维为 `num_envs`,模型参数、查表数据和 motion clip 等共享数组则没有 batch 维; -- Python 标量和固定 dtype NumPy scalar; -- 另一个 `@kernel_data` 类型。 +- `torch.Tensor`:term 拥有字段的默认类型,是宿主 numpy arena 上的零拷贝视图;lane 形状张量第一维为 `num_envs`,模型参数、查表数据和 motion clip 等共享张量则没有 batch 维; +- 普通 `np.ndarray`(SimBackend 读侧数组与 metric backing)与 Python 标量、固定 dtype NumPy scalar; +- 另一个 kernel-data record 类型。 需要对外发布逐环境指标时使用 `metric()` 字段。框架为其分配 `(num_envs, 1)` backing,term 在当前 lane 写入 `[0]`。 @@ -133,7 +133,7 @@ ActionCfg 在构造期返回冻结的 host 描述对象: ```python ActionTerm( action_space=space, - state=state, # @kernel_data BaseActionState 子类 + state=state, # kernel-data dataclass(BaseActionState 子类) process=process_action, # process_action(state, actions_batch) -> controls | None reset=reset_action, # reset_action(state, env_ids) -> None ) @@ -160,7 +160,7 @@ Reward term 返回当前 environment 的 scalar,Manager 最终按配置 weight Termination term 返回单个 environment 的 `bool`: ```python -@kernel_data +@dataclass(frozen=True, slots=True) class OutOfBoundsTermination: limit: np.float32 diff --git a/wiki/design/manager/term-required-queries.md b/wiki/design/manager/term-required-queries.md index bc60317d..cbcf41be 100644 --- a/wiki/design/manager/term-required-queries.md +++ b/wiki/design/manager/term-required-queries.md @@ -73,13 +73,13 @@ term 因此不需要在 dispatch 内部用字符串 key 访问 `ctx.sim`:query ## 5. 边界与已知限制 -- 收集范围是 reward、termination 和 observation term 的 `args`;command/action term 的 kernel_data +- 收集范围是 reward、termination 和 observation term 的 `args`;command/action term 的 kernel-data dataclass 走各自的 canonicalize 路径,不参与本机制。 - 合成 key `terms.args[i]` 依赖收集顺序(配置映射顺序),属于内部执行元数据,不是面向任务的稳定命名 契约;任务需要稳定 key 时应在 `queries.data` 中显式声明。 - 不支持单环境多机器人的 per-entity key 命名空间;canonical key 是单数角色("the robot")。 - model query 是静态数据,不经 read program:term 在 `__call__(env)` 时从 `env.model` 读出并烘焙进 - term 的 `@kernel_data` params,与本机制无关。 + term 的 kernel-data params,与本机制无关。 ## 验收标准 diff --git a/wiki/design/manager/torch-executor.md b/wiki/design/manager/torch-executor.md index ba099dd5..e2cb9652 100644 --- a/wiki/design/manager/torch-executor.md +++ b/wiki/design/manager/torch-executor.md @@ -21,9 +21,9 @@ RNG draw 顺序与 eager 完全一致;Inductor 重结合仅带来 ≤1e-8 量 ## 解析与执行的分工 -- `ManagerProgramResolver`(backend-neutral)负责输入准备:把 `@kernel_data` - term 参数与 `ManagerContext` 解析为 path-keyed 的 resolved tree(叶子持有 - live value 引用),构建 `ManagerLayout`、 +- `ManagerProgramResolver`(backend-neutral)负责输入准备:把 kernel-data + dataclass term 参数与 `ManagerContext` 解析为 path-keyed 的 resolved tree(叶子持有 + live value 引用;数组字段已是 torch tensor,执行器无需再包装),构建 `ManagerLayout`、 `ReadPlan`、per-term-kind typed call records(command hooks 携带 receiver 与 static function,reward/termination 携带 buffer column index 与 resolved argument 描述符等),并校验签名注解契约与签名契约。它不做任何代码生成、 diff --git a/wiki/design/motion-multi-clip-infra.md b/wiki/design/motion-multi-clip-infra.md index c13f157b..c921d2c6 100644 --- a/wiki/design/motion-multi-clip-infra.md +++ b/wiki/design/motion-multi-clip-infra.md @@ -34,7 +34,7 @@ main 上唯一的 motion 消费者是 `WbtMotionCommand`,"单 clip"假设同 ### 数据平面:`MotionLibrary` -位置 `motrix_envs/src/motrix_envs/motion/`,kernel_data 纯数值对象。 +位置 `motrix_envs/src/motrix_envs/motion/`,kernel-data dataclass 纯数值对象(torch tensor 字段)。 - 输入:有序的 MotrixMotion npz v1 文件列表(cfg 显式列表或目录展开,顺序固定可复现); - 逐文件:`MotrixMotion` 加载校验(schema、fps、finite、单位四元数)→ 按任务 joint/body 顺序重排 → From 8f7425f1283645e5460bd068a5faae70b30a6e5c Mon Sep 17 00:00:00 2001 From: zilch <147668916@qq.com> Date: Fri, 6 Nov 2026 13:54:45 +0800 Subject: [PATCH 26/41] feat(manager)!: exchange env state with the RL layer as torch tensors MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. --- .../src/motrix_env_core/array/env.py | 73 +++++++++++++------ .../src/motrix_env_core/manager/env.py | 57 ++++++++++----- .../src/motrix_env_core/manager/kernel.py | 22 +++--- .../motrix_env_core/torch_manager/runtime.py | 37 +++++----- motrix_env_core/tests/test_manager_env.py | 22 +++--- .../tests/test_manager_sim_backend.py | 2 +- .../src/motrix_rl_builtin/fastsac/wrap_np.py | 15 ++-- .../src/motrix_rl_rslrl/torch/wrap_np.py | 35 +++++---- .../src/motrix_rl_skrl/jax/train/ppo.py | 3 + .../src/motrix_rl_skrl/jax/wrap_np.py | 28 ++++--- .../src/motrix_rl_skrl/torch/wrap_np.py | 29 +++++--- 11 files changed, 203 insertions(+), 120 deletions(-) diff --git a/motrix_env_core/src/motrix_env_core/array/env.py b/motrix_env_core/src/motrix_env_core/array/env.py index d76d2ab8..6412ef8e 100644 --- a/motrix_env_core/src/motrix_env_core/array/env.py +++ b/motrix_env_core/src/motrix_env_core/array/env.py @@ -8,6 +8,7 @@ import gymnasium as gym import numpy as np +import torch from motrix_env_core.base import ABEnv, EnvCfg, ObsSpace from motrix_env_core.perf import Perf, perf_root @@ -25,11 +26,14 @@ @dataclass class NpObs: - policy: np.ndarray - value: np.ndarray | None = None + """One observation slot; numpy arrays for NumPy-frontend environments, + torch tensors for the manager pipeline.""" + + policy: "torch.Tensor | np.ndarray" + value: "torch.Tensor | np.ndarray | None" = None @property - def value_or_policy(self) -> np.ndarray: + def value_or_policy(self) -> "torch.Tensor | np.ndarray": return self.policy if self.value is None else self.value @@ -47,15 +51,24 @@ class FinalObs: @dataclass class ArrayEnvState: + """Environment transition state. + + ``obs``, ``reward``, ``terminated``, ``truncated`` and ``final_obs`` store + numpy arrays for NumPy-frontend environments and torch tensors for the + manager pipeline; each environment class consistently produces one of the + two, so consumers never branch on the element type. + """ + obs: NpObs - reward: np.ndarray - terminated: np.ndarray - truncated: np.ndarray + reward: "torch.Tensor | np.ndarray" + terminated: "torch.Tensor | np.ndarray" + truncated: "torch.Tensor | np.ndarray" episode_steps: np.ndarray # Per-term reward breakdown for the latest transition: term name to a - # ``(num_envs,)`` array. Freshly written every transition; empty when the - # environment does not decompose its reward. - reward_terms: dict[str, np.ndarray] = dataclasses.field(default_factory=dict) + # ``(num_envs,)`` array (numpy for the NumPy frontend, torch column views + # for the manager pipeline). Freshly written every transition; empty when + # the environment does not decompose its reward. + reward_terms: dict[str, "torch.Tensor | np.ndarray"] = dataclasses.field(default_factory=dict) # Live diagnostics view: per-environment quantities stay unreduced as # ``(num_envs,)`` arrays and are views into manager buffers that kernels # overwrite every step; batch-level gauges are scalars. Values always @@ -66,11 +79,11 @@ class ArrayEnvState: final_obs: FinalObs | None = None @property - def done(self) -> np.ndarray: + def done(self) -> "torch.Tensor | np.ndarray": """ Check if the environment is done. """ - return np.logical_or(self.terminated, self.truncated) + return self.terminated | self.truncated def replace(self, **updates) -> "ArrayEnvState": return dataclasses.replace(self, **updates) @@ -163,9 +176,9 @@ def init_state(self) -> ArrayEnvState: Create a new environment state """ obs = self._zeros_obs(self.observation_space) - reward = np.zeros((self._num_envs,), dtype=np.float32) - terminated = np.ones((self._num_envs,), dtype=bool) - truncated = np.zeros((self._num_envs,), dtype=bool) + reward = self._zeros_reward() + terminated = self._zeros_terminated() + truncated = self._zeros_truncated() episode_steps = np.zeros((self._num_envs,), dtype=np.uint64) self._state = self._new_state(obs, reward, terminated, truncated, episode_steps) self._reset_done_envs() @@ -207,6 +220,15 @@ def has_value_observation(self) -> bool: def _zeros_box(self, space: gym.spaces.Box) -> np.ndarray: return np.zeros((self._num_envs, *space.shape), dtype=space.dtype) + def _zeros_reward(self) -> np.ndarray: + return np.zeros((self._num_envs,), dtype=np.float32) + + def _zeros_terminated(self) -> np.ndarray: + return np.ones((self._num_envs,), dtype=bool) + + def _zeros_truncated(self) -> np.ndarray: + return np.zeros((self._num_envs,), dtype=bool) + def _zeros_obs(self, space: gym.Space | ObsSpace) -> NpObs: obs_space = _as_obs_space(space) value = None if obs_space.value is None else self._zeros_box(obs_space.value) @@ -227,9 +249,11 @@ def _reset_done_envs(self) -> None: """ state = self._state with self.perf.scope("done_mask"): - done = state.done + # np.asarray is a no-op for numpy state and a zero-copy view for + # torch-backed manager state; host control flow stays numpy. + done = np.asarray(state.done) assert done.shape == (self._num_envs,) - has_done = np.any(done) + has_done = done.any() if not has_done: return @@ -267,7 +291,14 @@ def _update_truncate(self): max_episode_steps = self._max_episode_steps() if not max_episode_steps: return - self._state.truncated = self._state.episode_steps >= max_episode_steps + # Assignment through the slot works for both numpy arrays and torch + # tensors (the latter receive the numpy mask via from_numpy). + truncated = self._state.episode_steps >= max_episode_steps + storage = self._state.truncated + if isinstance(storage, np.ndarray): + storage[:] = truncated + else: + storage[:] = torch.from_numpy(truncated) @abc.abstractmethod def create_renderer(self, config: RenderConfig) -> SimRenderer: @@ -305,9 +336,9 @@ def _prev_physics_step(self): if self._state.final_obs is not None: self._state = self._state.replace(final_obs=None) state = self._state - state.reward.fill(0.0) - state.terminated.fill(False) - state.truncated.fill(False) + state.reward[:] = 0.0 + state.terminated[:] = False + state.truncated[:] = False def compute_observation(self, state: ArrayEnvState) -> ArrayEnvState: """Compute observations for all environments after reset. @@ -345,7 +376,7 @@ def step(self, actions: np.ndarray) -> ArrayEnvState: self._state = self.compute_transition(self._state) self._state.episode_steps += 1 self._update_truncate() - trunc_ids = np.flatnonzero(self._state.truncated) + trunc_ids = np.flatnonzero(np.asarray(self._state.truncated)) if trunc_ids.size: with self.perf.scope("final_obs"): self._state = self._state.replace(final_obs=self._capture_final_obs(trunc_ids)) diff --git a/motrix_env_core/src/motrix_env_core/manager/env.py b/motrix_env_core/src/motrix_env_core/manager/env.py index 70ea2f5f..532f18e0 100644 --- a/motrix_env_core/src/motrix_env_core/manager/env.py +++ b/motrix_env_core/src/motrix_env_core/manager/env.py @@ -11,6 +11,7 @@ import gymnasium as gym import numpy as np +import torch from motrix_env_core.array.env import ArrayEnv, ArrayEnvState, EnvCfgType, FinalObs, NpObs from motrix_env_core.base import EnvCfg, ObsSpace @@ -205,10 +206,10 @@ class _ObsBufferSlot: kernel_outputs: KernelOutputs @classmethod - def bind(cls, observation: NpObs, reward: np.ndarray, terminated: np.ndarray) -> _ObsBufferSlot: + def bind(cls, observation: NpObs, reward: torch.Tensor, terminated: torch.Tensor) -> _ObsBufferSlot: value = observation.value if value is None: - value = np.empty((observation.policy.shape[0], 0), dtype=observation.policy.dtype) + value = torch.empty((observation.policy.shape[0], 0), dtype=observation.policy.dtype) return cls( observation=observation, kernel_outputs=KernelOutputs( @@ -225,11 +226,11 @@ class SwapObsBuffer: """Alternate two preallocated observation slots without copying. Each slot binds an exposed ``NpObs`` to ``KernelOutputs`` that write - into the same policy/value arrays. Swapping both together before the next + into the same policy/value tensors. Swapping both together before the next kernel call keeps ``obs_t`` unchanged while a transition consumer retains it and the kernel produces ``obs_t+1`` in the other slot. - Reward and termination arrays are shared because only observations need + Reward and termination tensors are shared because only observations need the extra step of lifetime. This preserves that lifetime without per-step observation allocations or copies. """ @@ -242,8 +243,8 @@ def create( cls, first: NpObs, second: NpObs, - reward: np.ndarray, - terminated: np.ndarray, + reward: torch.Tensor, + terminated: torch.Tensor, ) -> SwapObsBuffer: return cls((_ObsBufferSlot.bind(first, reward, terminated), _ObsBufferSlot.bind(second, reward, terminated))) @@ -410,6 +411,21 @@ def num_actuators(self) -> int: def physics_step(self) -> None: self.sim.step(self._cfg.sim_substeps) + # The manager pipeline's transition storage is torch-native: the compiled + # kernels write tensors directly and torch consumers pass them through + # without copies. Non-torch consumers take zero-copy ``.numpy()`` views. + def _zeros_box(self, space: gym.spaces.Box) -> torch.Tensor: + return torch.zeros((self.num_envs, *space.shape), dtype=torch.float32) + + def _zeros_reward(self) -> torch.Tensor: + return torch.zeros((self.num_envs,), dtype=torch.float32) + + def _zeros_terminated(self) -> torch.Tensor: + return torch.ones((self.num_envs,), dtype=torch.bool) + + def _zeros_truncated(self) -> torch.Tensor: + return torch.zeros((self.num_envs,), dtype=torch.bool) + def init_state(self) -> ArrayEnvState: startup_started = perf_counter() env_name = type(self).__name__ @@ -623,7 +639,7 @@ def _reset_done_envs(self) -> None: """ state = self._state with self.perf.scope("done_mask"): - done = state.done + done = np.asarray(state.done) sim_reset_ids = np.flatnonzero(self._sim_reset_requested[:, 0]) if sim_reset_ids.size: sim_reset_ids = sim_reset_ids[~done[sim_reset_ids]] @@ -650,11 +666,13 @@ def reset(self, env_ids: np.ndarray, sim_reset_ids: np.ndarray | None = None) -> assert self._state is not None if env_ids.size: with self.perf.scope("command_reset"): + # Host command hooks keep the numpy contract: hand them the + # zero-copy numpy view of the torch termination state. for command_term in self._command_terms.values(): command_term.reset( ResetContext( env_ids=env_ids, - terminated=self._state.terminated, + terminated=self._state.terminated.numpy(), metrics=self._state.metrics, ) ) @@ -669,19 +687,18 @@ def reset(self, env_ids: np.ndarray, sim_reset_ids: np.ndarray | None = None) -> def _make_metrics_view(self) -> dict[str, Any]: """Assemble the persistent live metrics view for the current state. - Values are views into the kernel reward/termination buffers and the - term-owned per-env arrays, so kernels write them in place every step - and nothing needs rebuilding per transition. Host-side writers may add - batch-level scalar entries (e.g. sampler statistics at reset time). - Consumers retain values across steps through + Values are numpy views into the kernel reward/termination buffers and + the term-owned per-env arrays, so kernels write them in place every + step and nothing needs rebuilding per transition. Host-side writers + may add batch-level scalar entries (e.g. sampler statistics at reset + time). Consumers retain values across steps through ``ArrayEnvState.process_metrics()``. """ assert self._compiled_manager_program is not None assert self._kernel_buffers is not None termination_masks = self._kernel_buffers[2] - metrics = { - term.name: termination_masks[:, term.index] for term in self._compiled_manager_program.layout.terminations - } + masks_view = termination_masks.numpy() + metrics = {term.name: masks_view[:, term.index] for term in self._compiled_manager_program.layout.terminations} metrics.update( { name: value[:, 0] if value.ndim == 2 and value.shape[1] == 1 else value @@ -737,12 +754,12 @@ def observation_groups(self) -> dict[str, ObservationGroupEntry]: """Return the environment-local runtime observation groups.""" return self._observation_groups - def _make_kernel_buffers(self, state: ArrayEnvState) -> tuple[np.ndarray, ...]: + def _make_kernel_buffers(self, state: ArrayEnvState) -> tuple[torch.Tensor, ...]: assert self._compiled_manager_program is not None layout = self._compiled_manager_program.layout - reward_terms = np.empty((self.num_envs, len(layout.rewards)), dtype=np.float32) - weighted_reward_terms = np.empty_like(reward_terms) - termination_masks = np.empty((self.num_envs, len(layout.terminations)), dtype=bool) + reward_terms = torch.empty((self.num_envs, len(layout.rewards)), dtype=torch.float32) + weighted_reward_terms = torch.empty_like(reward_terms) + termination_masks = torch.empty((self.num_envs, len(layout.terminations)), dtype=torch.bool) state.reward_terms = {term.name: weighted_reward_terms[:, term.index] for term in layout.rewards} buffers = ( reward_terms, diff --git a/motrix_env_core/src/motrix_env_core/manager/kernel.py b/motrix_env_core/src/motrix_env_core/manager/kernel.py index 15df1362..81fa8eec 100644 --- a/motrix_env_core/src/motrix_env_core/manager/kernel.py +++ b/motrix_env_core/src/motrix_env_core/manager/kernel.py @@ -13,12 +13,12 @@ class KernelOutputs(NamedTuple): - """Preallocated arrays written by one manager task kernel.""" + """Preallocated tensors written by one manager task kernel.""" - policy_obs: np.ndarray - value_obs: np.ndarray - reward: np.ndarray - terminated: np.ndarray + policy_obs: torch.Tensor + value_obs: torch.Tensor + reward: torch.Tensor + terminated: torch.Tensor @dataclass(frozen=True) @@ -36,7 +36,7 @@ class ManagerWarmupResult: def validate_kernel_context( trees: Any, - reward_weights: np.ndarray, + reward_weights: torch.Tensor, buffers: Any, outputs: KernelOutputs, num_envs: int, @@ -140,9 +140,9 @@ def _validate_value(value: Any, path: str, *, num_envs: int, batched: bool, writ def _validate_outputs(outputs: KernelOutputs, num_envs: int) -> None: _validate_value(outputs, "outputs", num_envs=num_envs, batched=True, writable=True) - if outputs.reward.shape != (num_envs,): - raise ValueError(f"outputs.reward must have shape ({num_envs},), got {outputs.reward.shape}.") - if outputs.terminated.shape != (num_envs,): - raise ValueError(f"outputs.terminated must have shape ({num_envs},), got {outputs.terminated.shape}.") - if outputs.terminated.dtype != np.bool_: + if tuple(outputs.reward.shape) != (num_envs,): + raise ValueError(f"outputs.reward must have shape ({num_envs},), got {tuple(outputs.reward.shape)}.") + if tuple(outputs.terminated.shape) != (num_envs,): + raise ValueError(f"outputs.terminated must have shape ({num_envs},), got {tuple(outputs.terminated.shape)}.") + if outputs.terminated.dtype is not torch.bool: raise TypeError(f"outputs.terminated must use bool dtype, got {outputs.terminated.dtype}.") diff --git a/motrix_env_core/src/motrix_env_core/torch_manager/runtime.py b/motrix_env_core/src/motrix_env_core/torch_manager/runtime.py index a9ddf5ae..894c6058 100644 --- a/motrix_env_core/src/motrix_env_core/torch_manager/runtime.py +++ b/motrix_env_core/src/motrix_env_core/torch_manager/runtime.py @@ -9,9 +9,10 @@ hooks, rewards, terminations, command ``advance`` hooks for evaluate; per-group observation terms for observe; command ``reset_env`` hooks plus simulator reset terms for reset — with one batched tensor call per term -instead of a per-lane loop. All writes go back into the persistent numpy -buffers the host lifecycle already owns, so ``ManagerEnv`` keeps its current -API and state surfaces unchanged. +instead of a per-lane loop. All writes go back into the persistent tensor +buffers the host lifecycle already owns (reset compact-write buffers stay +numpy, matching ``SimResetRuntime.buffers``), so ``ManagerEnv`` keeps its +current API and state surfaces unchanged. By default each kernel schedule is additionally fused with ``torch.compile`` into a compiled closure over the captured persistent tensors (one fused @@ -83,10 +84,10 @@ def _compile_options() -> dict[str, Any]: class TorchTaskProgram: """Drop-in replacement task program backed by Torch kernels.""" - evaluate_kernel: Callable[[tuple[Any, ...], np.ndarray, tuple[np.ndarray, ...], Any], None] + evaluate_kernel: Callable[[tuple[Any, ...], torch.Tensor, tuple[torch.Tensor, ...], Any], None] observe_kernel: Callable[[tuple[Any, ...], Any], None] reset_kernel: Callable[[tuple[Any, ...], np.ndarray, tuple[np.ndarray, ...]], None] - reward_weights: np.ndarray + reward_weights: torch.Tensor class _TensorCache: @@ -447,8 +448,8 @@ def __init__(self, env: ManagerEnv, program: ResolvedManagerProgram, compile: bo self.command_reset_hooks: tuple[CommandResetHook, ...] = program.command_reset_hooks self.observations: dict[str, tuple[ObservationTermCall, ...]] = program.observations reward_cfgs = env.cfg.reward_cfgs() - self.reward_weights = np.asarray( - [float(reward_cfgs[term.name].weight) for term in self.layout.rewards], dtype=np.float32 + self.reward_weights = torch.tensor( + [float(reward_cfgs[term.name].weight) for term in self.layout.rewards], dtype=torch.float32 ) # Reset-term compact-write layout mirrors SimResetRuntime.buffers: # one entry per declared write, ordered by term then output name. @@ -482,15 +483,15 @@ def task_program(self) -> TorchTaskProgram: def evaluate( self, inputs: tuple[Any, ...], - reward_weights: np.ndarray, - buffers: tuple[np.ndarray, np.ndarray, np.ndarray], + reward_weights: torch.Tensor, + buffers: tuple[torch.Tensor, torch.Tensor, torch.Tensor], outputs: KernelOutputs, ) -> None: kernel = self._compiled_kernels.get(("evaluate", id(inputs), id(buffers))) if kernel is None: kernel = self._build_compiled_evaluate(inputs, reward_weights, buffers) self._compiled_kernels[("evaluate", id(inputs), id(buffers))] = kernel - kernel(torch.from_numpy(outputs.reward), torch.from_numpy(outputs.terminated)) + kernel(outputs.reward, outputs.terminated) def observe(self, inputs: tuple[Any, ...], outputs: KernelOutputs) -> None: kernel = self._compiled_kernels.get(("observe", id(inputs))) @@ -499,8 +500,8 @@ def observe(self, inputs: tuple[Any, ...], outputs: KernelOutputs) -> None: self._compiled_kernels[("observe", id(inputs))] = kernel value_obs = outputs.value_obs kernel( - torch.from_numpy(outputs.policy_obs), - torch.from_numpy(value_obs) if value_obs is not None and value_obs.size else None, + outputs.policy_obs, + value_obs if value_obs is not None and value_obs.numel() else None, ) def reset(self, inputs: tuple[Any, ...], env_ids: np.ndarray, reset_buffers: tuple[np.ndarray, ...]) -> None: @@ -516,14 +517,14 @@ def reset(self, inputs: tuple[Any, ...], env_ids: np.ndarray, reset_buffers: tup def _build_compiled_evaluate( self, inputs: tuple[Any, ...], - reward_weights: np.ndarray, - buffers: tuple[np.ndarray, np.ndarray, np.ndarray], + reward_weights: torch.Tensor, + buffers: tuple[torch.Tensor, torch.Tensor, torch.Tensor], ) -> Callable[[torch.Tensor, torch.Tensor], None]: ctx = self._context_for(inputs) - weights = torch.from_numpy(np.asarray(reward_weights, dtype=np.float32)) - raw_terms = torch.from_numpy(buffers[0]) - weighted_terms = torch.from_numpy(buffers[1]) - termination_masks = torch.from_numpy(buffers[2]) + weights = reward_weights + raw_terms = buffers[0] + weighted_terms = buffers[1] + termination_masks = buffers[2] dt = ctx.dt update_hooks = self.command_updates advance_hooks = self.command_advances diff --git a/motrix_env_core/tests/test_manager_env.py b/motrix_env_core/tests/test_manager_env.py index c089a11c..bfe710b3 100644 --- a/motrix_env_core/tests/test_manager_env.py +++ b/motrix_env_core/tests/test_manager_env.py @@ -442,7 +442,7 @@ class _MultipleActionsCfg(ManagerActionsCfg): ("second", slice(1, 3)), ] - state.terminated[:] = [False, True] + state.terminated[:] = torch.tensor([False, True]) env._reset_done_envs() np.testing.assert_array_equal(env.action_terms["test"].state.reset_flags[:, 0], [False, True]) np.testing.assert_array_equal(env.action_terms["second"].state.reset_flags[:, 0], [False, True]) @@ -503,20 +503,20 @@ def test_step_preserves_previous_observation_with_alternating_preallocated_buffe env = _ManagerEnv() initial = env.init_state() initial_policy = initial.obs.policy - initial_snapshot = initial_policy.copy() + initial_snapshot = initial_policy.clone() first = env.step(np.full((env.num_envs, 1), 0.25, dtype=np.float32)) np.testing.assert_array_equal(initial_policy, initial_snapshot) - assert not np.shares_memory(initial_policy, first.obs.policy) + assert initial_policy.data_ptr() != first.obs.policy.data_ptr() first_policy = first.obs.policy - first_snapshot = first_policy.copy() + first_snapshot = first_policy.clone() second = env.step(np.full((env.num_envs, 1), 0.5, dtype=np.float32)) np.testing.assert_array_equal(first_policy, first_snapshot) - assert not np.shares_memory(first_policy, second.obs.policy) - assert np.shares_memory(initial_policy, second.obs.policy) + assert first_policy.data_ptr() != second.obs.policy.data_ptr() + assert initial_policy.data_ptr() == second.obs.policy.data_ptr() def test_manager_context_is_injected_once_and_reused_across_all_term_kinds() -> None: @@ -606,7 +606,7 @@ def test_command_term_updates_lifecycle_and_selected_reset_ids() -> None: command[:] = torch.tensor([[1.0], [2.0], [3.0]]) action = env.action_terms["test"] - state.terminated[:] = [False, True, False] + state.terminated[:] = torch.tensor([False, True, False]) env._reset_done_envs() np.testing.assert_array_equal(command[:, 0], [1.0, -1.0, 3.0]) @@ -662,8 +662,8 @@ def test_final_observations_capture_command_before_reset_and_survive_steps(with_ else: assert final.obs.value is None assert final.obs.value_or_policy is final.obs.policy - assert not np.shares_memory(final.obs.policy, state.obs.policy) - saved = final.obs.policy.copy() + assert final.obs.policy.data_ptr() != state.obs.policy.data_ptr() + saved = final.obs.policy.clone() next_state = env.step(np.full((3, 1), 0.25, dtype=np.float32)) assert next_state.final_obs is None assert state.final_obs is final @@ -687,7 +687,7 @@ def noisy_observe(inputs, outputs): np.testing.assert_array_equal( env._rand.state.numpy().view(np.uint64), rng_before.numpy().view(np.uint64) + np.uint64(1) ) - saved_final = final.obs.policy.copy() + saved_final = final.obs.policy.clone() np.testing.assert_array_equal( final.obs.policy[:, 0], (rng_before[[0], 0] + 1).to(torch.float32).numpy(), @@ -704,7 +704,7 @@ def test_post_reset_observation_does_not_run_command_evaluation() -> None: action = env.action_terms["test"] assert isinstance(action, ActionTerm) action.state.source[:, 0] = torch.tensor([0.25, 0.5, 0.75]) - state.terminated[:] = [False, True, False] + state.terminated[:] = torch.tensor([False, True, False]) env._refresh_sim_reads() env._reset_done_envs() env.compute_observation(state) diff --git a/motrix_env_core/tests/test_manager_sim_backend.py b/motrix_env_core/tests/test_manager_sim_backend.py index a3e79133..a0c03a78 100644 --- a/motrix_env_core/tests/test_manager_sim_backend.py +++ b/motrix_env_core/tests/test_manager_sim_backend.py @@ -336,7 +336,7 @@ def test_partial_reset_hands_canonical_rows_to_runtime_and_clears_ctrl() -> None env.step(np.full((3, 1), 0.5, dtype=np.float32)) assert runtime.ctrl_targets[1, 1] == np.float32(1.0) - env._state.terminated[:] = [False, True, False] + env._state.terminated[:] = torch.tensor([False, True, False]) env._reset_done_envs() assert len(runtime.reset_calls) == 1 diff --git a/motrix_rl_builtin/src/motrix_rl_builtin/fastsac/wrap_np.py b/motrix_rl_builtin/src/motrix_rl_builtin/fastsac/wrap_np.py index 20882d0b..36366973 100644 --- a/motrix_rl_builtin/src/motrix_rl_builtin/fastsac/wrap_np.py +++ b/motrix_rl_builtin/src/motrix_rl_builtin/fastsac/wrap_np.py @@ -1,10 +1,12 @@ # Copyright Motphys Technology Co., Ltd. 2025, 2026 # SPDX-License-Identifier: Apache-2.0 -"""Numpy-backend env wrapper for FastSAC. +"""NumPy-frontend env wrapper for FastSAC. Wraps a Motrix ``ArrayEnv`` as the torch-tensor vectorized environment consumed -by FastSAC. +by FastSAC. Manager environments already expose torch state; the tensors pass +through zero-copy (only a device move when training off-CPU). NumPy-frontend +(Direct) environments convert at this boundary. The actor maps its tanh output to the action range declared by the environment; the wrapper only clips to that standard contract before stepping. Done @@ -29,7 +31,9 @@ def __init__(self, env: ArrayEnv, device: torch.device, render: RenderConfig | N self._env: ArrayEnv = env self._renderer = create_renderer(env, render) - def _to_torch(self, arr: np.ndarray, dtype=torch.float32) -> torch.Tensor: + def _to_torch(self, arr, dtype=torch.float32) -> torch.Tensor: + if isinstance(arr, torch.Tensor): + return arr.to(device=self.device, dtype=dtype) return torch.as_tensor(np.asarray(arr), dtype=dtype, device=self.device) def reset(self) -> tuple[torch.Tensor, torch.Tensor]: @@ -37,8 +41,9 @@ def reset(self) -> tuple[torch.Tensor, torch.Tensor]: return self._to_torch(state.obs.policy), self._to_torch(state.obs.value_or_policy) def step(self, actions: torch.Tensor) -> FastSacStepResult: - np_actions = actions.detach().cpu().numpy().astype(np.float32) - np_actions = np.clip(np_actions, self._low, self._high) + # Zero-copy CPU view of the learner's actions; clipping allocates the + # buffer handed to the numpy-frontend action path. + np_actions = np.clip(actions.detach().cpu().numpy(), self._low, self._high) state = self._env.step(np_actions) return FastSacStepResult( obs=self._to_torch(state.obs.policy), diff --git a/motrix_rl_rslrl/src/motrix_rl_rslrl/torch/wrap_np.py b/motrix_rl_rslrl/src/motrix_rl_rslrl/torch/wrap_np.py index 4283a933..4973d757 100644 --- a/motrix_rl_rslrl/src/motrix_rl_rslrl/torch/wrap_np.py +++ b/motrix_rl_rslrl/src/motrix_rl_rslrl/torch/wrap_np.py @@ -1,27 +1,34 @@ # Copyright Motphys Technology Co., Ltd. 2025, 2026 # SPDX-License-Identifier: Apache-2.0 -"""VecEnv wrapper for adapting DirectEnv to RSLRL's VecEnv interface.""" +"""VecEnv wrapper for adapting ArrayEnv environments to RSLRL's VecEnv interface.""" import numpy as np import torch from rsl_rl.env.vec_env import VecEnv from tensordict import TensorDict -from motrix_env_core.array.env import NpObs -from motrix_env_core.direct.env import DirectEnv +from motrix_env_core.array.env import ArrayEnv, NpObs from motrix_env_core.renderer import RenderConfig, create_renderer +def _as_torch(value, device: torch.device) -> torch.Tensor: + """Zero-copy tensor view for both numpy and torch env state.""" + if isinstance(value, torch.Tensor): + return value.to(device) + return torch.from_numpy(np.asarray(value)).to(device) + + class RslrlNpEnvWrap(VecEnv): - """Adapter class that wraps DirectEnv to RSLRL's VecEnv interface. + """Adapter class that wraps ArrayEnv environments to RSLRL's VecEnv interface. RSLRL expects a VecEnv interface with specific methods for stepping, - resetting, and accessing observations. This adapter converts between - DirectEnv's ArrayEnvState format and RSLRL's expected format. + resetting, and accessing observations. This adapter converts between the + ``ArrayEnvState`` format and RSLRL's expected format; manager environments + already expose torch state and pass through zero-copy. """ - def __init__(self, env: DirectEnv, device: torch.device, render: RenderConfig | None = None): + def __init__(self, env: ArrayEnv, device: torch.device, render: RenderConfig | None = None): """Initialize the VecEnv adapter. Args: @@ -91,19 +98,19 @@ def step(self, actions: torch.Tensor) -> tuple[TensorDict, torch.Tensor, torch.T # Update episode length buffer self.episode_length_buf += 1 # Reset episode length for done environments - dones_np = state.done.astype(bool) - self.episode_length_buf[dones_np] = 0 + done_np = np.asarray(state.done, dtype=bool) + self.episode_length_buf[done_np] = 0 - rewards = torch.from_numpy(state.reward).to(self._device) + rewards = _as_torch(state.reward, self._device) # Merge terminated and truncated into dones - dones = torch.from_numpy(state.done.astype(np.float32)).to(self._device) + dones = _as_torch(np.asarray(state.done, dtype=np.float32), self._device) obs = self._to_tensordict(state.obs) # Build extras dict (RSLRL calls it "extras" not "infos") # time_outs: rows truncated without failing, for value bootstrapping. - extras = {"time_outs": torch.from_numpy(state.truncated & ~state.terminated).to(self._device)} + extras = {"time_outs": _as_torch(np.asarray(state.truncated & ~state.terminated, dtype=bool), self._device)} return obs, rewards, dones, extras @@ -141,9 +148,9 @@ def get_observations(self) -> TensorDict: return self._to_tensordict(self._state.obs) def _to_tensordict(self, obs: NpObs) -> TensorDict: - tensors = {"policy": torch.from_numpy(obs.policy).to(self._device)} + tensors = {"policy": _as_torch(obs.policy, self._device)} if obs.value is not None: - tensors["value"] = torch.from_numpy(obs.value).to(self._device) + tensors["value"] = _as_torch(obs.value, self._device) return TensorDict(tensors, batch_size=[self._num_envs], device=self._device) def render(self, *args, **kwargs) -> bool | None: diff --git a/motrix_rl_skrl/src/motrix_rl_skrl/jax/train/ppo.py b/motrix_rl_skrl/src/motrix_rl_skrl/jax/train/ppo.py index 8476caca..671fd9aa 100644 --- a/motrix_rl_skrl/src/motrix_rl_skrl/jax/train/ppo.py +++ b/motrix_rl_skrl/src/motrix_rl_skrl/jax/train/ppo.py @@ -84,6 +84,9 @@ def record_transition( if "Reward" in infos: for key, value in infos["Reward"].items(): + # Manager environments hand out torch tensors; take the + # zero-copy numpy view for the JAX statistics path. + value = np.asarray(value) self.tracking_data[f"Reward Instant / {key} (max)"].append(jnp.max(value)) self.tracking_data[f"Reward Instant / {key} (min)"].append(jnp.min(value)) self.tracking_data[f"Reward Instant / {key} (mean)"].append(jnp.mean(value)) diff --git a/motrix_rl_skrl/src/motrix_rl_skrl/jax/wrap_np.py b/motrix_rl_skrl/src/motrix_rl_skrl/jax/wrap_np.py index 8dc34104..e00e9d0e 100644 --- a/motrix_rl_skrl/src/motrix_rl_skrl/jax/wrap_np.py +++ b/motrix_rl_skrl/src/motrix_rl_skrl/jax/wrap_np.py @@ -9,27 +9,35 @@ import numpy as np from skrl.envs.wrappers.jax import Wrapper as SkrlWrapper -from motrix_env_core.direct.env import DirectEnv +from motrix_env_core.array.env import ArrayEnv from motrix_env_core.renderer import RenderConfig, VideoRecorder, create_renderer from motrix_env_core.sim.backend import SimRenderer from motrix_rl.utils import env_infos +def _as_jnp(value, dtype) -> jax.Array: + """Zero-copy JAX view over both numpy and torch env state.""" + return jnp.asarray(np.asarray(value), dtype=dtype) + + class SkrlNpWrapper(SkrlWrapper): """ Wrap the numpy-based environment to be compatible with skrl + + Manager environments expose torch state; ``np.asarray`` yields the + zero-copy numpy views the JAX path consumes. """ - _env: DirectEnv + _env: ArrayEnv _renderer: SimRenderer | VideoRecorder | None = None - def __init__(self, env: DirectEnv, render: RenderConfig | None = None): + def __init__(self, env: ArrayEnv, render: RenderConfig | None = None): super().__init__(env) self._renderer = create_renderer(env, render) def reset(self) -> tuple[jax.Array, Any]: state = self._env.init_state() - return jnp.asarray(state.obs.policy, dtype=jnp.float32), env_infos(state) + return _as_jnp(state.obs.policy, jnp.float32), env_infos(state) def step( self, actions: jax.Array @@ -40,18 +48,18 @@ def step( jax.Array, Any, ]: - actions = np.array(actions) + actions = np.asarray(actions) state = self._env.step(actions) return ( - jnp.asarray(state.obs.policy, dtype=jnp.float32), - jnp.asarray(state.reward.reshape(-1, 1), dtype=jnp.float32), - jnp.asarray(state.terminated.reshape(-1, 1), dtype=jnp.bool_), - jnp.asarray(state.truncated.reshape(-1, 1), dtype=jnp.bool_), + _as_jnp(state.obs.policy, jnp.float32), + _as_jnp(state.reward.reshape(-1, 1), jnp.float32), + _as_jnp(state.terminated.reshape(-1, 1), jnp.bool_), + _as_jnp(state.truncated.reshape(-1, 1), jnp.bool_), env_infos(state), ) def state(self) -> jax.Array: - return jnp.asarray(self._env.state.obs.value_or_policy, dtype=jnp.float32) + return _as_jnp(self._env.state.obs.value_or_policy, jnp.float32) def render(self, *args, **kwargs) -> bool | None: if self._renderer is None: diff --git a/motrix_rl_skrl/src/motrix_rl_skrl/torch/wrap_np.py b/motrix_rl_skrl/src/motrix_rl_skrl/torch/wrap_np.py index 4881bcce..3a92f142 100644 --- a/motrix_rl_skrl/src/motrix_rl_skrl/torch/wrap_np.py +++ b/motrix_rl_skrl/src/motrix_rl_skrl/torch/wrap_np.py @@ -4,25 +4,36 @@ from typing import Any import gymnasium +import numpy as np import torch from skrl.envs.wrappers.torch import Wrapper as SkrlWrapper -from motrix_env_core.direct.env import DirectEnv +from motrix_env_core.array.env import ArrayEnv from motrix_env_core.renderer import RenderConfig, VideoRecorder, create_renderer from motrix_env_core.sim.backend import SimRenderer from motrix_rl.device import resolve_default_device from motrix_rl.utils import env_infos +def _as_torch(value, dtype: torch.dtype, device: torch.device) -> torch.Tensor: + """Zero-copy tensor view for both numpy and torch env state.""" + if isinstance(value, torch.Tensor): + return value.to(dtype=dtype, device=device) + return torch.as_tensor(np.asarray(value), dtype=dtype, device=device) + + class SkrlNpWrapper(SkrlWrapper): """ Wrap the numpy-based environment to be compatible with skrl (PyTorch) + + Manager environments already expose torch state and pass through + zero-copy; NumPy-frontend (Direct) environments convert here. """ - _env: DirectEnv + _env: ArrayEnv _renderer: SimRenderer | VideoRecorder | None = None - def __init__(self, env: DirectEnv, render: RenderConfig | None = None): + def __init__(self, env: ArrayEnv, render: RenderConfig | None = None): super().__init__(env) # skrl's fallback for device-less envs is cuda-else-cpu; resolve with # the shared cuda -> mps -> cpu policy so Apple Silicon trains on MPS. @@ -35,7 +46,7 @@ def __init__(self, env: DirectEnv, render: RenderConfig | None = None): def reset(self) -> tuple[torch.Tensor, Any]: state = self._env.init_state() - return torch.tensor(state.obs.policy, dtype=torch.float32, device=self.device), env_infos(state) + return _as_torch(state.obs.policy, torch.float32, self.device), env_infos(state) def step( self, actions: torch.Tensor @@ -49,15 +60,15 @@ def step( actions = actions.cpu().numpy() state = self._env.step(actions) return ( - torch.tensor(state.obs.policy, dtype=torch.float32, device=self.device), - torch.tensor(state.reward.reshape(-1, 1), dtype=torch.float32, device=self.device), - torch.tensor(state.terminated.reshape(-1, 1), dtype=torch.bool, device=self.device), - torch.tensor(state.truncated.reshape(-1, 1), dtype=torch.bool, device=self.device), + _as_torch(state.obs.policy, torch.float32, self.device), + _as_torch(state.reward.reshape(-1, 1), torch.float32, self.device), + _as_torch(state.terminated.reshape(-1, 1), torch.bool, self.device), + _as_torch(state.truncated.reshape(-1, 1), torch.bool, self.device), env_infos(state), ) def state(self) -> torch.Tensor: - return torch.tensor(self._env.state.obs.value_or_policy, dtype=torch.float32, device=self.device) + return _as_torch(self._env.state.obs.value_or_policy, torch.float32, self.device) def render(self, *args, **kwargs) -> bool | None: if self._renderer is None: From 82b9d72fa9a4e7c336b25216f5a9a1d9e7b21a4f Mon Sep 17 00:00:00 2001 From: zilch <147668916@qq.com> Date: Fri, 6 Nov 2026 15:26:42 +0800 Subject: [PATCH 27/41] refactor(manager): resolve sim leaves as zero-copy tensor views 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. --- .../src/motrix_env_core/manager/env.py | 8 ++++---- .../src/motrix_env_core/manager/kernel.py | 2 +- .../manager/kernel_data/resolved.py | 6 ++++-- .../src/motrix_env_core/manager/resolver.py | 17 ++++++++++++++++- 4 files changed, 25 insertions(+), 8 deletions(-) diff --git a/motrix_env_core/src/motrix_env_core/manager/env.py b/motrix_env_core/src/motrix_env_core/manager/env.py index 532f18e0..35c62548 100644 --- a/motrix_env_core/src/motrix_env_core/manager/env.py +++ b/motrix_env_core/src/motrix_env_core/manager/env.py @@ -182,10 +182,10 @@ def termination_cfgs(self) -> dict[str, TerminationTermCfg]: class ReadPlan: """Simulator read program bound to the resolved context tree. - ``read`` executes the compiled read program; the simulator-backed ndarray - leaves of the resolved context tree are rewritten in place, so every other - resolved leaf and the tree objects themselves stay valid without - rebuilding. + ``read`` executes the compiled read program; the simulator arena the + context tree's sim leaves view (as zero-copy tensors) is rewritten in + place, so every resolved leaf and the tree objects themselves stay valid + without rebuilding. """ sim_data: PhysicsReadProgram diff --git a/motrix_env_core/src/motrix_env_core/manager/kernel.py b/motrix_env_core/src/motrix_env_core/manager/kernel.py index 81fa8eec..cb75e8a6 100644 --- a/motrix_env_core/src/motrix_env_core/manager/kernel.py +++ b/motrix_env_core/src/motrix_env_core/manager/kernel.py @@ -134,7 +134,7 @@ def _validate_value(value: Any, path: str, *, num_envs: int, batched: bool, writ return raise TypeError( - f"{path} must contain only ndarrays, numeric scalars, and fixed tuples; got {type(value).__name__}." + f"{path} must contain only tensors, ndarrays, numeric scalars, and fixed tuples; got {type(value).__name__}." ) diff --git a/motrix_env_core/src/motrix_env_core/manager/kernel_data/resolved.py b/motrix_env_core/src/motrix_env_core/manager/kernel_data/resolved.py index d59f1b75..72ff200a 100644 --- a/motrix_env_core/src/motrix_env_core/manager/kernel_data/resolved.py +++ b/motrix_env_core/src/motrix_env_core/manager/kernel_data/resolved.py @@ -15,8 +15,10 @@ class ResolvedLeaf: """One resolved leaf of a kernel-data tree. Attributes: - value: Live leaf value reference — an ndarray held and refreshed in - place by its owning runtime object, or an immutable scalar. + value: Live leaf value reference — a tensor or ndarray held and + refreshed in place by its owning runtime object, or an immutable + scalar. Simulator-backed leaves are zero-copy tensor views over + the read-program arena; metrics leaves stay ndarray. """ value: Any diff --git a/motrix_env_core/src/motrix_env_core/manager/resolver.py b/motrix_env_core/src/motrix_env_core/manager/resolver.py index 2aee0ea9..2c67991b 100644 --- a/motrix_env_core/src/motrix_env_core/manager/resolver.py +++ b/motrix_env_core/src/motrix_env_core/manager/resolver.py @@ -18,6 +18,7 @@ import hashlib import inspect import logging +import warnings from collections.abc import Callable, Mapping from dataclasses import dataclass from time import perf_counter @@ -46,6 +47,15 @@ logger = logging.getLogger(__name__) +# Simulator-backed read-only arena views are wrapped as tensors at resolution +# time; suppress PyTorch's non-writable-array warning once here instead of +# copying every view (mirrors the filter in the Torch executor). +warnings.filterwarnings( + "ignore", + message="The given NumPy array is not writable", + category=UserWarning, +) + # --------------------------------------------------------------------------- # Resolved layouts @@ -513,7 +523,12 @@ def _resolve_manager_context( ) -> ResolvedManagerContext: del observation_groups, reward_terms, termination_terms actions = {name: term.state for name, term in self._env.action_terms.items()} - sim = {key: binding.value for key, binding in self._sim_inputs.items()} + # Sim views enter the resolved tree as zero-copy tensor views over the + # read-program arena: the read program rewrites the arena contents in + # place, so the pre-built views stay valid for the program's lifetime + # and every upper layer sees one leaf type. Metrics stay ndarray per + # the existing numpy metrics contract. + sim = {key: torch.from_numpy(binding.value) for key, binding in self._sim_inputs.items()} commands = dict(self._env.command_terms) value = ManagerContext( env_id=torch.from_numpy(np.arange(self._env.num_envs, dtype=np.int64)), From 1b48ffc26b67261f26d644eaf4ac1e7f85032308 Mon Sep 17 00:00:00 2001 From: zilch <147668916@qq.com> Date: Fri, 6 Nov 2026 16:12:24 +0800 Subject: [PATCH 28/41] refactor(manager)!: pass carrier term arguments directly and drop TreeArg MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. --- .../src/motrix_env_core/manager/resolver.py | 68 ++++++++++++++----- .../src/motrix_env_core/mdp/terrain.py | 11 ++- .../motrix_env_core/torch_manager/runtime.py | 9 +-- motrix_env_core/tests/test_manager_env.py | 4 +- .../tests/test_term_arg_fingerprint.py | 56 +++++++++++++++ .../locomotion/ball_balance/mdp/reset.py | 4 +- .../locomotion/ball_balance/mdp/rewards.py | 4 +- .../humanoid/walk_manager_mdp/reset.py | 16 ++--- .../humanoid/walk_manager_mdp/terrain.py | 4 +- .../locomotion/wbt/mdp/observations.py | 4 +- .../motrix_envs/locomotion/wbt/mdp/rewards.py | 8 +-- wiki/design/manager/kernel-abi.md | 11 ++- wiki/design/manager/torch-executor.md | 9 ++- 13 files changed, 149 insertions(+), 59 deletions(-) create mode 100644 motrix_env_core/tests/test_term_arg_fingerprint.py diff --git a/motrix_env_core/src/motrix_env_core/manager/resolver.py b/motrix_env_core/src/motrix_env_core/manager/resolver.py index 2c67991b..89ef7faa 100644 --- a/motrix_env_core/src/motrix_env_core/manager/resolver.py +++ b/motrix_env_core/src/motrix_env_core/manager/resolver.py @@ -4,7 +4,7 @@ """Backend-neutral resolver for manager task programs. Lowers one constructed manager environment into its execution metadata: the -simulator read plan, the resolved context/term-argument trees (path-keyed +simulator read plan, the resolved context tree (path-keyed leaf descriptors with live value references), prepared term arguments, and one typed call record per manager term (command hooks, observations, rewards, terminations, simulator resets). Numerical executors — @@ -20,7 +20,7 @@ import logging import warnings from collections.abc import Callable, Mapping -from dataclasses import dataclass +from dataclasses import dataclass, fields from time import perf_counter from types import MappingProxyType from typing import Any, get_args, get_origin, get_type_hints @@ -151,13 +151,6 @@ class ConstantArg: value: Any -@dataclass(frozen=True) -class TreeArg: - """Term argument resolved as a structured kernel-data tree.""" - - tree_index: int - - @dataclass(frozen=True) class LeafArg: """Term argument resolved as one leaf value (numeric/ndarray argument).""" @@ -172,7 +165,7 @@ class SimArg: key: str -PreparedArg = ConstantArg | TreeArg | LeafArg | SimArg +PreparedArg = ConstantArg | LeafArg | SimArg @dataclass(frozen=True) @@ -297,7 +290,6 @@ class ResolvedManagerProgram: "SimInputLayout", "SimResetTermCall", "TerminationTermCall", - "TreeArg", ] @@ -335,6 +327,43 @@ def _tree_leaf_fingerprint(tree: ResolvedTree) -> tuple[Any, ...]: ) +def _carrier_fingerprint(value: Any) -> Any: + """Structural + content fingerprint of a raw carrier dataclass term argument. + + Recurses over dataclass fields recording field names, concrete types, + tensor/ndarray dtype+shape+content hashes, and Python scalar values. + Carriers previously resolved as kernel-data trees contributed no content + to the program fingerprint at all; this makes every carrier data change + (a grid array value, a scalar) invalidate compiled-program caches. + """ + if isinstance(value, torch.Tensor): + return ( + "tensor", + str(value.dtype), + tuple(value.shape), + hashlib.sha256(value.detach().contiguous().numpy().tobytes()).hexdigest(), + ) + if isinstance(value, np.ndarray): + return ( + "ndarray", + str(value.dtype), + value.shape, + hashlib.sha256(np.ascontiguousarray(value).tobytes()).hexdigest(), + ) + if is_kernel_data(value): + return ( + "dataclass", + type(value).__qualname__, + tuple((field.name, _carrier_fingerprint(getattr(value, field.name))) for field in fields(value)), + ) + if isinstance(value, (bool, int, float, str)): + return (type(value).__name__, value) + raise TypeError( + f"Carrier dataclass term argument holds field of unsupported type {type(value).__name__}; " + "allowed: torch.Tensor, ndarray, Python scalars, str, nested carrier dataclasses." + ) + + class ManagerProgramResolver: """Resolve one fixed manager config into executor-neutral metadata.""" @@ -708,15 +737,18 @@ def _resolve_dispatch_args(self, group: str, name: str, term: Any) -> tuple[Prep def _resolve_term_arg(self, source_name: str, value: Any) -> PreparedArg: """Resolve one manager term argument into a prepared descriptor. - Kernel-data dataclass arguments become structured trees; float-like - scalars and arrays/tensors become shared leaf slots read by reference; simulator queries - resolve to their read-program batched array; every other value is a - literal constant. + Carrier dataclass arguments pass through as constants: Dynamo traces + their attribute access directly, and a structural fingerprint over + their raw fields feeds the program fingerprint for cache + invalidation. Float-like scalars and arrays/tensors become shared + leaf slots read by reference; simulator queries resolve to their + read-program batched array; every other value is a literal constant. """ if is_kernel_data(value): - _, tree_def = flatten_kernel_data(value) - layout = self._kernel_data_lowering.lower(tree_def, context=source_name) - return TreeArg(self._register_tree(source_name, value, layout)) + self._layout_identity.append( + ("term_arg", source_name, type(value).__qualname__, _carrier_fingerprint(value)) + ) + return ConstantArg(value) if isinstance(value, (float, np.floating)): return LeafArg(ResolvedLeaf(np.float32(value))) if isinstance(value, np.ndarray): diff --git a/motrix_env_core/src/motrix_env_core/mdp/terrain.py b/motrix_env_core/src/motrix_env_core/mdp/terrain.py index 178b1481..8531b060 100644 --- a/motrix_env_core/src/motrix_env_core/mdp/terrain.py +++ b/motrix_env_core/src/motrix_env_core/mdp/terrain.py @@ -5,7 +5,6 @@ from dataclasses import dataclass -import numpy as np import torch @@ -20,14 +19,14 @@ class HeightFieldGrid: world ``xy`` to fractional grid coordinates. Attributes: - heights: ``np.ndarray`` of shape ``(nrow, ncol)`` holding the + heights: ``torch.Tensor`` of shape ``(nrow, ncol)`` holding the height-field sample values (row = y 方向, column = x 方向), 单位与引擎导出的 hfield 数据一致。 - origin: 长度为 2 的 ``np.ndarray``,网格首个采样点 ``(x, y)`` 的 + origin: 长度为 2 的 ``torch.Tensor``,网格首个采样点 ``(x, y)`` 的 世界坐标,是 bilinear 插值前的平移原点。 - spacing: 长度为 2 的 ``np.ndarray``,相邻采样点在世界坐标系下的 + spacing: 长度为 2 的 ``torch.Tensor``,相邻采样点在世界坐标系下的 间距 ``(dx, dy)``,用于把世界坐标映射到分数网格坐标。 - z0: 长度为 1 的 ``np.ndarray``,叠加在插值结果之上的基准高度偏移。 + z0: 长度为 1 的 ``torch.Tensor``,叠加在插值结果之上的基准高度偏移。 constant: 平地退化模式下的常数高度;仅在 ``enabled=False`` 时生效, 对应平地 ground geom 的世界 z 坐标。 enabled: 是否启用网格插值;``False`` 时高度查找直接返回 @@ -38,7 +37,7 @@ class HeightFieldGrid: origin: torch.Tensor spacing: torch.Tensor z0: torch.Tensor - constant: np.float32 + constant: float enabled: bool diff --git a/motrix_env_core/src/motrix_env_core/torch_manager/runtime.py b/motrix_env_core/src/motrix_env_core/torch_manager/runtime.py index 894c6058..9008db4e 100644 --- a/motrix_env_core/src/motrix_env_core/torch_manager/runtime.py +++ b/motrix_env_core/src/motrix_env_core/torch_manager/runtime.py @@ -4,7 +4,7 @@ """Generic Torch executor for resolved Manager programs. The runtime consumes the metadata produced by the manager resolver (read -plan, layout, resolved context/term-argument trees, and per-term-kind typed +plan, layout, the resolved context tree, and per-term-kind typed call records) and executes the manager term schedule — command ``update`` hooks, rewards, terminations, command ``advance`` hooks for evaluate; per-group observation terms for observe; command ``reset_env`` hooks plus @@ -45,7 +45,6 @@ SimArg, SimResetTermCall, TerminationTermCall, - TreeArg, ) from motrix_env_core.torch_manager.rand import draw_uniform @@ -745,11 +744,7 @@ def _invocation_args( assert self._tree_tensors is not None args: list[Any] = [] for arg in call.args: - if isinstance(arg, TreeArg): - leaves = self._tree_tensors[arg.tree_index] - host = self.program.prepared_trees[arg.tree_index].value - args.append(_Namespace(leaves, (), host, _leaf_lookup(leaves))) - elif isinstance(arg, LeafArg): + if isinstance(arg, LeafArg): args.append(self._cache.tensor(inputs, arg.leaf.value)) elif isinstance(arg, SimArg): assert self._context_tensors is not None diff --git a/motrix_env_core/tests/test_manager_env.py b/motrix_env_core/tests/test_manager_env.py index bfe710b3..a59399ed 100644 --- a/motrix_env_core/tests/test_manager_env.py +++ b/motrix_env_core/tests/test_manager_env.py @@ -96,7 +96,7 @@ def _second_action_space(_: ManagerEnv, __: np.ndarray | None) -> gym.spaces.Box @dataclass(frozen=True, slots=True) class _ObservationParams: offset: torch.Tensor - scale: np.float32 + scale: float def _injected_observation(ctx: ManagerContext, out: torch.Tensor, params: _ObservationParams) -> None: @@ -114,7 +114,7 @@ def __call__(self, env: ManagerEnv) -> ObsTerm: return ObsTerm( 2, _injected_observation, - _ObservationParams(torch.from_numpy(np.asarray([1.0], dtype=np.float32)), np.float32(self.scale)), + _ObservationParams(torch.from_numpy(np.asarray([1.0], dtype=np.float32)), float(self.scale)), ) diff --git a/motrix_env_core/tests/test_term_arg_fingerprint.py b/motrix_env_core/tests/test_term_arg_fingerprint.py new file mode 100644 index 00000000..9e8d8d28 --- /dev/null +++ b/motrix_env_core/tests/test_term_arg_fingerprint.py @@ -0,0 +1,56 @@ +# Copyright Motphys Technology Co., Ltd. 2025, 2026 +# SPDX-License-Identifier: Apache-2.0 + +"""Carrier dataclass term-argument fingerprint contract. + +Carrier dataclasses passed as manager term arguments resolve as constants; +their structural fingerprint must keep compiled-program caches invalidated by +any content change (tensor bytes or scalar values). +""" + +from dataclasses import dataclass + +import numpy as np +import pytest +import torch + +from motrix_env_core.manager.resolver import _carrier_fingerprint + + +@dataclass(frozen=True, slots=True) +class _Nested: + flag: bool + + +@dataclass(frozen=True, slots=True) +class _Carrier: + grid: torch.Tensor + label: str + offset: float + nested: _Nested + + +def test_fingerprint_changes_on_tensor_content_and_scalar() -> None: + carrier = _Carrier(torch.arange(6, dtype=torch.float32).reshape(2, 3), "a", 1.5, _Nested(True)) + base = _carrier_fingerprint(carrier) + edited_grid = _Carrier(carrier.grid.clone(), "a", 1.5, _Nested(True)) + edited_grid.grid[0, 0] += 1.0 + assert _carrier_fingerprint(edited_grid) != base + assert _carrier_fingerprint(_Carrier(carrier.grid, "b", 1.5, _Nested(True))) != base + assert _carrier_fingerprint(_Carrier(carrier.grid, "a", 2.0, _Nested(True))) != base + assert _carrier_fingerprint(_Carrier(carrier.grid, "a", 1.5, _Nested(False))) != base + + +def test_fingerprint_stable_for_equal_content() -> None: + first = _Carrier(torch.ones(3), "a", 1.5, _Nested(True)) + second = _Carrier(torch.ones(3), "a", 1.5, _Nested(True)) + assert _carrier_fingerprint(first) == _carrier_fingerprint(second) + + +def test_fingerprint_rejects_unsupported_field_type() -> None: + @dataclass(frozen=True, slots=True) + class _Bad: + leak: np.float32 + + with pytest.raises(TypeError, match="unsupported type"): + _carrier_fingerprint(_Bad(np.float32(1.0))) diff --git a/motrix_envs/src/motrix_envs/locomotion/ball_balance/mdp/reset.py b/motrix_envs/src/motrix_envs/locomotion/ball_balance/mdp/reset.py index 0bfbfbd6..4e1298af 100644 --- a/motrix_envs/src/motrix_envs/locomotion/ball_balance/mdp/reset.py +++ b/motrix_envs/src/motrix_envs/locomotion/ball_balance/mdp/reset.py @@ -132,7 +132,7 @@ class DofResetParams: default_pos: torch.Tensor joint_lower: torch.Tensor joint_upper: torch.Tensor - noise_scale: np.float32 + noise_scale: float def _reset_body_dof_pos(ctx: ManagerContext, sim_writes: Map[torch.Tensor], params: DofResetParams) -> None: @@ -173,7 +173,7 @@ def __call__(self, ctx) -> ResetTerm: default_pos=torch.from_numpy(default_pos), joint_lower=torch.from_numpy(joint_lower), joint_upper=torch.from_numpy(joint_upper), - noise_scale=np.float32(self.noise * self.noise_scale), + noise_scale=float(self.noise * self.noise_scale), ), writes={"position": JointPositionWrite(joints), "velocity": JointVelocityWrite(joints)}, ) diff --git a/motrix_envs/src/motrix_envs/locomotion/ball_balance/mdp/rewards.py b/motrix_envs/src/motrix_envs/locomotion/ball_balance/mdp/rewards.py index e58dddc6..74265f15 100644 --- a/motrix_envs/src/motrix_envs/locomotion/ball_balance/mdp/rewards.py +++ b/motrix_envs/src/motrix_envs/locomotion/ball_balance/mdp/rewards.py @@ -74,7 +74,7 @@ def __call__(self, ctx) -> RewardTerm: @dataclass(frozen=True, slots=True) class DofDefaultParams: reference: torch.Tensor - sigma: np.float32 + sigma: float def dof_default_reward(ctx: ManagerContext, params: DofDefaultParams) -> torch.Tensor: @@ -103,5 +103,5 @@ def __call__(self, ctx) -> RewardTerm: if missing: raise KeyError(f"key pose 'default' is missing joints: {missing}") reference = np.asarray([default[name] for name in joint_names], dtype=np.float32) - params = DofDefaultParams(torch.from_numpy(reference), np.float32(self.sigma)) + params = DofDefaultParams(torch.from_numpy(reference), float(self.sigma)) return RewardTerm(dof_default_reward, params) diff --git a/motrix_envs/src/motrix_envs/locomotion/humanoid/walk_manager_mdp/reset.py b/motrix_envs/src/motrix_envs/locomotion/humanoid/walk_manager_mdp/reset.py index 9cb581db..43c1f6dc 100644 --- a/motrix_envs/src/motrix_envs/locomotion/humanoid/walk_manager_mdp/reset.py +++ b/motrix_envs/src/motrix_envs/locomotion/humanoid/walk_manager_mdp/reset.py @@ -75,13 +75,13 @@ class WalkResetParams: default_joint_angles: torch.Tensor init_pose: torch.Tensor heightfield: HeightFieldGrid - spawn_range: np.float32 + spawn_range: float randomization_enabled: bool kp_default: torch.Tensor damping_default: torch.Tensor link_mass_default: torch.Tensor - base_link_index: np.int64 - base_mass_default: np.float32 + base_link_index: int + base_mass_default: float base_com_default: torch.Tensor friction_default: torch.Tensor kp_range: torch.Tensor @@ -232,7 +232,7 @@ def __call__(self, ctx: BuildContext) -> ResetTerm: np.concatenate([body.init_base_position, body.init_base_quat]).astype(np.float32) ), heightfield=ground_height_grid(ctx, self.ground_geom), - spawn_range=np.float32(self.spawn_xy_range), + spawn_range=float(self.spawn_xy_range), ) randomization = self.randomization @@ -256,8 +256,8 @@ def __call__(self, ctx: BuildContext) -> ResetTerm: kp_default=torch.from_numpy(np.asarray(others["actuator_kp"], dtype=np.float32)), damping_default=torch.from_numpy(np.asarray(others["randomize_actuator_kd"], dtype=np.float32)), link_mass_default=torch.from_numpy(link_masses), - base_link_index=np.int64(base_index), - base_mass_default=np.float32(link_masses[base_index]), + base_link_index=base_index, + base_mass_default=float(link_masses[base_index]), base_com_default=torch.from_numpy(np.asarray(others["randomize_base_com"], dtype=np.float32)), friction_default=torch.from_numpy(np.asarray(others["randomize_friction_default"], dtype=np.float32)), kp_range=torch.from_numpy(np.asarray(randomization.kp_scale_range, dtype=np.float32)), @@ -282,8 +282,8 @@ def __call__(self, ctx: BuildContext) -> ResetTerm: kp_default=torch.zeros(0, dtype=torch.float32), damping_default=torch.zeros(0, dtype=torch.float32), link_mass_default=torch.zeros(0, dtype=torch.float32), - base_link_index=np.int64(0), - base_mass_default=np.float32(0.0), + base_link_index=0, + base_mass_default=0.0, base_com_default=torch.zeros(3, dtype=torch.float32), friction_default=torch.zeros(3, dtype=torch.float32), kp_range=torch.ones(2, dtype=torch.float32), diff --git a/motrix_envs/src/motrix_envs/locomotion/humanoid/walk_manager_mdp/terrain.py b/motrix_envs/src/motrix_envs/locomotion/humanoid/walk_manager_mdp/terrain.py index 13aea2cb..d9f63534 100644 --- a/motrix_envs/src/motrix_envs/locomotion/humanoid/walk_manager_mdp/terrain.py +++ b/motrix_envs/src/motrix_envs/locomotion/humanoid/walk_manager_mdp/terrain.py @@ -23,7 +23,7 @@ def ground_height_grid(ctx, ground_geom: str) -> HeightFieldGrid: origin=torch.from_numpy(np.ascontiguousarray(hf["origin"], dtype=np.float32)), spacing=torch.from_numpy(np.ascontiguousarray(hf["spacing"], dtype=np.float32)), z0=torch.from_numpy(np.ascontiguousarray(hf["z0"], dtype=np.float32)), - constant=np.float32(0.0), + constant=0.0, enabled=True, ) spec = ctx.model.others["geoms"][ground_geom] @@ -32,6 +32,6 @@ def ground_height_grid(ctx, ground_geom: str) -> HeightFieldGrid: origin=torch.zeros((2,), dtype=torch.float32), spacing=torch.ones((2,), dtype=torch.float32), z0=torch.zeros((1,), dtype=torch.float32), - constant=np.float32(spec.local_pose[2]), + constant=float(spec.local_pose[2]), enabled=False, ) diff --git a/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/observations.py b/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/observations.py index 5b0fc888..b9ec3552 100644 --- a/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/observations.py +++ b/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/observations.py @@ -122,7 +122,7 @@ def __call__(self, ctx) -> ObsTerm: @dataclass(frozen=True, slots=True) class RelativePositionParams: reference: torch.Tensor - noise_amplitude: np.float32 + noise_amplitude: float def dof_pos_rel_obs(ctx: ManagerContext, out: torch.Tensor, params: RelativePositionParams) -> None: @@ -163,7 +163,7 @@ def __call__(self, ctx) -> ObsTerm: if missing: raise ValueError(f"RobotCfg {robot_name!r} key pose is missing queried joints: {missing}.") reference = np.asarray([positions[name] for name in queried_joints], dtype=np.float32) - params = RelativePositionParams(torch.from_numpy(reference), np.float32(self.noise.amplitude)) + params = RelativePositionParams(torch.from_numpy(reference), float(self.noise.amplitude)) return ObsTerm(reference.shape[0], dof_pos_rel_obs, params) diff --git a/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/rewards.py b/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/rewards.py index b7e06177..5c801061 100644 --- a/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/rewards.py +++ b/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/rewards.py @@ -150,8 +150,8 @@ def __call__(self, ctx) -> RewardTerm: class DofLimitParams: midpoint: torch.Tensor half_range: torch.Tensor - soft_limit: np.float32 - cap: np.float32 + soft_limit: float + cap: float def dof_limit_reward(ctx: ManagerContext, params: DofLimitParams) -> torch.Tensor: @@ -171,8 +171,8 @@ def __call__(self, ctx) -> RewardTerm: params = DofLimitParams( midpoint=(action.joint_lower + action.joint_upper) * 0.5, half_range=(action.joint_upper - action.joint_lower) * 0.5, - soft_limit=np.float32(self.soft_limit), - cap=np.float32(self.cap), + soft_limit=float(self.soft_limit), + cap=float(self.cap), ) return RewardTerm(dof_limit_reward, params) diff --git a/wiki/design/manager/kernel-abi.md b/wiki/design/manager/kernel-abi.md index c22aa386..229244de 100644 --- a/wiki/design/manager/kernel-abi.md +++ b/wiki/design/manager/kernel-abi.md @@ -77,8 +77,10 @@ rebuilt = unflatten_kernel_data(tree_def, leaves) Record 可以嵌套另一个 kernel-data record;不支持递归结构、dict、list、字符串、object array 或 backend handle (SimDataQuery 是 dataclass 但通过 read program 解析,不属于 kernel-data tree)。 -Logical tree 与 execution lowering 分离:同一 concrete value 在 host term 和 context 中保留相同的 tree schema, -resolver 产出 path-keyed 的 resolved tree(leaf 引用),运行时据此构建 namespace 视图。只有 ndarray/scalar leaves 可以成为 resolved leaves;term 内不应产生稳态临时大数组分配。 +Logical tree 与 execution lowering 分离:`ManagerContext` 及其 Map 容器中的 concrete value 保留 tree schema, +resolver 产出 path-keyed 的 resolved tree(leaf 引用),运行时据此构建 namespace 视图。只有 ndarray/scalar leaves 可以成为 resolved leaves;term 内不应产生稳态临时大数组分配。作为 +term 参数传入的 carrier dataclass 不走 tree 机制,按常量参数原样透传(字段仅限 +torch.Tensor、Python 标量、str 与嵌套 carrier dataclass)。 ## 5. Ownership 与 metric @@ -97,7 +99,10 @@ Term 可以读取其他 term 的数据,但不应调用其他 term 的行为方 - 校验 concrete runtime terms 和 `ManagerContext` 的 KernelData schema(leaf 类型); - 收集 read program 的 sim input map; -- 把 context 与 term 参数解析为 path-keyed resolved trees 和参数描述符(常量/叶子/树/sim key); +- 把 context 解析为 path-keyed resolved tree,把 term 参数解析为参数描述符 + (常量/叶子/sim key);carrier dataclass term 参数按常量透传,其结构化 + fingerprint(字段名/类型/张量 dtype+shape+内容哈希/标量值)计入 program + fingerprint; - 校验方法签名、注解契约和 layout; - 产出 `ManagerLayout`、`ReadPlan`、resolved trees 与 per-term-kind typed call records,供 Torch 执行器调度。 diff --git a/wiki/design/manager/torch-executor.md b/wiki/design/manager/torch-executor.md index e2cb9652..234c1720 100644 --- a/wiki/design/manager/torch-executor.md +++ b/wiki/design/manager/torch-executor.md @@ -21,9 +21,12 @@ RNG draw 顺序与 eager 完全一致;Inductor 重结合仅带来 ≤1e-8 量 ## 解析与执行的分工 -- `ManagerProgramResolver`(backend-neutral)负责输入准备:把 kernel-data - dataclass term 参数与 `ManagerContext` 解析为 path-keyed 的 resolved tree(叶子持有 - live value 引用;数组字段已是 torch tensor,执行器无需再包装),构建 `ManagerLayout`、 +- `ManagerProgramResolver`(backend-neutral)负责输入准备:`ManagerContext` + 解析为 path-keyed 的 resolved tree(叶子持有 live value 引用;数组字段已是 + torch tensor,执行器无需再包装);carrier dataclass term 参数按常量参数 + 原样透传,Dynamo 直接 trace 其属性访问,其结构化 fingerprint(字段名/类型/ + 张量 dtype+shape+内容哈希/标量值)计入 program fingerprint 以保持缓存失效 + 强度。构建 `ManagerLayout`、 `ReadPlan`、per-term-kind typed call records(command hooks 携带 receiver 与 static function,reward/termination 携带 buffer column index 与 resolved argument 描述符等),并校验签名注解契约与签名契约。它不做任何代码生成、 From 1a9f7a2a50d5edb1166f7fe1f8c3cf938eac9a2b Mon Sep 17 00:00:00 2001 From: zilch <147668916@qq.com> Date: Fri, 6 Nov 2026 16:33:06 +0800 Subject: [PATCH 29/41] refactor(manager)!: pass array term arguments through and drop LeafArg 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. --- .../src/motrix_env_core/manager/resolver.py | 43 ++++++++----------- .../src/motrix_env_core/mdp/observations.py | 21 +++++---- .../src/motrix_env_core/mdp/rewards.py | 5 +-- .../src/motrix_env_core/mdp/terminations.py | 3 +- .../motrix_env_core/torch_manager/runtime.py | 5 +-- motrix_env_core/tests/test_manager_env.py | 2 +- .../tests/test_term_arg_fingerprint.py | 29 +++++++++++-- .../ball_balance/mdp/observations.py | 5 +-- .../locomotion/ball_balance/mdp/rewards.py | 6 +-- .../ball_balance/mdp/terminations.py | 7 ++- .../humanoid/walk_manager_mdp/observations.py | 5 +-- .../humanoid/walk_manager_mdp/rewards.py | 14 +++--- .../locomotion/wbt/mdp/observations.py | 4 +- .../motrix_envs/locomotion/wbt/mdp/reset.py | 2 +- .../motrix_envs/locomotion/wbt/mdp/rewards.py | 17 ++++---- .../locomotion/wbt/mdp/terminations.py | 11 +++-- 16 files changed, 92 insertions(+), 87 deletions(-) diff --git a/motrix_env_core/src/motrix_env_core/manager/resolver.py b/motrix_env_core/src/motrix_env_core/manager/resolver.py index 89ef7faa..dde903cb 100644 --- a/motrix_env_core/src/motrix_env_core/manager/resolver.py +++ b/motrix_env_core/src/motrix_env_core/manager/resolver.py @@ -151,13 +151,6 @@ class ConstantArg: value: Any -@dataclass(frozen=True) -class LeafArg: - """Term argument resolved as one leaf value (numeric/ndarray argument).""" - - leaf: ResolvedLeaf - - @dataclass(frozen=True) class SimArg: """Term argument resolving to one sim read-program batched array.""" @@ -165,7 +158,7 @@ class SimArg: key: str -PreparedArg = ConstantArg | LeafArg | SimArg +PreparedArg = ConstantArg | SimArg @dataclass(frozen=True) @@ -273,7 +266,6 @@ class ResolvedManagerProgram: "CommandResetHook", "CommandUpdateHook", "ConstantArg", - "LeafArg", "ManagerLayout", "ManagerProgramResolver", "MetricFieldLayout", @@ -328,13 +320,14 @@ def _tree_leaf_fingerprint(tree: ResolvedTree) -> tuple[Any, ...]: def _carrier_fingerprint(value: Any) -> Any: - """Structural + content fingerprint of a raw carrier dataclass term argument. + """Structural + content fingerprint of a constant term argument. Recurses over dataclass fields recording field names, concrete types, tensor/ndarray dtype+shape+content hashes, and Python scalar values. Carriers previously resolved as kernel-data trees contributed no content - to the program fingerprint at all; this makes every carrier data change - (a grid array value, a scalar) invalidate compiled-program caches. + to the program fingerprint at all, and array arguments previously lived + as shared leaves outside the fingerprint; this makes every content + change (a grid array value, a scalar) invalidate compiled-program caches. """ if isinstance(value, torch.Tensor): return ( @@ -359,8 +352,8 @@ def _carrier_fingerprint(value: Any) -> Any: if isinstance(value, (bool, int, float, str)): return (type(value).__name__, value) raise TypeError( - f"Carrier dataclass term argument holds field of unsupported type {type(value).__name__}; " - "allowed: torch.Tensor, ndarray, Python scalars, str, nested carrier dataclasses." + f"Term argument of unsupported type {type(value).__name__}; " + "allowed: torch.Tensor, ndarray, Python scalars, str, carrier dataclasses." ) @@ -737,22 +730,24 @@ def _resolve_dispatch_args(self, group: str, name: str, term: Any) -> tuple[Prep def _resolve_term_arg(self, source_name: str, value: Any) -> PreparedArg: """Resolve one manager term argument into a prepared descriptor. - Carrier dataclass arguments pass through as constants: Dynamo traces - their attribute access directly, and a structural fingerprint over - their raw fields feeds the program fingerprint for cache - invalidation. Float-like scalars and arrays/tensors become shared - leaf slots read by reference; simulator queries resolve to their - read-program batched array; every other value is a literal constant. + Carrier dataclass and tensor/ndarray arguments pass through as + constants: Dynamo traces their attribute access and tensor ops + directly, and a structural fingerprint over their raw fields (dtype, + shape, content hash) feeds the program fingerprint for cache + invalidation. Float-like numpy scalars normalize to Python floats; + simulator queries resolve to their read-program batched array; every + other value is a literal constant. """ if is_kernel_data(value): self._layout_identity.append( ("term_arg", source_name, type(value).__qualname__, _carrier_fingerprint(value)) ) return ConstantArg(value) - if isinstance(value, (float, np.floating)): - return LeafArg(ResolvedLeaf(np.float32(value))) - if isinstance(value, np.ndarray): - return LeafArg(ResolvedLeaf(value)) + if isinstance(value, (torch.Tensor, np.ndarray)): + self._layout_identity.append(("term_arg", source_name, _carrier_fingerprint(value))) + return ConstantArg(value) + if isinstance(value, np.floating): + return ConstantArg(float(value)) if isinstance(value, SimDataQuery): # A query passed as a term argument resolves to its registered # read-program batched view; the dispatch receives the array in diff --git a/motrix_env_core/src/motrix_env_core/mdp/observations.py b/motrix_env_core/src/motrix_env_core/mdp/observations.py index 00408c3c..6c1211ed 100644 --- a/motrix_env_core/src/motrix_env_core/mdp/observations.py +++ b/motrix_env_core/src/motrix_env_core/mdp/observations.py @@ -7,7 +7,6 @@ from typing import TYPE_CHECKING -import numpy as np import torch from motrix_env_core.config import configclass @@ -118,8 +117,8 @@ def __call__(self, ctx: BuildContext) -> ObsTerm: len(body.joint_names), body_joint_vel_obs, BodyJointVelocityQuery(body=body.base_link_name), - np.float32(self.scale), - np.float32(self.noise.amplitude), + float(self.scale), + float(self.noise.amplitude), ) @@ -151,8 +150,8 @@ def __call__(self, ctx: BuildContext) -> ObsTerm: body_linear_velocity_obs, LinkQuaternionQuery(link=body.base_link_name), LinkLinearVelocityQuery(link=body.base_link_name), - np.float32(self.scale), - np.float32(self.noise.amplitude), + float(self.scale), + float(self.noise.amplitude), ) @@ -184,8 +183,8 @@ def __call__(self, ctx: BuildContext) -> ObsTerm: body_angular_velocity_obs, LinkQuaternionQuery(link=body.base_link_name), LinkAngularVelocityQuery(link=body.base_link_name), - np.float32(self.scale), - np.float32(self.noise.amplitude), + float(self.scale), + float(self.noise.amplitude), ) @@ -211,7 +210,7 @@ def __call__(self, ctx: BuildContext) -> ObsTerm: 3, body_projected_gravity_obs, LinkQuaternionQuery(link=body.base_link_name), - np.float32(self.noise.amplitude), + float(self.noise.amplitude), ) @@ -247,9 +246,9 @@ def __call__(self, ctx: BuildContext) -> ObsTerm: int(body.init_joint_pos.shape[0]), body_joint_pos_rel_obs, BodyJointPositionQuery(body=body.base_link_name), - body.init_joint_pos, - np.float32(self.scale), - np.float32(self.noise.amplitude), + torch.from_numpy(body.init_joint_pos), + float(self.scale), + float(self.noise.amplitude), ) diff --git a/motrix_env_core/src/motrix_env_core/mdp/rewards.py b/motrix_env_core/src/motrix_env_core/mdp/rewards.py index 631401fe..4220c00f 100644 --- a/motrix_env_core/src/motrix_env_core/mdp/rewards.py +++ b/motrix_env_core/src/motrix_env_core/mdp/rewards.py @@ -3,7 +3,6 @@ """Reusable reward terms for manager-based environments.""" -import numpy as np import torch from motrix_env_core.config import configclass @@ -70,7 +69,7 @@ def __call__(self, ctx: BuildContext) -> RewardTerm: link = ctx.model.bodies[self.body].base_link_name return RewardTerm( tracking_lin_vel_xy_reward, - np.float32(self.sigma), + float(self.sigma), LinkQuaternionQuery(link=link), LinkLinearVelocityQuery(link=link), self.command_name, @@ -102,7 +101,7 @@ def __call__(self, ctx: BuildContext) -> RewardTerm: link = ctx.model.bodies[self.body].base_link_name return RewardTerm( tracking_ang_vel_z_reward, - np.float32(self.sigma), + float(self.sigma), LinkQuaternionQuery(link=link), LinkAngularVelocityQuery(link=link), self.command_name, diff --git a/motrix_env_core/src/motrix_env_core/mdp/terminations.py b/motrix_env_core/src/motrix_env_core/mdp/terminations.py index 519da342..36cceaf0 100644 --- a/motrix_env_core/src/motrix_env_core/mdp/terminations.py +++ b/motrix_env_core/src/motrix_env_core/mdp/terminations.py @@ -3,7 +3,6 @@ """Reusable termination terms for manager-based environments.""" -import numpy as np import torch from motrix_env_core.config import configclass @@ -64,7 +63,7 @@ def __call__(self, ctx: BuildContext) -> TerminationTerm: return TerminationTerm( bad_dof_velocity_termination, BodyJointVelocityQuery(body=body.base_link_name), - np.float32(self.threshold), + float(self.threshold), metric_names=("dof_vel_abs_max",), ) diff --git a/motrix_env_core/src/motrix_env_core/torch_manager/runtime.py b/motrix_env_core/src/motrix_env_core/torch_manager/runtime.py index 9008db4e..60d4e9b3 100644 --- a/motrix_env_core/src/motrix_env_core/torch_manager/runtime.py +++ b/motrix_env_core/src/motrix_env_core/torch_manager/runtime.py @@ -39,7 +39,6 @@ CommandAdvanceHook, CommandResetHook, CommandUpdateHook, - LeafArg, ObservationTermCall, RewardTermCall, SimArg, @@ -744,9 +743,7 @@ def _invocation_args( assert self._tree_tensors is not None args: list[Any] = [] for arg in call.args: - if isinstance(arg, LeafArg): - args.append(self._cache.tensor(inputs, arg.leaf.value)) - elif isinstance(arg, SimArg): + if isinstance(arg, SimArg): assert self._context_tensors is not None args.append(self._context_tensors[self._sim_paths[arg.key]]) else: diff --git a/motrix_env_core/tests/test_manager_env.py b/motrix_env_core/tests/test_manager_env.py index a59399ed..a2df00a0 100644 --- a/motrix_env_core/tests/test_manager_env.py +++ b/motrix_env_core/tests/test_manager_env.py @@ -183,7 +183,7 @@ def __call__(self, env: ManagerEnv) -> TerminationTerm: del env return TerminationTerm( _injected_termination, - np.float32(self.threshold), + float(self.threshold), metric_names=("source_at_termination",), ) diff --git a/motrix_env_core/tests/test_term_arg_fingerprint.py b/motrix_env_core/tests/test_term_arg_fingerprint.py index 9e8d8d28..12300122 100644 --- a/motrix_env_core/tests/test_term_arg_fingerprint.py +++ b/motrix_env_core/tests/test_term_arg_fingerprint.py @@ -1,11 +1,12 @@ # Copyright Motphys Technology Co., Ltd. 2025, 2026 # SPDX-License-Identifier: Apache-2.0 -"""Carrier dataclass term-argument fingerprint contract. +"""Constant term-argument fingerprint contract. -Carrier dataclasses passed as manager term arguments resolve as constants; -their structural fingerprint must keep compiled-program caches invalidated by -any content change (tensor bytes or scalar values). +Carrier dataclasses and raw tensors/ndarrays passed as manager term +arguments resolve as constants; their structural fingerprint must keep +compiled-program caches invalidated by any content change (tensor bytes or +scalar values). """ from dataclasses import dataclass @@ -47,6 +48,26 @@ def test_fingerprint_stable_for_equal_content() -> None: assert _carrier_fingerprint(first) == _carrier_fingerprint(second) +def test_raw_tensor_arg_fingerprint_flips_on_content_edit() -> None: + arg = torch.arange(4, dtype=torch.float32) + base = _carrier_fingerprint(arg) + assert _carrier_fingerprint(torch.arange(4, dtype=torch.float32)) == base + assert _carrier_fingerprint(torch.arange(4, dtype=torch.float64)) != base + assert _carrier_fingerprint(torch.arange(5, dtype=torch.float32)) != base + edited = arg.clone() + edited[0] += 1.0 + assert _carrier_fingerprint(edited) != base + + +def test_raw_ndarray_arg_fingerprint_flips_on_content_edit() -> None: + arg = np.arange(4, dtype=np.float32) + base = _carrier_fingerprint(arg) + assert _carrier_fingerprint(np.arange(4, dtype=np.float32)) == base + edited = arg.copy() + edited[0] += 1.0 + assert _carrier_fingerprint(edited) != base + + def test_fingerprint_rejects_unsupported_field_type() -> None: @dataclass(frozen=True, slots=True) class _Bad: diff --git a/motrix_envs/src/motrix_envs/locomotion/ball_balance/mdp/observations.py b/motrix_envs/src/motrix_envs/locomotion/ball_balance/mdp/observations.py index 07407af8..dd2a01d3 100644 --- a/motrix_envs/src/motrix_envs/locomotion/ball_balance/mdp/observations.py +++ b/motrix_envs/src/motrix_envs/locomotion/ball_balance/mdp/observations.py @@ -3,7 +3,6 @@ """Observation terms for the manager-based ball-balance task.""" -import numpy as np import torch from motrix_env_core.config import configclass @@ -30,7 +29,7 @@ class BallRelativePositionObsCfg(ObservationTermCfg): noise: UniformNoiseCfg = UniformNoiseCfg() def __call__(self, ctx) -> ObsTerm: - return ObsTerm(3, ball_relative_position_obs, np.float32(self.noise.amplitude)) + return ObsTerm(3, ball_relative_position_obs, float(self.noise.amplitude)) def ball_relative_velocity_obs(ctx: ManagerContext, out: torch.Tensor, noise_amplitude: float) -> None: @@ -46,7 +45,7 @@ class BallRelativeVelocityObsCfg(ObservationTermCfg): noise: UniformNoiseCfg = UniformNoiseCfg() def __call__(self, ctx) -> ObsTerm: - return ObsTerm(3, ball_relative_velocity_obs, np.float32(self.noise.amplitude)) + return ObsTerm(3, ball_relative_velocity_obs, float(self.noise.amplitude)) def ball_position_obs(ctx: ManagerContext, out: torch.Tensor) -> None: diff --git a/motrix_envs/src/motrix_envs/locomotion/ball_balance/mdp/rewards.py b/motrix_envs/src/motrix_envs/locomotion/ball_balance/mdp/rewards.py index 74265f15..216eb3ff 100644 --- a/motrix_envs/src/motrix_envs/locomotion/ball_balance/mdp/rewards.py +++ b/motrix_envs/src/motrix_envs/locomotion/ball_balance/mdp/rewards.py @@ -31,7 +31,7 @@ class UprightRewardCfg(RewardTermCfg): sigma: float def __call__(self, ctx) -> RewardTerm: - return RewardTerm(upright_reward, np.float32(self.sigma)) + return RewardTerm(upright_reward, float(self.sigma)) def base_height_reward(ctx: ManagerContext, target_z: float, sigma: float) -> torch.Tensor: @@ -47,7 +47,7 @@ class BaseHeightRewardCfg(RewardTermCfg): sigma: float def __call__(self, ctx) -> RewardTerm: - return RewardTerm(base_height_reward, np.float32(self.target_z), np.float32(self.sigma)) + return RewardTerm(base_height_reward, float(self.target_z), float(self.sigma)) def ball_under_feet_reward(ctx: ManagerContext, sigma: float) -> torch.Tensor: @@ -68,7 +68,7 @@ class BallUnderFeetRewardCfg(RewardTermCfg): sigma: float def __call__(self, ctx) -> RewardTerm: - return RewardTerm(ball_under_feet_reward, np.float32(self.sigma)) + return RewardTerm(ball_under_feet_reward, float(self.sigma)) @dataclass(frozen=True, slots=True) diff --git a/motrix_envs/src/motrix_envs/locomotion/ball_balance/mdp/terminations.py b/motrix_envs/src/motrix_envs/locomotion/ball_balance/mdp/terminations.py index fb2f4bee..4895226e 100644 --- a/motrix_envs/src/motrix_envs/locomotion/ball_balance/mdp/terminations.py +++ b/motrix_envs/src/motrix_envs/locomotion/ball_balance/mdp/terminations.py @@ -3,7 +3,6 @@ """Termination terms for the manager-based ball-balance task.""" -import numpy as np import torch from motrix_env_core.config import configclass @@ -29,7 +28,7 @@ class BadBaseZTerminationCfg(_BallBalanceTerminationCfg): def __call__(self, ctx) -> TerminationTerm: return TerminationTerm( bad_base_z_termination, - np.float32(self.threshold), + float(self.threshold), metric_names=("base_z",), ) @@ -51,7 +50,7 @@ class BadOrientationTerminationCfg(_BallBalanceTerminationCfg): def __call__(self, ctx) -> TerminationTerm: return TerminationTerm( bad_orientation_termination, - np.float32(self.threshold), + float(self.threshold), metric_names=("base_tilt",), ) @@ -75,6 +74,6 @@ class BallEscapedTerminationCfg(_BallBalanceTerminationCfg): def __call__(self, ctx) -> TerminationTerm: return TerminationTerm( ball_escaped_termination, - np.float32(self.threshold), + float(self.threshold), metric_names=("ball_feet_distance",), ) diff --git a/motrix_envs/src/motrix_envs/locomotion/humanoid/walk_manager_mdp/observations.py b/motrix_envs/src/motrix_envs/locomotion/humanoid/walk_manager_mdp/observations.py index 265b7712..cb534d33 100644 --- a/motrix_envs/src/motrix_envs/locomotion/humanoid/walk_manager_mdp/observations.py +++ b/motrix_envs/src/motrix_envs/locomotion/humanoid/walk_manager_mdp/observations.py @@ -3,7 +3,6 @@ """Observation terms for the humanoid velocity-tracking task.""" -import numpy as np import torch from motrix_env_core.config import configclass @@ -11,7 +10,7 @@ from motrix_env_core.manager.observations import ObservationTermCfg, ObsTerm -def gait_phase_obs(ctx: ManagerContext, out: torch.Tensor, offset: np.int64, size: np.int64) -> None: +def gait_phase_obs(ctx: ManagerContext, out: torch.Tensor, offset: int, size: int) -> None: walk = ctx.commands["walk"] out.copy_(walk.sin_cos[:, offset : offset + size]) @@ -29,4 +28,4 @@ class GaitPhaseObsCfg(ObservationTermCfg): size: int = 2 def __call__(self, ctx: BuildContext) -> ObsTerm: - return ObsTerm(self.size, gait_phase_obs, np.int64(self.offset), np.int64(self.size)) + return ObsTerm(self.size, gait_phase_obs, self.offset, self.size) diff --git a/motrix_envs/src/motrix_envs/locomotion/humanoid/walk_manager_mdp/rewards.py b/motrix_envs/src/motrix_envs/locomotion/humanoid/walk_manager_mdp/rewards.py index ebee3d8d..6d5a2ffa 100644 --- a/motrix_envs/src/motrix_envs/locomotion/humanoid/walk_manager_mdp/rewards.py +++ b/motrix_envs/src/motrix_envs/locomotion/humanoid/walk_manager_mdp/rewards.py @@ -125,8 +125,8 @@ def __call__(self, ctx) -> RewardTerm: raise ValueError("FeetPhaseRewardCfg requires sole_l_site, sole_r_site, and ground_geom.") return RewardTerm( feet_phase_reward, - np.float32(self.swing_height), - np.float32(self.feet_phase_sigma), + float(self.swing_height), + float(self.feet_phase_sigma), ground_height_grid(ctx, self.ground_geom), SitePositionQuery(site=self.sole_l_site), SitePositionQuery(site=self.sole_r_site), @@ -163,8 +163,8 @@ def __call__(self, ctx) -> RewardTerm: # resolved) so it stays aligned with body.init_joint_pos and weights. return RewardTerm( pose_reward, - body.init_joint_pos, - weights, + torch.from_numpy(body.init_joint_pos), + torch.from_numpy(weights), JointPositionQuery(joints=body.joint_names), ) @@ -195,7 +195,7 @@ def __call__(self, ctx) -> RewardTerm: body = ctx.model.bodies["robot"] return RewardTerm( penalty_close_feet_xy_reward, - np.float32(self.close_feet_threshold), + float(self.close_feet_threshold), LinkQuaternionQuery(link=body.base_link_name), BatchLinkPositionQuery(links=ctx.cfg.scene.objs.robot.resolved_foot_link_names), ) @@ -209,7 +209,7 @@ def penalty_feet_ori_reward( gravity = foot_quat.new_zeros(*foot_quat.shape[:2], 3) gravity[..., 2] = -1.0 gx, gy, gz = rotate_inverse_components(foot_quat.reshape(-1, 4), gravity.reshape(-1, 3)) - reference = torch.as_tensor(default_foot_gravity) + reference = default_foot_gravity total = None for foot in range(num_feet): # Flat layout is lane-major (lane * F + foot): strided slices pick one foot. @@ -232,7 +232,7 @@ def __call__(self, ctx) -> RewardTerm: robot = ctx.cfg.scene.objs.robot # Default foot gravity at the init key pose, from BodyModel's # compile-time init-pose FK snapshot. - default_foot_gravity = _default_foot_gravity(ctx, robot) + default_foot_gravity = torch.from_numpy(_default_foot_gravity(ctx, robot)) return RewardTerm( penalty_feet_ori_reward, default_foot_gravity, diff --git a/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/observations.py b/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/observations.py index b9ec3552..7f756a04 100644 --- a/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/observations.py +++ b/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/observations.py @@ -76,7 +76,7 @@ class MotionReferenceOrientationObsCfg(ObservationTermCfg): noise: UniformNoiseCfg = UniformNoiseCfg() def __call__(self, ctx) -> ObsTerm: - return ObsTerm(6, motion_reference_orientation_obs, np.float32(self.noise.amplitude)) + return ObsTerm(6, motion_reference_orientation_obs, float(self.noise.amplitude)) def robot_body_position_in_reference_frame_obs(ctx: ManagerContext, out: torch.Tensor) -> None: @@ -180,4 +180,4 @@ class DofVelObsCfg(ObservationTermCfg): def __call__(self, ctx) -> ObsTerm: num_joints = len(ctx.model.bodies["robot"].joint_names) - return ObsTerm(num_joints, dof_vel_obs, np.float32(self.noise.amplitude)) + return ObsTerm(num_joints, dof_vel_obs, float(self.noise.amplitude)) diff --git a/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/reset.py b/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/reset.py index 6926dfd3..7acb2b6d 100644 --- a/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/reset.py +++ b/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/reset.py @@ -190,6 +190,6 @@ def __call__(self, ctx) -> ResetTerm: joints = ctx.cfg.commands.motion.joint_names return ResetTerm( _reset_body_dof_pos, - np.float32(self.noise * self.noise_scale), + float(self.noise * self.noise_scale), writes={"position": JointPositionWrite(joints), "velocity": JointVelocityWrite(joints)}, ) diff --git a/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/rewards.py b/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/rewards.py index 5c801061..5185da5a 100644 --- a/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/rewards.py +++ b/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/rewards.py @@ -6,7 +6,6 @@ from dataclasses import dataclass from typing import cast -import numpy as np import torch from motrix_env_core.config import configclass @@ -27,7 +26,7 @@ class GlobalRefPositionRewardCfg(RewardTermCfg): sigma: float def __call__(self, ctx) -> RewardTerm: - return RewardTerm(global_ref_position_reward, np.float32(self.sigma)) + return RewardTerm(global_ref_position_reward, float(self.sigma)) def global_ref_orientation_reward(ctx: ManagerContext, sigma: float) -> torch.Tensor: @@ -44,7 +43,7 @@ class GlobalRefOrientationRewardCfg(RewardTermCfg): sigma: float def __call__(self, ctx) -> RewardTerm: - return RewardTerm(global_ref_orientation_reward, np.float32(self.sigma)) + return RewardTerm(global_ref_orientation_reward, float(self.sigma)) def relative_body_position_reward(ctx: ManagerContext, sigma: float) -> torch.Tensor: @@ -59,7 +58,7 @@ class RelativeBodyPositionRewardCfg(RewardTermCfg): sigma: float def __call__(self, ctx) -> RewardTerm: - return RewardTerm(relative_body_position_reward, np.float32(self.sigma)) + return RewardTerm(relative_body_position_reward, float(self.sigma)) def ee_body_pos_reward(ctx: ManagerContext, body_indices: tuple[int, ...], sigma: float) -> torch.Tensor: @@ -93,7 +92,7 @@ def __call__(self, ctx) -> RewardTerm: ) tracked_body_names = ctx.cfg.commands.motion.tracked_body_names body_indices = tuple(tracked_body_names.index(name) for name in self.body_names) - return RewardTerm(ee_body_pos_reward, body_indices, np.float32(self.sigma)) + return RewardTerm(ee_body_pos_reward, body_indices, float(self.sigma)) def relative_body_orientation_reward(ctx: ManagerContext, sigma: float) -> torch.Tensor: @@ -109,7 +108,7 @@ class RelativeBodyOrientationRewardCfg(RewardTermCfg): sigma: float def __call__(self, ctx) -> RewardTerm: - return RewardTerm(relative_body_orientation_reward, np.float32(self.sigma)) + return RewardTerm(relative_body_orientation_reward, float(self.sigma)) def global_body_linear_velocity_reward(ctx: ManagerContext, sigma: float) -> torch.Tensor: @@ -126,7 +125,7 @@ class GlobalBodyLinearVelocityRewardCfg(RewardTermCfg): sigma: float def __call__(self, ctx) -> RewardTerm: - return RewardTerm(global_body_linear_velocity_reward, np.float32(self.sigma)) + return RewardTerm(global_body_linear_velocity_reward, float(self.sigma)) def global_body_angular_velocity_reward(ctx: ManagerContext, sigma: float) -> torch.Tensor: @@ -143,7 +142,7 @@ class GlobalBodyAngularVelocityRewardCfg(RewardTermCfg): sigma: float def __call__(self, ctx) -> RewardTerm: - return RewardTerm(global_body_angular_velocity_reward, np.float32(self.sigma)) + return RewardTerm(global_body_angular_velocity_reward, float(self.sigma)) @dataclass(frozen=True, slots=True) @@ -189,4 +188,4 @@ class UndesiredContactsRewardCfg(RewardTermCfg): threshold: float def __call__(self, ctx) -> RewardTerm: - return RewardTerm(undesired_contacts_reward, np.float32(self.threshold)) + return RewardTerm(undesired_contacts_reward, float(self.threshold)) diff --git a/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/terminations.py b/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/terminations.py index 601af67a..8538a16d 100644 --- a/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/terminations.py +++ b/motrix_envs/src/motrix_envs/locomotion/wbt/mdp/terminations.py @@ -3,7 +3,6 @@ """WBT termination terms and raw diagnostic metrics.""" -import numpy as np import torch from motrix_env_core.config import configclass @@ -30,7 +29,7 @@ class BadRefZTerminationCfg(_WbtTerminationCfg): def __call__(self, ctx) -> TerminationTerm: return TerminationTerm( bad_ref_z_termination, - np.float32(self.threshold), + float(self.threshold), metric_names=("ref_z_abs_err",), ) @@ -50,7 +49,7 @@ class BadRefOrientationTerminationCfg(_WbtTerminationCfg): def __call__(self, ctx) -> TerminationTerm: return TerminationTerm( bad_ref_orientation_termination, - np.float32(self.threshold), + float(self.threshold), metric_names=("ref_ori_abs_err",), ) @@ -89,7 +88,7 @@ def __call__(self, ctx) -> TerminationTerm: return TerminationTerm( bad_motion_body_position_termination, body_indices, - np.float32(self.threshold), + float(self.threshold), metric_names=("motion_body_pos_err",), ) @@ -120,7 +119,7 @@ def __call__(self, ctx) -> TerminationTerm: return TerminationTerm( bad_body_z_termination, body_indices, - np.float32(self.threshold), + float(self.threshold), metric_names=("body_z_max_err",), ) @@ -142,6 +141,6 @@ class BadDofPositionTerminationCfg(_WbtTerminationCfg): def __call__(self, ctx) -> TerminationTerm: return TerminationTerm( bad_dof_position_termination, - np.float32(self.threshold), + float(self.threshold), metric_names=("dof_limit_violation_max",), ) From a12e76cb983625b26a5ae30928d779fefec7db57 Mon Sep 17 00:00:00 2001 From: zilch <147668916@qq.com> Date: Fri, 6 Nov 2026 20:08:53 +0800 Subject: [PATCH 30/41] refactor(core)!: parameterize the env state by array type 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. --- .../src/motrix_env_core/array/env.py | 78 ++++++++++--------- .../src/motrix_env_core/direct/__init__.py | 4 +- .../src/motrix_env_core/direct/env.py | 4 +- .../src/motrix_env_core/manager/env.py | 30 +++---- motrix_env_core/tests/test_array_env.py | 6 +- .../locomotion/quadruped/walk_np.py | 4 +- .../franka_lift_cube/franka_lift_cube_np.py | 4 +- .../franka_open_cabinet_np.py | 4 +- .../rm65_insert_peg/insert_peg_np.py | 4 +- .../rm65_open_cabinet/rm65_open_cabinet_np.py | 4 +- .../shadow_hand/shadow_hand_np.py | 4 +- .../tests/test_basic_direct_contract.py | 6 +- motrix_envs/tests/test_basic_env_lifecycle.py | 4 +- .../test_manipulation_direct_contract.py | 8 +- .../tests/test_builtin_sim_backend.py | 6 +- .../tests/test_fastsac_collector.py | 4 +- .../tests/test_fastsac_final_obs.py | 6 +- .../src/motrix_rl_rslrl/torch/wrap_np.py | 4 +- .../tests/test_rslrl_sim_backend.py | 4 +- motrix_rl_skrl/tests/test_skrl_sim_backend.py | 4 +- wiki/design/array-env-final-observation.md | 2 +- wiki/design/manager/runtime.md | 2 +- 22 files changed, 102 insertions(+), 94 deletions(-) diff --git a/motrix_env_core/src/motrix_env_core/array/env.py b/motrix_env_core/src/motrix_env_core/array/env.py index 6412ef8e..fff7bb95 100644 --- a/motrix_env_core/src/motrix_env_core/array/env.py +++ b/motrix_env_core/src/motrix_env_core/array/env.py @@ -23,22 +23,28 @@ # (back-compat). EnvCfgType = TypeVar("EnvCfgType", bound=EnvCfg) +# The array flavor an environment's state carries end to end: numpy arrays for +# NumPy-frontend environments, torch tensors for the manager pipeline. Each +# environment consistently produces one of the two; the lifecycle is +# duck-typed on it and never branches per element type. +T = TypeVar("T", np.ndarray, torch.Tensor) + @dataclass -class NpObs: +class EnvObs(Generic[T]): """One observation slot; numpy arrays for NumPy-frontend environments, torch tensors for the manager pipeline.""" - policy: "torch.Tensor | np.ndarray" - value: "torch.Tensor | np.ndarray | None" = None + policy: T + value: T | None = None @property - def value_or_policy(self) -> "torch.Tensor | np.ndarray": + def value_or_policy(self) -> T: return self.policy if self.value is None else self.value @dataclass -class FinalObs: +class FinalObs(Generic[T]): """Owned pre-reset observations of the current transition's truncated rows. Rows in ``obs`` align with ``env_ids``. Arrays remain valid across steps; @@ -46,11 +52,11 @@ class FinalObs: """ env_ids: np.ndarray - obs: NpObs + obs: EnvObs[T] @dataclass -class ArrayEnvState: +class ArrayEnvState(Generic[T]): """Environment transition state. ``obs``, ``reward``, ``terminated``, ``truncated`` and ``final_obs`` store @@ -59,16 +65,16 @@ class ArrayEnvState: two, so consumers never branch on the element type. """ - obs: NpObs - reward: "torch.Tensor | np.ndarray" - terminated: "torch.Tensor | np.ndarray" - truncated: "torch.Tensor | np.ndarray" + obs: EnvObs[T] + reward: T + terminated: T + truncated: T episode_steps: np.ndarray # Per-term reward breakdown for the latest transition: term name to a # ``(num_envs,)`` array (numpy for the NumPy frontend, torch column views # for the manager pipeline). Freshly written every transition; empty when # the environment does not decompose its reward. - reward_terms: dict[str, "torch.Tensor | np.ndarray"] = dataclasses.field(default_factory=dict) + reward_terms: dict[str, T] = dataclasses.field(default_factory=dict) # Live diagnostics view: per-environment quantities stay unreduced as # ``(num_envs,)`` arrays and are views into manager buffers that kernels # overwrite every step; batch-level gauges are scalars. Values always @@ -76,16 +82,16 @@ class ArrayEnvState: # :meth:`process_metrics`, which returns a fresh dict of reduced scalars. metrics: dict[str, float | np.ndarray] = dataclasses.field(default_factory=dict) # None when no row truncated or the implementation does not capture finals. - final_obs: FinalObs | None = None + final_obs: FinalObs[T] | None = None @property - def done(self) -> "torch.Tensor | np.ndarray": + def done(self) -> T: """ Check if the environment is done. """ return self.terminated | self.truncated - def replace(self, **updates) -> "ArrayEnvState": + def replace(self, **updates) -> "ArrayEnvState[T]": return dataclasses.replace(self, **updates) def process_metrics(self) -> dict[str, float]: @@ -118,14 +124,14 @@ def _as_obs_space(space: gym.Space | ObsSpace) -> ObsSpace: return ObsSpace(policy=space) -def _as_obs(obs: NpObs | np.ndarray) -> NpObs: - if isinstance(obs, NpObs): +def _as_obs(obs: "EnvObs[T] | np.ndarray") -> "EnvObs[T]": + if isinstance(obs, EnvObs): return obs assert isinstance(obs, np.ndarray) - return NpObs(policy=obs) + return EnvObs(policy=obs) -class ArrayEnv(ABEnv, Generic[EnvCfgType]): +class ArrayEnv(ABEnv, Generic[EnvCfgType, T]): """Simulator-agnostic NumPy frontend environment lifecycle. ``ArrayEnv`` owns the step / auto-reset / truncation lifecycle shared by @@ -135,7 +141,7 @@ class ArrayEnv(ABEnv, Generic[EnvCfgType]): """ _cfg: EnvCfgType - _state: ArrayEnvState = None + _state: ArrayEnvState[T] = None _render_spacing: float def __init__(self, cfg: EnvCfgType, num_envs: int = 1): @@ -147,7 +153,7 @@ def __init__(self, cfg: EnvCfgType, num_envs: int = 1): self._render_spacing = cfg.render_spacing @property - def state(self) -> ArrayEnvState: + def state(self) -> ArrayEnvState[T]: """ Get the current environment state """ @@ -171,7 +177,7 @@ def render_spacing(self) -> float: def num_envs(self) -> int: return self._num_envs - def init_state(self) -> ArrayEnvState: + def init_state(self) -> ArrayEnvState[T]: """ Create a new environment state """ @@ -190,12 +196,12 @@ def init_state(self) -> ArrayEnvState: def _new_state( self, - obs: NpObs, - reward: np.ndarray, - terminated: np.ndarray, - truncated: np.ndarray, + obs: EnvObs[T], + reward: T, + terminated: T, + truncated: T, episode_steps: np.ndarray, - ) -> ArrayEnvState: + ) -> ArrayEnvState[T]: """Assemble the environment state; subclasses add simulator-owned fields.""" return ArrayEnvState( obs=obs, @@ -229,12 +235,12 @@ def _zeros_terminated(self) -> np.ndarray: def _zeros_truncated(self) -> np.ndarray: return np.zeros((self._num_envs,), dtype=bool) - def _zeros_obs(self, space: gym.Space | ObsSpace) -> NpObs: + def _zeros_obs(self, space: gym.Space | ObsSpace) -> EnvObs[T]: obs_space = _as_obs_space(space) value = None if obs_space.value is None else self._zeros_box(obs_space.value) - return NpObs(policy=self._zeros_box(obs_space.policy), value=value) + return EnvObs(policy=self._zeros_box(obs_space.policy), value=value) - def _assign_obs(self, dst: NpObs, mask: np.ndarray, src: NpObs | np.ndarray): + def _assign_obs(self, dst: EnvObs[T], mask: np.ndarray, src: "EnvObs[T] | np.ndarray"): src = _as_obs(src) dst.policy[mask] = src.policy if dst.value is not None: @@ -272,7 +278,7 @@ def reset(self, env_ids: np.ndarray) -> None: writes have completed. """ - def _capture_final_obs(self, env_ids: np.ndarray) -> FinalObs | None: + def _capture_final_obs(self, env_ids: np.ndarray) -> "FinalObs[T] | None": """Capture truncated rows before reset; default implementation is unsupported. Called after transition and truncation updates, before any reset write. @@ -305,7 +311,7 @@ def create_renderer(self, config: RenderConfig) -> SimRenderer: """Create the backend-owned renderer bound to this environment.""" @abc.abstractmethod - def apply_action(self, actions: np.ndarray, state: ArrayEnvState) -> ArrayEnvState: + def apply_action(self, actions: np.ndarray, state: ArrayEnvState[T]) -> ArrayEnvState[T]: """ Apply the action to the environment @@ -315,7 +321,7 @@ def apply_action(self, actions: np.ndarray, state: ArrayEnvState) -> ArrayEnvSta """ @abc.abstractmethod - def compute_transition(self, state: ArrayEnvState) -> ArrayEnvState: + def compute_transition(self, state: ArrayEnvState[T]) -> ArrayEnvState[T]: """ Update the environment state after physics step @@ -340,7 +346,7 @@ def _prev_physics_step(self): state.terminated[:] = False state.truncated[:] = False - def compute_observation(self, state: ArrayEnvState) -> ArrayEnvState: + def compute_observation(self, state: ArrayEnvState[T]) -> ArrayEnvState[T]: """Compute observations for all environments after reset. Called once after the first reset of :meth:`init_state` and once per @@ -353,7 +359,7 @@ def compute_observation(self, state: ArrayEnvState) -> ArrayEnvState: return state @perf_root("step") - def step(self, actions: np.ndarray) -> ArrayEnvState: + def step(self, actions: np.ndarray) -> ArrayEnvState[T]: """Advance one control step. Timing: ``apply_action`` -> ``physics_step`` -> ``compute_transition`` @@ -388,4 +394,4 @@ def step(self, actions: np.ndarray) -> ArrayEnvState: return self._state -__all__ = ["ArrayEnv", "ArrayEnvState", "EnvCfgType", "FinalObs", "NpObs"] +__all__ = ["ArrayEnv", "ArrayEnvState", "EnvCfgType", "FinalObs", "EnvObs"] diff --git a/motrix_env_core/src/motrix_env_core/direct/__init__.py b/motrix_env_core/src/motrix_env_core/direct/__init__.py index 44a13912..a461ea71 100644 --- a/motrix_env_core/src/motrix_env_core/direct/__init__.py +++ b/motrix_env_core/src/motrix_env_core/direct/__init__.py @@ -1,7 +1,7 @@ # Copyright Motphys Technology Co., Ltd. 2025, 2026 # SPDX-License-Identifier: Apache-2.0 -from motrix_env_core.array.env import ArrayEnvState, NpObs +from motrix_env_core.array.env import ArrayEnvState, EnvObs from motrix_env_core.direct.env import DirectEnv, DirectEnvCfg -__all__ = ["ArrayEnvState", "DirectEnv", "DirectEnvCfg", "NpObs"] +__all__ = ["ArrayEnvState", "DirectEnv", "DirectEnvCfg", "EnvObs"] diff --git a/motrix_env_core/src/motrix_env_core/direct/env.py b/motrix_env_core/src/motrix_env_core/direct/env.py index c47a42db..775677b9 100644 --- a/motrix_env_core/src/motrix_env_core/direct/env.py +++ b/motrix_env_core/src/motrix_env_core/direct/env.py @@ -13,6 +13,8 @@ from typing import TypeVar +import numpy as np + from motrix_env_core.array.env import ArrayEnv from motrix_env_core.array.env import ArrayEnvState as ArrayEnvState from motrix_env_core.base import EnvCfg @@ -32,7 +34,7 @@ class DirectEnvCfg(EnvCfg): DirectEnvCfgType = TypeVar("DirectEnvCfgType", bound=DirectEnvCfg) -class DirectEnv(ArrayEnv[DirectEnvCfgType]): +class DirectEnv(ArrayEnv[DirectEnvCfgType, np.ndarray]): """Direct-workflow frontend that exposes only the simulator backend. Concrete environments own their model/data query declarations and compile diff --git a/motrix_env_core/src/motrix_env_core/manager/env.py b/motrix_env_core/src/motrix_env_core/manager/env.py index 35c62548..9da14f61 100644 --- a/motrix_env_core/src/motrix_env_core/manager/env.py +++ b/motrix_env_core/src/motrix_env_core/manager/env.py @@ -13,7 +13,7 @@ import numpy as np import torch -from motrix_env_core.array.env import ArrayEnv, ArrayEnvState, EnvCfgType, FinalObs, NpObs +from motrix_env_core.array.env import ArrayEnv, ArrayEnvState, EnvCfgType, EnvObs, FinalObs from motrix_env_core.base import EnvCfg, ObsSpace from motrix_env_core.config import configclass from motrix_env_core.config.sim_reset import ManagerResetCfg @@ -202,11 +202,11 @@ def read( @dataclass(frozen=True) class _ObsBufferSlot: - observation: NpObs + observation: EnvObs[torch.Tensor] kernel_outputs: KernelOutputs @classmethod - def bind(cls, observation: NpObs, reward: torch.Tensor, terminated: torch.Tensor) -> _ObsBufferSlot: + def bind(cls, observation: EnvObs[torch.Tensor], reward: torch.Tensor, terminated: torch.Tensor) -> _ObsBufferSlot: value = observation.value if value is None: value = torch.empty((observation.policy.shape[0], 0), dtype=observation.policy.dtype) @@ -225,7 +225,7 @@ def bind(cls, observation: NpObs, reward: torch.Tensor, terminated: torch.Tensor class SwapObsBuffer: """Alternate two preallocated observation slots without copying. - Each slot binds an exposed ``NpObs`` to ``KernelOutputs`` that write + Each slot binds an exposed ``EnvObs`` to ``KernelOutputs`` that write into the same policy/value tensors. Swapping both together before the next kernel call keeps ``obs_t`` unchanged while a transition consumer retains it and the kernel produces ``obs_t+1`` in the other slot. @@ -241,8 +241,8 @@ class SwapObsBuffer: @classmethod def create( cls, - first: NpObs, - second: NpObs, + first: EnvObs[torch.Tensor], + second: EnvObs[torch.Tensor], reward: torch.Tensor, terminated: torch.Tensor, ) -> SwapObsBuffer: @@ -252,13 +252,13 @@ def create( def current(self) -> _ObsBufferSlot: return self._slots[self._index] - def swap(self) -> tuple[NpObs, KernelOutputs]: + def swap(self) -> tuple[EnvObs[torch.Tensor], KernelOutputs]: self._index = 1 - self._index slot = self.current return slot.observation, slot.kernel_outputs -class ManagerEnv(ArrayEnv[EnvCfgType]): +class ManagerEnv(ArrayEnv[EnvCfgType, torch.Tensor]): """NumPy environment whose task program is generated from manager config.""" def __init__( @@ -426,7 +426,7 @@ def _zeros_terminated(self) -> torch.Tensor: def _zeros_truncated(self) -> torch.Tensor: return torch.zeros((self.num_envs,), dtype=torch.bool) - def init_state(self) -> ArrayEnvState: + def init_state(self) -> ArrayEnvState[torch.Tensor]: startup_started = perf_counter() env_name = type(self).__name__ if self._task_program is None: @@ -488,7 +488,7 @@ def _execute_observe_kernel(self, inputs: tuple[ResolvedTree, ...]) -> None: assert self._kernel_outputs is not None self._task_program.observe_kernel(inputs, self._kernel_outputs) - def _capture_final_obs(self, env_ids: np.ndarray) -> FinalObs: + def _capture_final_obs(self, env_ids: np.ndarray) -> FinalObs[torch.Tensor]: """Observe pre-reset inputs after command on_transition. The compiled observe pass is batch-wide. Dedicated outputs preserve the @@ -501,7 +501,7 @@ def _capture_final_obs(self, env_ids: np.ndarray) -> FinalObs: obs = slot.observation return FinalObs( env_ids=env_ids, - obs=NpObs( + obs=EnvObs( policy=obs.policy[env_ids], value=None if obs.value is None else obs.value[env_ids], ), @@ -593,7 +593,7 @@ def observation_space(self) -> ObsSpace: def action_space(self) -> gym.spaces.Box: return self._action_space - def apply_action(self, actions: np.ndarray, state: ArrayEnvState) -> ArrayEnvState: + def apply_action(self, actions: np.ndarray, state: ArrayEnvState[torch.Tensor]) -> ArrayEnvState[torch.Tensor]: expected_shape = (self.num_envs, self.action_space.shape[0]) if actions.shape != expected_shape: raise ValueError(f"Expected action shape {expected_shape}, got {actions.shape}.") @@ -707,7 +707,7 @@ def _make_metrics_view(self) -> dict[str, Any]: ) return metrics - def compute_transition(self, state: ArrayEnvState) -> ArrayEnvState: + def compute_transition(self, state: ArrayEnvState[torch.Tensor]) -> ArrayEnvState[torch.Tensor]: with self.perf.scope("read_inputs"): self._refresh_sim_reads() with self.perf.scope("evaluate"): @@ -717,7 +717,7 @@ def compute_transition(self, state: ArrayEnvState) -> ArrayEnvState: command_term.on_transition() return state - def compute_observation(self, state: ArrayEnvState) -> ArrayEnvState: + def compute_observation(self, state: ArrayEnvState[torch.Tensor]) -> ArrayEnvState[torch.Tensor]: # During initialisation, ArrayEnv computes the first observation # before the manager's kernel outputs exist. The manager-specific # init_state() performs that first kernel pass itself. @@ -754,7 +754,7 @@ def observation_groups(self) -> dict[str, ObservationGroupEntry]: """Return the environment-local runtime observation groups.""" return self._observation_groups - def _make_kernel_buffers(self, state: ArrayEnvState) -> tuple[torch.Tensor, ...]: + def _make_kernel_buffers(self, state: ArrayEnvState[torch.Tensor]) -> tuple[torch.Tensor, ...]: assert self._compiled_manager_program is not None layout = self._compiled_manager_program.layout reward_terms = torch.empty((self.num_envs, len(layout.rewards)), dtype=torch.float32) diff --git a/motrix_env_core/tests/test_array_env.py b/motrix_env_core/tests/test_array_env.py index 387a0896..fc7b0d2d 100644 --- a/motrix_env_core/tests/test_array_env.py +++ b/motrix_env_core/tests/test_array_env.py @@ -7,7 +7,7 @@ import numpy as np import pytest -from motrix_env_core.array.env import ArrayEnv, ArrayEnvState, FinalObs, NpObs +from motrix_env_core.array.env import ArrayEnv, ArrayEnvState, EnvObs, FinalObs from motrix_env_core.base import EnvCfg, ObsSpace from motrix_env_core.config.scene import SceneCfg from motrix_env_core.direct.env import DirectEnv @@ -80,7 +80,7 @@ def _capture_final_obs(self, env_ids: np.ndarray) -> FinalObs: self.capture_rewards.append(self.state.reward[env_ids].copy()) policy = self.counters[env_ids, None].copy() value = np.concatenate((policy, -policy), axis=1) if self.has_value_observation else None - return FinalObs(env_ids=env_ids.copy(), obs=NpObs(policy=policy, value=value)) + return FinalObs(env_ids=env_ids.copy(), obs=EnvObs(policy=policy, value=value)) def _step(env: _CounterEnv) -> ArrayEnvState: @@ -138,7 +138,7 @@ def test_final_obs_contains_only_truncated_rows_before_auto_reset(with_value: bo final = state.final_obs assert isinstance(final, FinalObs) - assert isinstance(final.obs, NpObs) + assert isinstance(final.obs, EnvObs) np.testing.assert_array_equal(state.truncated, [True, False, True, False]) np.testing.assert_array_equal(state.terminated, [False, True, False, False]) np.testing.assert_array_equal(final.env_ids, [0, 2]) diff --git a/motrix_envs/src/motrix_envs/locomotion/quadruped/walk_np.py b/motrix_envs/src/motrix_envs/locomotion/quadruped/walk_np.py index c5a0c0eb..603ccc97 100644 --- a/motrix_envs/src/motrix_envs/locomotion/quadruped/walk_np.py +++ b/motrix_envs/src/motrix_envs/locomotion/quadruped/walk_np.py @@ -7,7 +7,7 @@ import numpy as np from motrix_robots.quadruped import QuadrupedRobotCfg -from motrix_env_core.array.env import ArrayEnvState, NpObs +from motrix_env_core.array.env import ArrayEnvState, EnvObs from motrix_env_core.base import ObsSpace from motrix_env_core.direct.env import DirectEnv from motrix_env_core.mdp.action_space import asymmetric_residual_action_space @@ -376,7 +376,7 @@ def compute_observation(self, state: ArrayEnvState) -> ArrayEnvState: f"expected policy[{expected_policy}]; value={value.shape}, " f"expected value[{expected_value}]" ) - return state.replace(obs=NpObs(policy=policy, value=value)) + return state.replace(obs=EnvObs(policy=policy, value=value)) def compute_transition(self, state: ArrayEnvState) -> ArrayEnvState: self.sim_data.execute() diff --git a/motrix_envs/src/motrix_envs/manipulation/franka_lift_cube/franka_lift_cube_np.py b/motrix_envs/src/motrix_envs/manipulation/franka_lift_cube/franka_lift_cube_np.py index fd9e9f77..cd04c2d2 100755 --- a/motrix_envs/src/motrix_envs/manipulation/franka_lift_cube/franka_lift_cube_np.py +++ b/motrix_envs/src/motrix_envs/manipulation/franka_lift_cube/franka_lift_cube_np.py @@ -5,7 +5,7 @@ import numpy as np from motrix_env_core import registry -from motrix_env_core.array.env import ArrayEnvState, NpObs +from motrix_env_core.array.env import ArrayEnvState, EnvObs from motrix_env_core.direct.env import DirectEnv from motrix_env_core.sim import ( BodyAngularVelocityWrite, @@ -154,7 +154,7 @@ def compute_observation(self, state: ArrayEnvState): # name is applied via setattr so static audits can tell this # sanctioned observation-stage write apart from the forbidden # transition-time writes. - setattr(state, "obs", NpObs(policy=obs.astype(np.float32))) + setattr(state, "obs", EnvObs(policy=obs.astype(np.float32))) return state def compute_transition(self, state: ArrayEnvState): diff --git a/motrix_envs/src/motrix_envs/manipulation/franka_open_cabinet/franka_open_cabinet_np.py b/motrix_envs/src/motrix_envs/manipulation/franka_open_cabinet/franka_open_cabinet_np.py index ce71ad77..12a76d1a 100755 --- a/motrix_envs/src/motrix_envs/manipulation/franka_open_cabinet/franka_open_cabinet_np.py +++ b/motrix_envs/src/motrix_envs/manipulation/franka_open_cabinet/franka_open_cabinet_np.py @@ -5,7 +5,7 @@ import numpy as np from motrix_env_core import registry -from motrix_env_core.array.env import ArrayEnvState, NpObs +from motrix_env_core.array.env import ArrayEnvState, EnvObs from motrix_env_core.direct.env import DirectEnv from motrix_env_core.math import quaternion from motrix_env_core.sim import ( @@ -199,7 +199,7 @@ def compute_observation(self, state: ArrayEnvState): # name is applied via setattr so static audits can tell this # sanctioned observation-stage write apart from the forbidden # transition-time writes. - setattr(state, "obs", NpObs(policy=np.clip(obs, -5, 5))) + setattr(state, "obs", EnvObs(policy=np.clip(obs, -5, 5))) return state def compute_transition(self, state: ArrayEnvState): diff --git a/motrix_envs/src/motrix_envs/manipulation/rm65_insert_peg/insert_peg_np.py b/motrix_envs/src/motrix_envs/manipulation/rm65_insert_peg/insert_peg_np.py index 89517fff..2b32b003 100644 --- a/motrix_envs/src/motrix_envs/manipulation/rm65_insert_peg/insert_peg_np.py +++ b/motrix_envs/src/motrix_envs/manipulation/rm65_insert_peg/insert_peg_np.py @@ -5,7 +5,7 @@ import numpy as np from motrix_env_core import registry -from motrix_env_core.array.env import ArrayEnvState, NpObs +from motrix_env_core.array.env import ArrayEnvState, EnvObs from motrix_env_core.direct.env import DirectEnv from motrix_env_core.math import quaternion from motrix_env_core.sim import ( @@ -329,7 +329,7 @@ def compute_observation(self, state: ArrayEnvState): # name is applied via setattr so static audits can tell this # sanctioned observation-stage write apart from the forbidden # transition-time writes. - setattr(state, "obs", NpObs(policy=obs.astype(np.float32))) + setattr(state, "obs", EnvObs(policy=obs.astype(np.float32))) return state def compute_transition(self, state: ArrayEnvState): diff --git a/motrix_envs/src/motrix_envs/manipulation/rm65_open_cabinet/rm65_open_cabinet_np.py b/motrix_envs/src/motrix_envs/manipulation/rm65_open_cabinet/rm65_open_cabinet_np.py index 5e4e77ad..b9c7cbb8 100755 --- a/motrix_envs/src/motrix_envs/manipulation/rm65_open_cabinet/rm65_open_cabinet_np.py +++ b/motrix_envs/src/motrix_envs/manipulation/rm65_open_cabinet/rm65_open_cabinet_np.py @@ -5,7 +5,7 @@ import numpy as np from motrix_env_core import registry -from motrix_env_core.array.env import ArrayEnvState, NpObs +from motrix_env_core.array.env import ArrayEnvState, EnvObs from motrix_env_core.direct.env import DirectEnv from motrix_env_core.math import quaternion from motrix_env_core.sim import ( @@ -542,7 +542,7 @@ def compute_observation(self, state: ArrayEnvState): # name is applied via setattr so static audits can tell this # sanctioned observation-stage write apart from the forbidden # transition-time writes. - setattr(state, "obs", NpObs(policy=np.clip(obs, -5, 5))) + setattr(state, "obs", EnvObs(policy=np.clip(obs, -5, 5))) return state def compute_transition(self, state: ArrayEnvState): diff --git a/motrix_envs/src/motrix_envs/manipulation/shadow_hand/shadow_hand_np.py b/motrix_envs/src/motrix_envs/manipulation/shadow_hand/shadow_hand_np.py index 587f0cf7..5f3e8beb 100644 --- a/motrix_envs/src/motrix_envs/manipulation/shadow_hand/shadow_hand_np.py +++ b/motrix_envs/src/motrix_envs/manipulation/shadow_hand/shadow_hand_np.py @@ -5,7 +5,7 @@ import numpy as np from motrix_env_core import registry -from motrix_env_core.array.env import ArrayEnvState, NpObs +from motrix_env_core.array.env import ArrayEnvState, EnvObs from motrix_env_core.direct.env import DirectEnv from motrix_env_core.math import quaternion, utils from motrix_env_core.sim import ( @@ -259,7 +259,7 @@ def compute_observation(self, state: ArrayEnvState): # name is applied via setattr so static audits can tell this # sanctioned observation-stage write apart from the forbidden # transition-time writes. - setattr(state, "obs", NpObs(policy=obs)) + setattr(state, "obs", EnvObs(policy=obs)) return state def compute_transition(self, state: ArrayEnvState): diff --git a/motrix_envs/tests/test_basic_direct_contract.py b/motrix_envs/tests/test_basic_direct_contract.py index 37d2039a..d5f85418 100644 --- a/motrix_envs/tests/test_basic_direct_contract.py +++ b/motrix_envs/tests/test_basic_direct_contract.py @@ -19,7 +19,7 @@ import pytest from motrix_env_core import registry -from motrix_env_core.array.env import ArrayEnvState, NpObs +from motrix_env_core.array.env import ArrayEnvState, EnvObs # One registered name per basic DirectEnv family. _BASIC_ENVS = [ @@ -64,7 +64,7 @@ def recording_execute(env_ids=None): def _assert_finite_state(state: ArrayEnvState, num_envs: int, obs_dim: Iterable[int]) -> None: - assert isinstance(state.obs, NpObs) + assert isinstance(state.obs, EnvObs) assert state.obs.policy.shape == (num_envs, *obs_dim) assert not np.isnan(state.obs.policy).any() assert state.reward.shape == (num_envs,) @@ -134,7 +134,7 @@ def test_compute_observation_reads_cache_without_executing(env_name: str) -> Non # Observation must be served purely from cache. assert calls == [] - obs = observed.obs.policy if isinstance(observed.obs, NpObs) else observed.obs + obs = observed.obs.policy if isinstance(observed.obs, EnvObs) else observed.obs assert obs.shape == (num_envs, *env.observation_space.shape) assert not np.isnan(obs).any() # Reward / termination / reward_terms must not be modified by observation. diff --git a/motrix_envs/tests/test_basic_env_lifecycle.py b/motrix_envs/tests/test_basic_env_lifecycle.py index c489383b..eb7cf600 100644 --- a/motrix_envs/tests/test_basic_env_lifecycle.py +++ b/motrix_envs/tests/test_basic_env_lifecycle.py @@ -5,7 +5,7 @@ import pytest from motrix_env_core import registry -from motrix_env_core.array.env import NpObs +from motrix_env_core.array.env import EnvObs @pytest.mark.parametrize("num_envs", [1, 2]) @@ -22,7 +22,7 @@ def test_basic_environment_step_preserves_numpy_observation_contract(env_name: s state = env.step(action) assert state is not None - assert isinstance(state.obs, NpObs) + assert isinstance(state.obs, EnvObs) assert state.obs.policy.shape == (num_envs, *env.observation_space.shape) if env_name == "franka-open-cabinet": # Episode-scoped gripper state lives on the env instance, not in info. diff --git a/motrix_envs/tests/test_manipulation_direct_contract.py b/motrix_envs/tests/test_manipulation_direct_contract.py index 78b8a0e7..ece5d455 100644 --- a/motrix_envs/tests/test_manipulation_direct_contract.py +++ b/motrix_envs/tests/test_manipulation_direct_contract.py @@ -19,7 +19,7 @@ import pytest from motrix_env_core import registry -from motrix_env_core.array.env import ArrayEnvState, NpObs +from motrix_env_core.array.env import ArrayEnvState, EnvObs _MANIPULATION_ENVS = [ "franka-lift-cube", @@ -51,7 +51,7 @@ def recording_execute(env_ids=None): def _assert_finite_state(state: ArrayEnvState, num_envs: int, obs_dim: Iterable[int]) -> None: - assert isinstance(state.obs, NpObs) + assert isinstance(state.obs, EnvObs) assert state.obs.policy.shape == (num_envs, *obs_dim) assert not np.isnan(state.obs.policy).any() assert state.reward.shape == (num_envs,) @@ -126,9 +126,9 @@ def test_compute_observation_reads_cache_without_executing(env_name: str) -> Non # Observation must be served purely from cache. assert calls == [] - # step() normalizes obs to NpObs; a raw compute_observation call may return + # step() normalizes obs to EnvObs; a raw compute_observation call may return # either the dataclass or the bare policy array. - obs = observed.obs.policy if isinstance(observed.obs, NpObs) else observed.obs + obs = observed.obs.policy if isinstance(observed.obs, EnvObs) else observed.obs assert obs.shape == (num_envs, *env.observation_space.shape) assert not np.isnan(obs).any() # Reward / termination / reward terms must not be modified by observation. diff --git a/motrix_rl_builtin/tests/test_builtin_sim_backend.py b/motrix_rl_builtin/tests/test_builtin_sim_backend.py index 036b0f06..0180bb44 100644 --- a/motrix_rl_builtin/tests/test_builtin_sim_backend.py +++ b/motrix_rl_builtin/tests/test_builtin_sim_backend.py @@ -17,7 +17,7 @@ from motrix_rl_builtin.fastsac.async_impl.worker import actor_param_numel, run_collector_process from motrix_rl_builtin.fastsac.wrap import FastSacEnvWrap, FastSacStepResult -from motrix_env_core.array.env import ArrayEnvState, NpObs +from motrix_env_core.array.env import ArrayEnvState, EnvObs from motrix_env_core.base import EnvCfg from motrix_env_core.config.scene import SceneCfg from motrix_env_core.direct.env import DirectEnv @@ -48,7 +48,7 @@ def _make_env(sim_backend: str): if sim_backend == "np": env = Mock(spec=DirectEnv) state = ArrayEnvState( - obs=NpObs(np.zeros((_NUM_ENVS, _OBS_DIM), dtype=np.float32)), + obs=EnvObs(np.zeros((_NUM_ENVS, _OBS_DIM), dtype=np.float32)), reward=np.ones(_NUM_ENVS, dtype=np.float32), terminated=np.zeros(_NUM_ENVS, dtype=bool), truncated=np.zeros(_NUM_ENVS, dtype=bool), @@ -169,7 +169,7 @@ def compute_transition(self, state): def _make_state(self, value: float) -> ArrayEnvState: return ArrayEnvState( - obs=NpObs(np.full((self.num_envs, _OBS_DIM), value, dtype=np.float32)), + obs=EnvObs(np.full((self.num_envs, _OBS_DIM), value, dtype=np.float32)), reward=np.full(self.num_envs, value, dtype=np.float32), terminated=np.zeros(self.num_envs, dtype=bool), truncated=np.zeros(self.num_envs, dtype=bool), diff --git a/motrix_rl_builtin/tests/test_fastsac_collector.py b/motrix_rl_builtin/tests/test_fastsac_collector.py index 6d396298..bd096426 100644 --- a/motrix_rl_builtin/tests/test_fastsac_collector.py +++ b/motrix_rl_builtin/tests/test_fastsac_collector.py @@ -20,7 +20,7 @@ from motrix_rl_builtin.fastsac.networks import Actor from motrix_rl_builtin.fastsac.wrap import FastSacStepResult -from motrix_env_core.array.env import FinalObs, NpObs +from motrix_env_core.array.env import EnvObs, FinalObs from motrix_env_core.perf import Perf _NUM_ENVS = 8 @@ -396,7 +396,7 @@ def step(self, actions): reward=torch.tensor([1.0, 2.0]), terminated=torch.zeros(2, dtype=torch.bool), truncated=torch.tensor([False, True]), - final_obs=FinalObs(np.array([1]), NpObs(policy=np.full((1, 3), 7.0, dtype=np.float32))), + final_obs=FinalObs(np.array([1]), EnvObs(policy=np.full((1, 3), 7.0, dtype=np.float32))), info={"Reward": {"alive": np.ones(2)}, "metrics": {"error": np.ones(2)}}, ) return self.outputs diff --git a/motrix_rl_builtin/tests/test_fastsac_final_obs.py b/motrix_rl_builtin/tests/test_fastsac_final_obs.py index 09047991..05002a02 100644 --- a/motrix_rl_builtin/tests/test_fastsac_final_obs.py +++ b/motrix_rl_builtin/tests/test_fastsac_final_obs.py @@ -27,13 +27,13 @@ from motrix_rl_builtin.fastsac.buffer import SimpleReplayBuffer from motrix_rl_builtin.fastsac.wrap_np import FastSacNpEnvWrap -from motrix_env_core.array.env import ArrayEnvState, FinalObs, NpObs +from motrix_env_core.array.env import ArrayEnvState, EnvObs, FinalObs def test_numpy_adapter_reads_final_result_without_consuming_environment_state(): - final = FinalObs(np.array([0]), NpObs(np.array([[7.0]], dtype=np.float32))) + final = FinalObs(np.array([0]), EnvObs(np.array([[7.0]], dtype=np.float32))) state = ArrayEnvState( - obs=NpObs(np.zeros((1, 1), dtype=np.float32)), + obs=EnvObs(np.zeros((1, 1), dtype=np.float32)), reward=np.zeros(1, dtype=np.float32), terminated=np.zeros(1, dtype=bool), truncated=np.ones(1, dtype=bool), diff --git a/motrix_rl_rslrl/src/motrix_rl_rslrl/torch/wrap_np.py b/motrix_rl_rslrl/src/motrix_rl_rslrl/torch/wrap_np.py index 4973d757..2663f465 100644 --- a/motrix_rl_rslrl/src/motrix_rl_rslrl/torch/wrap_np.py +++ b/motrix_rl_rslrl/src/motrix_rl_rslrl/torch/wrap_np.py @@ -8,7 +8,7 @@ from rsl_rl.env.vec_env import VecEnv from tensordict import TensorDict -from motrix_env_core.array.env import ArrayEnv, NpObs +from motrix_env_core.array.env import ArrayEnv, EnvObs from motrix_env_core.renderer import RenderConfig, create_renderer @@ -147,7 +147,7 @@ def get_observations(self) -> TensorDict: return self._to_tensordict(self._state.obs) - def _to_tensordict(self, obs: NpObs) -> TensorDict: + def _to_tensordict(self, obs: EnvObs) -> TensorDict: tensors = {"policy": _as_torch(obs.policy, self._device)} if obs.value is not None: tensors["value"] = _as_torch(obs.value, self._device) diff --git a/motrix_rl_rslrl/tests/test_rslrl_sim_backend.py b/motrix_rl_rslrl/tests/test_rslrl_sim_backend.py index 6acca46f..fe69d2d4 100644 --- a/motrix_rl_rslrl/tests/test_rslrl_sim_backend.py +++ b/motrix_rl_rslrl/tests/test_rslrl_sim_backend.py @@ -10,7 +10,7 @@ import pytest import torch -from motrix_env_core.array.env import ArrayEnvState, NpObs +from motrix_env_core.array.env import ArrayEnvState, EnvObs from motrix_env_core.direct.env import DirectEnv from motrix_env_motrixsim.torch_env import TorchEnv, TorchEnvState, TorchObs @@ -35,7 +35,7 @@ def _make_env(sim_backend: str): if sim_backend == "np": env = Mock(spec=DirectEnv) state = ArrayEnvState( - obs=NpObs(np.zeros((_NUM_ENVS, _OBS_DIM), dtype=np.float32)), + obs=EnvObs(np.zeros((_NUM_ENVS, _OBS_DIM), dtype=np.float32)), reward=np.ones(_NUM_ENVS, dtype=np.float32), terminated=np.zeros(_NUM_ENVS, dtype=bool), truncated=np.zeros(_NUM_ENVS, dtype=bool), diff --git a/motrix_rl_skrl/tests/test_skrl_sim_backend.py b/motrix_rl_skrl/tests/test_skrl_sim_backend.py index 81457de1..149f8d4b 100644 --- a/motrix_rl_skrl/tests/test_skrl_sim_backend.py +++ b/motrix_rl_skrl/tests/test_skrl_sim_backend.py @@ -10,7 +10,7 @@ import pytest import torch -from motrix_env_core.array.env import ArrayEnvState, NpObs +from motrix_env_core.array.env import ArrayEnvState, EnvObs from motrix_env_core.direct.env import DirectEnv from motrix_env_motrixsim.torch_env import TorchEnv, TorchEnvState, TorchObs @@ -37,7 +37,7 @@ def _make_env(sim_backend: str): if sim_backend == "np": env = Mock(spec=DirectEnv) state = ArrayEnvState( - obs=NpObs(np.zeros((_NUM_ENVS, _OBS_DIM), dtype=np.float32)), + obs=EnvObs(np.zeros((_NUM_ENVS, _OBS_DIM), dtype=np.float32)), reward=np.ones(_NUM_ENVS, dtype=np.float32), terminated=np.zeros(_NUM_ENVS, dtype=bool), truncated=np.zeros(_NUM_ENVS, dtype=bool), diff --git a/wiki/design/array-env-final-observation.md b/wiki/design/array-env-final-observation.md index a787eff9..ff2e9540 100644 --- a/wiki/design/array-env-final-observation.md +++ b/wiki/design/array-env-final-observation.md @@ -6,7 +6,7 @@ ArrayEnv 在时间截断自动 reset 前捕获当前 transition 的终态观测 ## 结果结构与生命周期 -`FinalObs` 为普通 dataclass,包含 `env_ids: np.ndarray` 和 `obs: NpObs`。观测行顺序与 env_ids 一致,policy/value 保持原始观测结构;没有独立 value 时仍为 None,消费者使用 `value_or_policy`。行数组为独立快照,可跨 step 保留;不需要重复复制 policy 作为 value。消费者不得修改共享结果。 +`FinalObs` 为普通 dataclass,包含 `env_ids: np.ndarray` 和 `obs: EnvObs`。观测行顺序与 env_ids 一致,policy/value 保持原始观测结构;没有独立 value 时仍为 None,消费者使用 `value_or_policy`。行数组为独立快照,可跨 step 保留;不需要重复复制 policy 作为 value。消费者不得修改共享结果。 `ArrayEnvState.final_obs` 为可选字段。初始化为 None,每次 step 在推进环境前通过 replace 清空当前结果;此前返回的 state 持有的终态快照不因此清空。step 在 compute_transition、episode_steps 和 truncation 更新后、任何 auto-reset 写入前调用 `_capture_final_obs(env_ids)`。仅有 truncated 行时调用;没有截断或实现不支持捕获时为 None。读取无消费副作用。 diff --git a/wiki/design/manager/runtime.md b/wiki/design/manager/runtime.md index 124a474c..11fad5ab 100644 --- a/wiki/design/manager/runtime.md +++ b/wiki/design/manager/runtime.md @@ -170,7 +170,7 @@ Observation terms use the host-side ``ObservationTerm(size, dispatch, *args)`` f ## 5. 执行与性能边界 -- 只有纯数组计算进入 term dispatch;不得传入 `SceneModel`、`SceneData`、`ArrayEnvState`、`NpObs`、config 或 `info`; +- 只有纯数组计算进入 term dispatch;不得传入 `SceneModel`、`SceneData`、`ArrayEnvState`、`EnvObs`、config 或 `info`; - term 一次调用处理整批张量,写入预分配 output 和 term-owned buffers; - Torch 执行器默认以 `torch.compile` 编译整段 term 调度(可回退 eager),不自动转换任意 Python/NumPy 环境方法; - 单个环境实例不修改进程级 Torch 线程数,线程策略由训练 runtime 管理; From 37753f8a997648b4b03345f71717902a7e78a374 Mon Sep 17 00:00:00 2001 From: zilch <147668916@qq.com> Date: Fri, 6 Nov 2026 20:25:40 +0800 Subject: [PATCH 31/41] refactor(kernel_data)!: collapse the dual tree schemas into one walk 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. --- .../manager/kernel_data/__init__.py | 16 +- .../manager/kernel_data/canonical.py | 12 +- .../manager/kernel_data/lowering.py | 261 ------------------ .../manager/kernel_data/tree.py | 87 +++--- .../src/motrix_env_core/manager/resolver.py | 95 +++++-- motrix_env_core/tests/test_kernel_data.py | 39 +-- 6 files changed, 123 insertions(+), 387 deletions(-) delete mode 100644 motrix_env_core/src/motrix_env_core/manager/kernel_data/lowering.py diff --git a/motrix_env_core/src/motrix_env_core/manager/kernel_data/__init__.py b/motrix_env_core/src/motrix_env_core/manager/kernel_data/__init__.py index d6ef1772..0fa68723 100644 --- a/motrix_env_core/src/motrix_env_core/manager/kernel_data/__init__.py +++ b/motrix_env_core/src/motrix_env_core/manager/kernel_data/__init__.py @@ -1,14 +1,6 @@ # Copyright Motphys Technology Co., Ltd. 2025, 2026 from motrix_env_core.manager.kernel_data.canonical import canonicalize_kernel_data -from motrix_env_core.manager.kernel_data.lowering import ( - KernelDataLayout, - KernelDataLowering, - KernelLeafLayout, - KernelMapLayout, - KernelRecordLayout, - iter_layout_leaves, -) from motrix_env_core.manager.kernel_data.map import Map from motrix_env_core.manager.kernel_data.resolved import ResolvedLeaf, ResolvedTree from motrix_env_core.manager.kernel_data.tree import ( @@ -22,21 +14,17 @@ is_kernel_data, is_kernel_data_type, iter_tree_leaves, + leaf_value_type, unflatten_kernel_data, ) __all__ = [ "BaseDef", - "KernelDataLayout", "Map", - "KernelDataLowering", - "KernelLeafLayout", "KernelMapDef", "FieldDef", "LeafDef", "MapEntryDef", - "KernelMapLayout", - "KernelRecordLayout", "TreeClassDef", "ResolvedLeaf", "ResolvedTree", @@ -44,7 +32,7 @@ "flatten_kernel_data", "is_kernel_data", "is_kernel_data_type", - "iter_layout_leaves", "iter_tree_leaves", + "leaf_value_type", "unflatten_kernel_data", ] diff --git a/motrix_env_core/src/motrix_env_core/manager/kernel_data/canonical.py b/motrix_env_core/src/motrix_env_core/manager/kernel_data/canonical.py index 541be3f7..fb1a2e8f 100644 --- a/motrix_env_core/src/motrix_env_core/manager/kernel_data/canonical.py +++ b/motrix_env_core/src/motrix_env_core/manager/kernel_data/canonical.py @@ -1,15 +1,15 @@ -# Copyright Motrix Technology Co., Ltd. 2025, 2026 -# SPDX-License-Identifier: Apache-2.0 +# Copyright Motphys Technology Co., Ltd. 2025, 2026 from typing import TypeVar, cast import numpy as np import torch -from motrix_env_core.manager.kernel_data.lowering import KernelDataLowering from motrix_env_core.manager.kernel_data.tree import ( flatten_kernel_data, is_kernel_data, + iter_tree_leaves, + leaf_value_type, unflatten_kernel_data, ) @@ -21,7 +21,11 @@ def canonicalize_kernel_data(value: KernelDataType, *, context: str) -> KernelDa if not is_kernel_data(value): raise TypeError(f"{context} must return a concrete kernel-data dataclass, got {type(value).__name__}.") leaves, tree_def = flatten_kernel_data(value) - KernelDataLowering().lower(tree_def, context=context) + for leaf_def in iter_tree_leaves(tree_def): + if leaf_value_type(leaf_def.annotation) is None: + raise TypeError( + f"{context} KernelData {'.'.join(leaf_def.path)} expected a tensor, an ndarray, or a numeric scalar." + ) copied = tuple( leaf.clone() if isinstance(leaf, torch.Tensor) else leaf.copy() if isinstance(leaf, np.ndarray) else leaf for leaf in leaves diff --git a/motrix_env_core/src/motrix_env_core/manager/kernel_data/lowering.py b/motrix_env_core/src/motrix_env_core/manager/kernel_data/lowering.py deleted file mode 100644 index 9d83e7ea..00000000 --- a/motrix_env_core/src/motrix_env_core/manager/kernel_data/lowering.py +++ /dev/null @@ -1,261 +0,0 @@ -# Copyright Motphys Technology Co., Ltd. 2025, 2026 -# SPDX-License-Identifier: Apache-2.0 - -from __future__ import annotations - -import hashlib -from dataclasses import dataclass -from typing import Any - -import numpy as np -import torch - -from motrix_env_core.manager.fingerprint import function_fingerprint -from motrix_env_core.manager.kernel_data.methods import dispatch_methods -from motrix_env_core.manager.kernel_data.tree import KernelMapDef, LeafDef, TreeClassDef - -_LOWERING_SCHEMA_VERSION = 3 -_LAYOUT_CACHE: dict[ - tuple[str, bool], - KernelRecordLayout, -] = {} - - -@dataclass(frozen=True) -class KernelLeafLayout: - """Lowered ABI description of one logical kernel-data leaf. - - Attributes: - tree_def: Logical leaf definition that owns the path and annotation. - value_type: Concrete Python/NumPy type accepted for the leaf value. - slot_index: Zero-based declaration index relative to the root layout. - """ - - tree_def: LeafDef - value_type: type - slot_index: int - - @property - def path(self) -> tuple[str, ...]: - """Return the logical path inherited from ``tree_def``.""" - return self.tree_def.path - - -@dataclass(frozen=True) -class KernelFieldLayout: - """Named lowered child of a ``KernelRecordLayout``. - - Attributes: - name: Python dataclass field name on the logical record. - child: Lowered layout for the field value. - """ - - name: str - child: KernelDataLayout - - -@dataclass(frozen=True) -class KernelMapEntryLayout: - """Named lowered child of a ``KernelMapLayout``. - - Attributes: - key: Immutable string key used for compile-time map lookup. - child: Lowered layout for the value associated with ``key``. - """ - - key: str - child: KernelDataLayout - - -@dataclass(frozen=True) -class KernelMapLayout: - """Lowered schema for one heterogeneous ``Map``. - - Attributes: - tree_def: Value-derived logical map definition. - entries: Lowered children in the same order as the logical map entries. - fingerprint: Deterministic identity including keys and child schemas. - """ - - tree_def: KernelMapDef - entries: tuple[KernelMapEntryLayout, ...] - fingerprint: str - - -@dataclass(frozen=True) -class KernelRecordLayout: - """Lowered schema for one concrete kernel-data class tree. - - Attributes: - tree_def: Logical record definition being lowered. - fields: Lowered children in logical dataclass declaration order. - fingerprint: Deterministic identity including child types and - dispatch-method code fingerprints. - """ - - tree_def: TreeClassDef - fields: tuple[KernelFieldLayout, ...] - fingerprint: str - - @property - def logical_type(self) -> type: - """Return the host kernel-data type represented by this layout.""" - return self.tree_def.logical_type - - -KernelDataLayout = KernelRecordLayout | KernelMapLayout | KernelLeafLayout - - -def iter_layout_leaves(layout: KernelDataLayout) -> tuple[KernelLeafLayout, ...]: - """Return lowered leaves in stable logical declaration order.""" - if isinstance(layout, KernelLeafLayout): - return (layout,) - if isinstance(layout, KernelMapLayout): - return tuple(leaf for entry in layout.entries for leaf in iter_layout_leaves(entry.child)) - assert isinstance(layout, KernelRecordLayout) - return tuple(leaf for field in layout.fields for leaf in iter_layout_leaves(field.child)) - - -def _tree_method_fingerprint(tree_def: TreeClassDef | KernelMapDef | LeafDef) -> Any: - """Include nested records and map values before consulting the layout cache.""" - if isinstance(tree_def, LeafDef): - return () - if isinstance(tree_def, KernelMapDef): - return tuple((entry.key, _tree_method_fingerprint(entry.tree_def)) for entry in tree_def.entries) - return ( - tuple( - (name, function_fingerprint(method)) - for name, method in sorted(dispatch_methods(tree_def.logical_type).items()) - ), - tuple((field.name, _tree_method_fingerprint(field.tree_def)) for field in tree_def.fields), - ) - - -class KernelDataLowering: - """Validate one logical ``TreeClassDef`` and derive its leaf layout. - - The lowering checks that every leaf annotation is a supported - ndarray/scalar type and produces a deterministic structural fingerprint - (field names, types, and dispatch-method code) used for - resolved-program invalidation. - """ - - def lower( - self, - tree_def: TreeClassDef, - *, - context: str, - ) -> KernelRecordLayout: - method_key = hashlib.sha256(repr(_tree_method_fingerprint(tree_def)).encode()).hexdigest() - cache_key = f"{tree_def.fingerprint}:{method_key}" - cached = _LAYOUT_CACHE.get(cache_key) - if cached is not None: - return cached - layout, _ = self._lower_record(tree_def, 0, context) - _LAYOUT_CACHE[cache_key] = layout - return layout - - def _lower_record( - self, - tree_def: TreeClassDef, - slot_index: int, - context: str, - ) -> tuple[KernelRecordLayout, int]: - field_layouts: list[KernelFieldLayout] = [] - fingerprint_fields: list[Any] = [] - for field in tree_def.fields: - child: KernelDataLayout - if isinstance(field.tree_def, TreeClassDef): - child, slot_index = self._lower_record(field.tree_def, slot_index, context) - elif isinstance(field.tree_def, KernelMapDef): - child, slot_index = self._lower_map(field.tree_def, slot_index, context) - else: - assert isinstance(field.tree_def, LeafDef) - value_type = self._leaf_type(field.tree_def.annotation) - if value_type is None: - raise TypeError( - f"{context} KernelData {tree_def.logical_type.__name__}." - f"{'.'.join(field.tree_def.path)} expected a tensor, an ndarray, or a numeric scalar." - ) - child = KernelLeafLayout(field.tree_def, value_type, slot_index) - slot_index += 1 - field_layouts.append(KernelFieldLayout(field.name, child)) - fingerprint_fields.append((field.name, self._fingerprint_part(child))) - identity = f"{tree_def.logical_type.__module__}.{tree_def.logical_type.__qualname__}" - method_fingerprints = tuple( - (name, function_fingerprint(method)) - for name, method in sorted(dispatch_methods(tree_def.logical_type).items()) - ) - raw_fingerprint = repr( - (_LOWERING_SCHEMA_VERSION, tree_def.fingerprint, identity, tuple(fingerprint_fields), method_fingerprints) - ) - fingerprint = hashlib.sha256(raw_fingerprint.encode()).hexdigest() - return KernelRecordLayout(tree_def, tuple(field_layouts), fingerprint), slot_index - - def _lower_map( - self, - tree_def: KernelMapDef, - slot_index: int, - context: str, - ) -> tuple[KernelMapLayout, int]: - entry_layouts: list[KernelMapEntryLayout] = [] - fingerprint_entries: list[Any] = [] - for entry in tree_def.entries: - if isinstance(entry.tree_def, TreeClassDef): - child, slot_index = self._lower_record(entry.tree_def, slot_index, context) - elif isinstance(entry.tree_def, KernelMapDef): - child, slot_index = self._lower_map(entry.tree_def, slot_index, context) - else: - assert isinstance(entry.tree_def, LeafDef) - value_type = self._leaf_type(entry.tree_def.annotation) - if value_type is None: - raise TypeError( - f"{context} Map {'.'.join(entry.tree_def.path)} expected a tensor, an ndarray, or a numeric " - "scalar." - ) - child = KernelLeafLayout(entry.tree_def, value_type, slot_index) - slot_index += 1 - entry_layouts.append(KernelMapEntryLayout(entry.key, child)) - fingerprint_entries.append((entry.key, self._fingerprint_part(child))) - raw_fingerprint = repr((_LOWERING_SCHEMA_VERSION, tree_def.fingerprint, tuple(fingerprint_entries))) - fingerprint = hashlib.sha256(raw_fingerprint.encode()).hexdigest() - return KernelMapLayout(tree_def, tuple(entry_layouts), fingerprint), slot_index - - @staticmethod - def _fingerprint_part(layout: KernelDataLayout) -> Any: - if isinstance(layout, KernelLeafLayout): - return ( - "leaf", - layout.path, - f"{layout.value_type.__module__}.{layout.value_type.__qualname__}", - layout.slot_index, - ) - if isinstance(layout, KernelMapLayout): - return ( - "map", - layout.tree_def.path, - layout.fingerprint, - tuple((entry.key, KernelDataLowering._fingerprint_part(entry.child)) for entry in layout.entries), - ) - assert isinstance(layout, KernelRecordLayout) - return ( - "record", - f"{layout.logical_type.__module__}.{layout.logical_type.__qualname__}", - layout.fingerprint, - tuple((field.name, KernelDataLowering._fingerprint_part(field.child)) for field in layout.fields), - ) - - @staticmethod - def _leaf_type(annotation: Any) -> type[Any] | None: - if annotation in (np.ndarray, torch.Tensor): - return annotation - if annotation in {bool, int, float}: - return annotation - if isinstance(annotation, type) and issubclass(annotation, np.generic): - try: - dtype = np.dtype(annotation) - except TypeError: - return None - if dtype.type is annotation and dtype.kind in "biuf": - return annotation - return None diff --git a/motrix_env_core/src/motrix_env_core/manager/kernel_data/tree.py b/motrix_env_core/src/motrix_env_core/manager/kernel_data/tree.py index 3f782194..b2f95391 100644 --- a/motrix_env_core/src/motrix_env_core/manager/kernel_data/tree.py +++ b/motrix_env_core/src/motrix_env_core/manager/kernel_data/tree.py @@ -5,7 +5,7 @@ import hashlib from dataclasses import dataclass, fields, is_dataclass -from typing import Annotated, Any, get_args, get_origin, get_type_hints +from typing import Any, get_origin, get_type_hints import numpy as np import torch @@ -13,8 +13,7 @@ from motrix_env_core.manager.kernel_data.map import Map from motrix_env_core.sim.read import SimDataQuery -_TREE_SCHEMA_VERSION = 2 -_TREE_DEF_CACHE: dict[type, TreeClassDef] = {} +_TREE_SCHEMA_VERSION = 3 def is_kernel_data_type(value_type: object) -> bool: @@ -35,6 +34,22 @@ def is_kernel_data(value: object) -> bool: ) +def leaf_value_type(annotation: Any) -> type | None: + """Return the concrete value type accepted for one leaf annotation.""" + if annotation in (np.ndarray, torch.Tensor): + return annotation + if annotation in {bool, int, float}: + return annotation + if isinstance(annotation, type) and issubclass(annotation, np.generic): + try: + dtype = np.dtype(annotation) + except TypeError: + return None + if dtype.type is annotation and dtype.kind in "biuf": + return annotation + return None + + class BaseDef: """Base class for nodes contained in a ``TreeClassDef`` schema tree.""" @@ -45,7 +60,7 @@ class LeafDef(BaseDef): Attributes: path: Field and map-key components from the tree root to this leaf. - annotation: Declared leaf annotation, including any ``Annotated`` metadata. + annotation: Declared leaf annotation. """ path: tuple[str, ...] @@ -53,7 +68,7 @@ class LeafDef(BaseDef): @dataclass(frozen=True) -class FieldDef: +class FieldDef(BaseDef): """Named child of a ``TreeClassDef``. Attributes: @@ -92,7 +107,7 @@ def unflatten(self, leaves: tuple[Any, ...]) -> Any: @dataclass(frozen=True) -class MapEntryDef: +class MapEntryDef(BaseDef): """Named child of a value-dependent ``KernelMapDef``. Attributes: @@ -123,13 +138,7 @@ def flatten_kernel_data(value: Any) -> tuple[tuple[Any, ...], TreeClassDef]: """Return zero-copy leaves and the concrete class tree definition for one kernel-data value.""" if not is_kernel_data(value): raise TypeError(f"KernelData tree root must be a dataclass value, got {type(value).__name__}.") - tree_def, dynamic = _build_class_def(type(value), value, (), ()) - if not dynamic: - cached = _TREE_DEF_CACHE.get(type(value)) - if cached is None: - _TREE_DEF_CACHE[type(value)] = tree_def - else: - tree_def = cached + tree_def = _build_class_def(type(value), value, (), ()) return _flatten(value, tree_def), tree_def @@ -197,44 +206,26 @@ def _build_class_def( value: Any | None, path: tuple[str, ...], stack: tuple[type, ...], -) -> tuple[TreeClassDef, bool]: +) -> TreeClassDef: if logical_type in stack: chain = " -> ".join(item.__name__ for item in (*stack, logical_type)) raise TypeError(f"KernelData cycle detected: {chain}.") annotations = get_type_hints(logical_type, include_extras=True) field_defs: list[FieldDef] = [] fingerprint_fields: list[Any] = [] - dynamic = False for data_field in fields(logical_type): annotation = annotations[data_field.name] field_path = (*path, data_field.name) - node_annotation = _node_annotation(annotation) field_value = getattr(value, data_field.name) if value is not None else None child: BaseDef - if node_annotation is Map or get_origin(node_annotation) is Map: - if get_origin(annotation) is Annotated: - raise TypeError( - f"KernelData {logical_type.__name__}.{'.'.join(field_path)} map fields cannot declare metadata." - ) + if annotation is Map or get_origin(annotation) is Map: if not isinstance(field_value, Map): raise TypeError( f"KernelData {logical_type.__name__}.{'.'.join(field_path)} requires a concrete Map value." ) - child, _ = _build_map_def(field_value, field_path, (*stack, logical_type)) - dynamic = True - elif is_kernel_data_type(node_annotation): - if get_origin(annotation) is Annotated: - raise TypeError( - f"KernelData {logical_type.__name__}.{'.'.join(field_path)} class fields cannot declare " - "leaf metadata." - ) - child, child_dynamic = _build_class_def( - node_annotation, - field_value, - field_path, - (*stack, logical_type), - ) - dynamic = dynamic or child_dynamic + child = _build_map_def(field_value, field_path, (*stack, logical_type)) + elif is_kernel_data_type(annotation): + child = _build_class_def(annotation, field_value, field_path, (*stack, logical_type)) else: child = LeafDef(field_path, annotation) field_defs.append(FieldDef(data_field.name, child)) @@ -242,29 +233,29 @@ def _build_class_def( identity = f"{logical_type.__module__}.{logical_type.__qualname__}" raw_fingerprint = repr((_TREE_SCHEMA_VERSION, identity, tuple(fingerprint_fields))) fingerprint = hashlib.sha256(raw_fingerprint.encode()).hexdigest() - return TreeClassDef(logical_type, tuple(field_defs), fingerprint), dynamic + return TreeClassDef(logical_type, tuple(field_defs), fingerprint) def _build_map_def( value: Map, path: tuple[str, ...], stack: tuple[type, ...], -) -> tuple[KernelMapDef, bool]: +) -> KernelMapDef: entries: list[MapEntryDef] = [] fingerprint_entries: list[Any] = [] for key, entry_value in zip(value.keys_tuple, value.values_tuple, strict=True): entry_path = (*path, key) if isinstance(entry_value, Map): - child, _ = _build_map_def(entry_value, entry_path, stack) + child = _build_map_def(entry_value, entry_path, stack) elif is_kernel_data(entry_value): - child, _ = _build_class_def(type(entry_value), entry_value, entry_path, stack) + child = _build_class_def(type(entry_value), entry_value, entry_path, stack) else: child = LeafDef(entry_path, type(entry_value)) entries.append(MapEntryDef(key, child)) fingerprint_entries.append((key, _tree_fingerprint_part(child))) raw_fingerprint = repr((_TREE_SCHEMA_VERSION, "map", path, tuple(fingerprint_entries))) fingerprint = hashlib.sha256(raw_fingerprint.encode()).hexdigest() - return KernelMapDef(path, tuple(entries), fingerprint), True + return KernelMapDef(path, tuple(entries), fingerprint) def _tree_fingerprint_part(tree_def: BaseDef) -> Any: @@ -285,20 +276,6 @@ def _tree_fingerprint_part(tree_def: BaseDef) -> Any: def _annotation_identity(annotation: Any) -> Any: - if get_origin(annotation) is Annotated: - value_type, *metadata = get_args(annotation) - return "annotated", _annotation_identity(value_type), tuple(_metadata_identity(value) for value in metadata) if isinstance(annotation, type): return f"{annotation.__module__}.{annotation.__qualname__}" return repr(annotation) - - -def _metadata_identity(value: Any) -> Any: - value_type = type(value) - if hasattr(value, "value"): - return f"{value_type.__module__}.{value_type.__qualname__}", value.value - return repr(value) - - -def _node_annotation(annotation: Any) -> Any: - return get_args(annotation)[0] if get_origin(annotation) is Annotated else annotation diff --git a/motrix_env_core/src/motrix_env_core/manager/resolver.py b/motrix_env_core/src/motrix_env_core/manager/resolver.py index dde903cb..f7e3137f 100644 --- a/motrix_env_core/src/motrix_env_core/manager/resolver.py +++ b/motrix_env_core/src/motrix_env_core/manager/resolver.py @@ -30,16 +30,19 @@ from motrix_env_core.manager.context import ManagerContext from motrix_env_core.manager.env import ManagerEnv, ReadPlan +from motrix_env_core.manager.fingerprint import function_fingerprint from motrix_env_core.manager.kernel_data import ( - KernelDataLayout, - KernelDataLowering, Map, ResolvedLeaf, ResolvedTree, + TreeClassDef, flatten_kernel_data, is_kernel_data, - iter_layout_leaves, + iter_tree_leaves, + leaf_value_type, ) +from motrix_env_core.manager.kernel_data.methods import dispatch_methods +from motrix_env_core.manager.kernel_data.tree import BaseDef, KernelMapDef, LeafDef from motrix_env_core.manager.observations import ObservationGroupEntry, ObsTerm from motrix_env_core.manager.rewards import RewardTerm from motrix_env_core.manager.terminations import TerminationTerm @@ -47,6 +50,8 @@ logger = logging.getLogger(__name__) +_PROGRAM_FINGERPRINT_VERSION = 4 + # Simulator-backed read-only arena views are wrapped as tensors at resolution # time; suppress PyTorch's non-writable-array warning once here instead of # copying every view (mirrors the filter in the Torch executor). @@ -319,6 +324,44 @@ def _tree_leaf_fingerprint(tree: ResolvedTree) -> tuple[Any, ...]: ) +def _kernel_data_program_fingerprint(tree_def: TreeClassDef) -> str: + """Deterministic kernel-data identity for resolved-program invalidation. + + Combines the structural tree fingerprint with per-leaf value types and + the code fingerprints of ``@dispatch`` methods on every nested record, so + schema, leaf-type, or method-code changes each invalidate compiled + program caches. + """ + raw = repr((_PROGRAM_FINGERPRINT_VERSION, _program_fingerprint_part(tree_def))) + return hashlib.sha256(raw.encode()).hexdigest() + + +def _program_fingerprint_part(tree_def: BaseDef) -> Any: + if isinstance(tree_def, LeafDef): + value_type = leaf_value_type(tree_def.annotation) + return ( + "leaf", + tree_def.path, + None if value_type is None else f"{value_type.__module__}.{value_type.__qualname__}", + ) + if isinstance(tree_def, KernelMapDef): + return ( + "map", + tree_def.fingerprint, + tuple((entry.key, _program_fingerprint_part(entry.tree_def)) for entry in tree_def.entries), + ) + assert isinstance(tree_def, TreeClassDef) + return ( + "record", + tree_def.fingerprint, + tuple( + (name, function_fingerprint(method)) + for name, method in sorted(dispatch_methods(tree_def.logical_type).items()) + ), + tuple((field.name, _program_fingerprint_part(field.tree_def)) for field in tree_def.fields), + ) + + def _carrier_fingerprint(value: Any) -> Any: """Structural + content fingerprint of a constant term argument. @@ -365,7 +408,6 @@ def __init__(self, env: ManagerEnv): self._sim_inputs: Mapping[str, ManagerSimInput] self._sim_inputs, sim_input_fingerprint = _manager_sim_inputs(env.sim_data) self._prepared_trees: list[ResolvedTree] = [] - self._kernel_data_lowering = KernelDataLowering() self._layout_identity: list[Any] = [sim_input_fingerprint] def build(self) -> ResolvedManagerProgram: @@ -562,14 +604,11 @@ def _resolve_manager_context( dt=np.float32(self._env.cfg.ctrl_dt), sim_reset_requested=torch.from_numpy(self._env._sim_reset_requested), ) - _, tree_def = flatten_kernel_data(value) - layout = self._kernel_data_lowering.lower( - tree_def, - context="ManagerContext", - ) - tree_index = self._register_tree("manager_context", value, layout) + tree_index, tree_def = self._register_tree("manager_context", value) tree = self._prepared_trees[tree_index] - self._layout_identity.append(("manager_context", layout.fingerprint, _tree_leaf_fingerprint(tree))) + self._layout_identity.append( + ("manager_context", _kernel_data_program_fingerprint(tree_def), _tree_leaf_fingerprint(tree)) + ) self._layout_identity.append( ( "metrics", @@ -759,33 +798,33 @@ def _register_tree( self, source_name: str, value: Any, - layout: KernelDataLayout, - ) -> int: - """Validate one kernel-data value against its layout and record its tree. + ) -> tuple[int, TreeClassDef]: + """Validate one kernel-data value against its tree schema and record its tree. Leaves are captured by reference: array leaves are refreshed in place by their owning runtime objects and scalar/shared leaves are reused unchanged. """ values, tree_def = flatten_kernel_data(value) - if tree_def.fingerprint != layout.tree_def.fingerprint: - raise TypeError(f"{source_name} value has a different KernelData schema.") leaves: dict[tuple[str, ...], ResolvedLeaf] = {} - for prepared_field, field_value in zip(iter_layout_leaves(layout), values, strict=True): - if prepared_field.value_type in (np.ndarray, torch.Tensor): - if not isinstance(field_value, prepared_field.value_type): - raise TypeError( - f"{source_name} field {'.'.join(prepared_field.path)!r} must be " - f"{prepared_field.value_type.__name__}." - ) - elif type(field_value) is not prepared_field.value_type: + for leaf_def, field_value in zip(iter_tree_leaves(tree_def), values, strict=True): + value_type = leaf_value_type(leaf_def.annotation) + if value_type is None: + raise TypeError( + f"{source_name} KernelData {'.'.join(leaf_def.path)} expected a tensor, an ndarray, or a numeric " + "scalar." + ) + if value_type in (np.ndarray, torch.Tensor): + if not isinstance(field_value, value_type): + raise TypeError(f"{source_name} field {'.'.join(leaf_def.path)!r} must be {value_type.__name__}.") + elif type(field_value) is not value_type: raise TypeError( - f"{source_name} field {'.'.join(prepared_field.path)!r} must return " - f"{prepared_field.value_type.__name__}, got {type(field_value).__name__}." + f"{source_name} field {'.'.join(leaf_def.path)!r} must return " + f"{value_type.__name__}, got {type(field_value).__name__}." ) - leaves[prepared_field.path] = ResolvedLeaf(field_value) + leaves[leaf_def.path] = ResolvedLeaf(field_value) self._prepared_trees.append(ResolvedTree(value, leaves)) - return len(self._prepared_trees) - 1 + return len(self._prepared_trees) - 1, tree_def def _validate_observation_spaces(self, groups: dict[str, ObservationGroupEntry]) -> None: expected_policy = self._env.policy_observation_space.shape diff --git a/motrix_env_core/tests/test_kernel_data.py b/motrix_env_core/tests/test_kernel_data.py index 452a4435..84aa6ff1 100644 --- a/motrix_env_core/tests/test_kernel_data.py +++ b/motrix_env_core/tests/test_kernel_data.py @@ -1,5 +1,4 @@ # Copyright Motphys Technology Co., Ltd. 2025, 2026 -# SPDX-License-Identifier: Apache-2.0 from dataclasses import FrozenInstanceError, dataclass @@ -8,11 +7,12 @@ import torch from motrix_env_core.manager.kernel_data import ( - KernelDataLowering, Map, + canonicalize_kernel_data, flatten_kernel_data, is_kernel_data, - iter_layout_leaves, + iter_tree_leaves, + leaf_value_type, unflatten_kernel_data, ) @@ -44,11 +44,6 @@ class _MappedValues: values: Map -@dataclass(frozen=True, slots=True) -class _ArrayMapValues: - values: Map[np.ndarray] - - def test_kernel_data_is_frozen_and_compiles_nested_inherited_leaves() -> None: value = _Tracking( _Robot(np.asarray([[1.0, 2.0]], dtype=np.float32), _Pose(np.asarray([[3.0]]), np.asarray([4.0]))), @@ -58,14 +53,15 @@ def test_kernel_data_is_frozen_and_compiles_nested_inherited_leaves() -> None: value.phase = np.float32(1.0) # type: ignore[misc] _, tree_def = flatten_kernel_data(value) - layout = KernelDataLowering().lower(tree_def, context="test") - leaves = iter_layout_leaves(layout) + leaves = iter_tree_leaves(tree_def) assert [leaf.path for leaf in leaves] == [ ("robot", "velocity"), ("robot", "pose", "position"), ("robot", "pose", "orientation"), ("phase",), ] + assert leaf_value_type(leaves[0].annotation) is np.ndarray + assert leaf_value_type(leaves[-1].annotation) is np.float32 def test_kernel_data_flatten_is_zero_copy_and_rebuilds_host_value() -> None: @@ -76,7 +72,7 @@ def test_kernel_data_flatten_is_zero_copy_and_rebuilds_host_value() -> None: leaves, tree_def = flatten_kernel_data(value) assert leaves[:3] == (velocity, position, orientation) rebuilt = unflatten_kernel_data(tree_def, leaves) - KernelDataLowering().lower(tree_def, context="test") + canonicalize_kernel_data(value, context="test") assert rebuilt.robot.pose.position is position assert rebuilt.robot.pose.orientation is orientation assert rebuilt.phase == np.float32(0.25) @@ -91,8 +87,7 @@ class _Tensors: value = _Tensors(buffer) leaves, tree_def = flatten_kernel_data(value) assert leaves == (buffer,) - layout = KernelDataLowering().lower(tree_def, context="test") - assert iter_layout_leaves(layout)[0].value_type is torch.Tensor + assert leaf_value_type(iter_tree_leaves(tree_def)[0].annotation) is torch.Tensor def test_kernel_data_schema_fingerprint_ignores_values() -> None: @@ -106,24 +101,19 @@ def test_kernel_data_schema_fingerprint_ignores_values() -> None: ) _, first = flatten_kernel_data(first_value) _, second = flatten_kernel_data(second_value) - assert first is second assert first.fingerprint == second.fingerprint - first_layout = KernelDataLowering().lower(first, context="first") - second_layout = KernelDataLowering().lower(second, context="second") - assert first_layout is second_layout - assert first_layout.logical_type.__module__ == _Tracking.__module__ + assert first.logical_type.__module__ == _Tracking.__module__ -def test_map_schema_changes_tree_and_layout_identity() -> None: +def test_map_schema_changes_tree_identity() -> None: first = _MappedValues(Map({"value": np.float32(1.0)})) second = _MappedValues(Map({"value": np.asarray([[1.0]], dtype=np.float32)})) _, first_tree = flatten_kernel_data(first) _, second_tree = flatten_kernel_data(second) - first_layout = KernelDataLowering().lower(first_tree, context="first") - second_layout = KernelDataLowering().lower(second_tree, context="second") assert first_tree.fingerprint != second_tree.fingerprint - assert first_layout.fingerprint != second_layout.fingerprint - assert first_layout.fields[0].child.fingerprint != second_layout.fields[0].child.fingerprint + first_map = first_tree.fields[0].tree_def + second_map = second_tree.fields[0].tree_def + assert first_map.fingerprint != second_map.fingerprint def test_kernel_data_rejects_unsupported_leaf() -> None: @@ -131,9 +121,8 @@ def test_kernel_data_rejects_unsupported_leaf() -> None: class _Invalid: label: str - _, tree_def = flatten_kernel_data(_Invalid("invalid")) with pytest.raises(TypeError, match="expected a tensor"): - KernelDataLowering().lower(tree_def, context="invalid") + canonicalize_kernel_data(_Invalid("invalid"), context="invalid") def test_ndarrays_and_tensors_are_not_kernel_data_values() -> None: From 8327bee57b0304103add820375c5006e4223326d Mon Sep 17 00:00:00 2001 From: zilch <147668916@qq.com> Date: Fri, 6 Nov 2026 21:12:47 +0800 Subject: [PATCH 32/41] feat(manager)!: exchange step actions as torch tensors 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. --- bench/env_benchmark/runner.py | 2 +- bench/env_benchmark/workload.py | 15 +++++++-- .../src/motrix_env_core/array/env.py | 11 +++++-- .../src/motrix_env_core/manager/actions.py | 10 +++--- .../src/motrix_env_core/manager/env.py | 9 ++++-- .../src/motrix_env_core/mdp/action.py | 22 +++++++------ motrix_env_core/tests/test_action_state.py | 2 +- motrix_env_core/tests/test_manager_env.py | 32 +++++++++---------- .../tests/test_manager_sim_backend.py | 10 +++--- .../tests/test_torch_manager_actions.py | 6 ++-- motrix_envs/tests/test_walk_randomization.py | 16 +++++----- .../src/motrix_rl_builtin/fastsac/wrap_np.py | 25 +++++++++++---- .../src/motrix_rl_rslrl/torch/wrap_np.py | 12 ++++--- .../src/motrix_rl_skrl/torch/wrap_np.py | 8 +++-- motrixlab/src/motrixlab/commands/view.py | 7 ++++ tests/test_bench_workload.py | 2 +- 16 files changed, 117 insertions(+), 72 deletions(-) diff --git a/bench/env_benchmark/runner.py b/bench/env_benchmark/runner.py index eff5331f..e943d0ce 100644 --- a/bench/env_benchmark/runner.py +++ b/bench/env_benchmark/runner.py @@ -68,7 +68,7 @@ def run_steps( expected_spec = None names = [args.env] + ([args.compare_env] if args.compare_env is not None else []) with redirect_stdout(sys.stderr) if args.json else redirect_stdout(sys.stdout): - np, _ = workload.load_runtime() + np, _, _ = workload.load_runtime() for name in names: env, actions, expected_spec = workload.prepare_env(name, args, num_envs, actions, expected_spec) try: diff --git a/bench/env_benchmark/workload.py b/bench/env_benchmark/workload.py index 660d9768..20a8c713 100644 --- a/bench/env_benchmark/workload.py +++ b/bench/env_benchmark/workload.py @@ -18,8 +18,9 @@ def load_runtime(): import numpy as np from motrix_env_core import registry + from motrix_env_core.manager import ManagerEnv - return np, registry + return np, registry, ManagerEnv def prepare_env( @@ -30,7 +31,7 @@ def prepare_env( expected_action_spec=None, ): """Create an environment and compatible fixed actions, checking comparisons.""" - np, registry = load_runtime() + np, registry, manager_env_type = load_runtime() kwargs = {"mode": args.mode, "num_envs": num_envs} if args.sim_backend is not None: kwargs["sim_backend"] = args.sim_backend @@ -56,6 +57,16 @@ def prepare_env( low = np.where(np.isfinite(space.low), space.low, -1.0) high = np.where(np.isfinite(space.high), space.high, 1.0) actions = np.random.default_rng(args.seed).uniform(low, high, size=shape).astype(space.dtype) + elif not isinstance(actions, np.ndarray): + # Compare runs pass the previous env's actions through; canonicalize + # back to numpy so every env receives identically valued actions. + actions = np.asarray(actions) + # Manager environments consume tensor actions; hand each env a + # zero-copy tensor view of the same canonical numpy values. + if isinstance(env, manager_env_type): + import torch + + actions = torch.from_numpy(np.ascontiguousarray(actions)) return env, actions, action_spec except BaseException: close_env(env) diff --git a/motrix_env_core/src/motrix_env_core/array/env.py b/motrix_env_core/src/motrix_env_core/array/env.py index fff7bb95..abe74461 100644 --- a/motrix_env_core/src/motrix_env_core/array/env.py +++ b/motrix_env_core/src/motrix_env_core/array/env.py @@ -311,12 +311,14 @@ def create_renderer(self, config: RenderConfig) -> SimRenderer: """Create the backend-owned renderer bound to this environment.""" @abc.abstractmethod - def apply_action(self, actions: np.ndarray, state: ArrayEnvState[T]) -> ArrayEnvState[T]: + def apply_action(self, actions: T, state: ArrayEnvState[T]) -> ArrayEnvState[T]: """ Apply the action to the environment Args: - actions (np.ndarray): The actions to apply + actions (T): The actions to apply — the same array flavor the + environment carries end to end (numpy arrays for the NumPy + frontend, torch tensors for the manager pipeline) state (ArrayEnvState): The environment state to apply the actions. """ @@ -359,9 +361,12 @@ def compute_observation(self, state: ArrayEnvState[T]) -> ArrayEnvState[T]: return state @perf_root("step") - def step(self, actions: np.ndarray) -> ArrayEnvState[T]: + def step(self, actions: T) -> ArrayEnvState[T]: """Advance one control step. + ``actions`` flows the environment's array flavor: numpy arrays for + NumPy-frontend environments, torch tensors for the manager pipeline. + Timing: ``apply_action`` -> ``physics_step`` -> ``compute_transition`` (the only full simulator-data refresh) -> episode bookkeeping -> optional pre-reset ``final_obs`` capture for truncated rows -> diff --git a/motrix_env_core/src/motrix_env_core/manager/actions.py b/motrix_env_core/src/motrix_env_core/manager/actions.py index f0353089..56efb2f8 100644 --- a/motrix_env_core/src/motrix_env_core/manager/actions.py +++ b/motrix_env_core/src/motrix_env_core/manager/actions.py @@ -67,17 +67,17 @@ def __init__(self, action_space: gym.spaces.Box, state: ActionState) -> None: self.action_space = action_space self.state = canonicalize_kernel_data(state, context="ActionTerm state") - def process(self, actions: np.ndarray) -> np.ndarray | None: - """Record raw actions and process one batched action slice.""" + def process(self, actions: torch.Tensor) -> torch.Tensor | None: + """Record raw tensor actions and process one batched action slice.""" state = self.state ptr = (int(state.action_ptr[0]) + 1) % state.action_queue.shape[1] state.action_ptr[0] = ptr - state.action_queue.numpy()[:, ptr] = actions + state.action_queue[:, ptr] = actions return self._process(actions) @abc.abstractmethod - def _process(self, actions: np.ndarray) -> np.ndarray | None: - """Process one batched action slice into route-local controls.""" + def _process(self, actions: torch.Tensor) -> torch.Tensor | None: + """Process one batched tensor action slice into route-local controls.""" def reset(self, env_ids: np.ndarray) -> None: """Clear raw history and reset action-specific state for selected rows.""" diff --git a/motrix_env_core/src/motrix_env_core/manager/env.py b/motrix_env_core/src/motrix_env_core/manager/env.py index 9da14f61..d40b37c3 100644 --- a/motrix_env_core/src/motrix_env_core/manager/env.py +++ b/motrix_env_core/src/motrix_env_core/manager/env.py @@ -593,10 +593,10 @@ def observation_space(self) -> ObsSpace: def action_space(self) -> gym.spaces.Box: return self._action_space - def apply_action(self, actions: np.ndarray, state: ArrayEnvState[torch.Tensor]) -> ArrayEnvState[torch.Tensor]: + def apply_action(self, actions: torch.Tensor, state: ArrayEnvState[torch.Tensor]) -> ArrayEnvState[torch.Tensor]: expected_shape = (self.num_envs, self.action_space.shape[0]) if actions.shape != expected_shape: - raise ValueError(f"Expected action shape {expected_shape}, got {actions.shape}.") + raise ValueError(f"Expected action shape {expected_shape}, got {tuple(actions.shape)}.") for name, term in self._action_terms.items(): controls = term.process(actions[:, self._action_slices[name]]) actuators = self._action_actuators[name] @@ -604,7 +604,10 @@ def apply_action(self, actions: np.ndarray, state: ArrayEnvState[torch.Tensor]) assert controls is None continue assert controls is not None - self._action_writes.buffer(name)[:] = controls + # The compiled write program owns numpy buffers. ``.numpy()`` is a + # zero-copy view of the controls tensor; the slice assignment is + # the single torch→numpy copy on this path. + self._action_writes.buffer(name)[:] = controls.numpy() self._action_writes.execute() return state diff --git a/motrix_env_core/src/motrix_env_core/mdp/action.py b/motrix_env_core/src/motrix_env_core/mdp/action.py index e1e7c35b..92f089d0 100644 --- a/motrix_env_core/src/motrix_env_core/mdp/action.py +++ b/motrix_env_core/src/motrix_env_core/mdp/action.py @@ -61,21 +61,23 @@ class JointPositionActionTerm(ActionTerm): def __init__(self, space: gym.spaces.Box, state: JointPositionActionState) -> None: super().__init__(space, state) - def _process(self, actions: np.ndarray) -> np.ndarray: - """Map raw policy actions to targets using manager-owned action history.""" + def _process(self, actions: torch.Tensor) -> torch.Tensor: + """Map raw policy actions to targets using manager-owned action history. + + Stays on tensors end to end: the delayed-slot lookup uses + ``torch.take_along_dim`` and the affine target map multiplies torch + buffers directly, mirroring + ``motrix_env_core.torch_manager.actions.process_joint_position_actions``. + """ state = self.state ptr = int(state.action_ptr[0]) target_raw = actions if state.delay_hi > 0: - queue = state.action_queue.numpy() + queue = state.action_queue width = queue.shape[1] - delay_indices = (ptr - state.delay_steps.numpy()) % width - target_raw = np.take_along_axis( - queue, - delay_indices[:, None, None], - axis=1, - )[:, 0, :] - return target_raw * state.action_scales.numpy() + state.default_angles.numpy() + delay_indices = (ptr - state.delay_steps) % width + target_raw = torch.take_along_dim(queue, delay_indices[:, None, None], dim=1)[:, 0, :] + return target_raw * state.action_scales + state.default_angles def _reset(self, env_ids: np.ndarray) -> None: """Resample per-episode delay after the base term clears raw history.""" diff --git a/motrix_env_core/tests/test_action_state.py b/motrix_env_core/tests/test_action_state.py index 55556477..7dea4885 100644 --- a/motrix_env_core/tests/test_action_state.py +++ b/motrix_env_core/tests/test_action_state.py @@ -70,7 +70,7 @@ def test_joint_position_action_applies_per_lane_delay() -> None: ) term = JointPositionActionTerm(gym.spaces.Box(-np.inf, np.inf, shape=(1,)), state) - np.testing.assert_array_equal(term._process(np.asarray([[0.0], [0.0]], dtype=np.float32)), [[140.0], [102.0]]) + np.testing.assert_array_equal(term._process(torch.zeros(2, 1, dtype=torch.float32)).numpy(), [[140.0], [102.0]]) def test_joint_position_action_resamples_delay_on_reset(monkeypatch) -> None: diff --git a/motrix_env_core/tests/test_manager_env.py b/motrix_env_core/tests/test_manager_env.py index a2df00a0..545be58f 100644 --- a/motrix_env_core/tests/test_manager_env.py +++ b/motrix_env_core/tests/test_manager_env.py @@ -62,8 +62,8 @@ class _TestActionTerm(ActionTerm): def __init__(self, action_space: gym.spaces.Box, state: _TestActionState) -> None: super().__init__(action_space, state) - def _process(self, actions: np.ndarray) -> None: - self.state.source.copy_(torch.from_numpy(actions)) + def _process(self, actions: torch.Tensor) -> None: + self.state.source.copy_(actions) def _reset(self, env_ids: np.ndarray) -> None: self.state.reset_flags.fill_(False) @@ -74,8 +74,8 @@ class _SecondActionTerm(ActionTerm): def __init__(self, action_space: gym.spaces.Box, state: _SecondActionState) -> None: super().__init__(action_space, state) - def _process(self, actions: np.ndarray) -> None: - self.state.source.copy_(torch.from_numpy(actions)) + def _process(self, actions: torch.Tensor) -> None: + self.state.source.copy_(actions) def _reset(self, env_ids: np.ndarray) -> None: self.state.reset_flags.fill_(False) @@ -365,7 +365,7 @@ def record_read() -> None: monkeypatch.setattr(env, "_refresh_sim_reads", record_read) - state = env.step(np.ones((env.num_envs, 1), dtype=np.float32)) + state = env.step(torch.ones((env.num_envs, 1), dtype=torch.float32)) np.testing.assert_array_equal(state.terminated, True) assert read_count == 1 @@ -423,7 +423,7 @@ class _MultipleActionsCfg(ManagerActionsCfg): cfg = _ManagerEnvCfg(actions=_MultipleActionsCfg(), observations=_MultipleActionObservationsCfg()) env = ManagerEnv(cfg, num_envs=2) state = env.init_state() - actions = np.asarray([[0.25, 1.0, 2.0], [0.5, -1.0, -2.0]], dtype=np.float32) + actions = torch.tensor([[0.25, 1.0, 2.0], [0.5, -1.0, -2.0]], dtype=torch.float32) env.apply_action(actions, state) @@ -453,7 +453,7 @@ def test_multiple_action_terms_validate_total_action_shape() -> None: state = env.init_state() with pytest.raises(ValueError, match=r"Expected action shape \(2, 1\), got \(2, 2\)"): - env.apply_action(np.zeros((2, 2), dtype=np.float32), state) + env.apply_action(torch.zeros((2, 2), dtype=torch.float32), state) def test_rand_value_owns_seeded_per_environment_states() -> None: @@ -473,7 +473,7 @@ def test_step_perf_records_standard_and_manager_phases() -> None: state = env.init_state() assert state.episode_steps.shape == (env.num_envs,) assert set(state.reward_terms) == {"source"} - actions = np.zeros((env.num_envs, 1), dtype=np.float32) + actions = torch.zeros((env.num_envs, 1), dtype=torch.float32) env.step(actions) assert env.perf.snapshot() == () @@ -505,14 +505,14 @@ def test_step_preserves_previous_observation_with_alternating_preallocated_buffe initial_policy = initial.obs.policy initial_snapshot = initial_policy.clone() - first = env.step(np.full((env.num_envs, 1), 0.25, dtype=np.float32)) + first = env.step(torch.full((env.num_envs, 1), 0.25, dtype=torch.float32)) np.testing.assert_array_equal(initial_policy, initial_snapshot) assert initial_policy.data_ptr() != first.obs.policy.data_ptr() first_policy = first.obs.policy first_snapshot = first_policy.clone() - second = env.step(np.full((env.num_envs, 1), 0.5, dtype=np.float32)) + second = env.step(torch.full((env.num_envs, 1), 0.5, dtype=torch.float32)) np.testing.assert_array_equal(first_policy, first_snapshot) assert first_policy.data_ptr() != second.obs.policy.data_ptr() @@ -617,7 +617,7 @@ def test_command_term_updates_lifecycle_and_selected_reset_ids() -> None: def test_done_envs_reset_before_observation_and_preserve_transition_outputs() -> None: env = _ManagerEnv(num_envs=3) state = env.init_state() - state = env.step(np.asarray([[0.25], [0.75], [0.25]], dtype=np.float32)) + state = env.step(torch.tensor([[0.25], [0.75], [0.25]], dtype=torch.float32)) np.testing.assert_allclose(state.obs.policy[[0, 2], 0], 1.75) np.testing.assert_array_equal(state.obs.policy[1], [3.25, -1.0, 1.5, 1.0]) @@ -629,7 +629,7 @@ def test_done_envs_reset_before_observation_and_preserve_transition_outputs() -> def test_truncated_envs_reset_before_observation() -> None: env = _ManagerEnv(num_envs=2, max_episode_seconds=0.01) - state = env.step(np.asarray([[0.25], [0.25]], dtype=np.float32)) + state = env.step(torch.full((2, 1), 0.25, dtype=torch.float32)) np.testing.assert_array_equal(state.truncated, True) np.testing.assert_array_equal(state.episode_steps, 0) @@ -649,7 +649,7 @@ def test_final_observations_capture_command_before_reset_and_survive_steps(with_ assert initial.final_obs is None initial.episode_steps[:] = [2, 0, 2] env.command_terms["counter"].command[:] = 2.0 - state = env.step(np.full((3, 1), 0.25, dtype=np.float32)) + state = env.step(torch.full((3, 1), 0.25, dtype=torch.float32)) final = state.final_obs assert final is not None np.testing.assert_array_equal(final.env_ids, [0, 2]) @@ -664,10 +664,10 @@ def test_final_observations_capture_command_before_reset_and_survive_steps(with_ assert final.obs.value_or_policy is final.obs.policy assert final.obs.policy.data_ptr() != state.obs.policy.data_ptr() saved = final.obs.policy.clone() - next_state = env.step(np.full((3, 1), 0.25, dtype=np.float32)) + next_state = env.step(torch.full((3, 1), 0.25, dtype=torch.float32)) assert next_state.final_obs is None assert state.final_obs is final - env.step(np.full((3, 1), 0.25, dtype=np.float32)) + env.step(torch.full((3, 1), 0.25, dtype=torch.float32)) np.testing.assert_array_equal(final.obs.policy, saved) @@ -809,7 +809,7 @@ class _RenamedRewardsCfg(ManagerRewardsCfg): env = ManagerEnv(_ManagerEnvCfg(rewards=_RenamedRewardsCfg()), num_envs=2) state = env.init_state() - actions = np.full((env.num_envs, 1), 0.5, dtype=np.float32) + actions = torch.full((env.num_envs, 1), 0.5, dtype=torch.float32) env.apply_action(actions, state) env.compute_transition(state) diff --git a/motrix_env_core/tests/test_manager_sim_backend.py b/motrix_env_core/tests/test_manager_sim_backend.py index a0c03a78..1f6dc0b2 100644 --- a/motrix_env_core/tests/test_manager_sim_backend.py +++ b/motrix_env_core/tests/test_manager_sim_backend.py @@ -208,9 +208,9 @@ class _CtrlActionTerm(ActionTerm): def __init__(self, action_space: gym.spaces.Box, state: _CtrlActionState) -> None: super().__init__(action_space, state) - def _process(self, actions: np.ndarray) -> np.ndarray: - self.state.source.copy_(torch.from_numpy(actions)) - return actions * np.float32(2.0) + def _process(self, actions: torch.Tensor) -> torch.Tensor: + self.state.source.copy_(actions) + return actions * 2.0 def _reset(self, env_ids: np.ndarray) -> None: self.state.source[env_ids] = 0.0 @@ -314,7 +314,7 @@ def test_manager_env_steps_on_fake_backend() -> None: env = _make_env() runtime = _FakeBackend.last - state = env.step(np.full((3, 1), 0.5, dtype=np.float32)) + state = env.step(torch.full((3, 1), 0.5, dtype=torch.float32)) assert runtime.step_calls == [env.cfg.sim_substeps] # ctrl routing: only the "routed" actuator receives the doubled action. @@ -333,7 +333,7 @@ def test_partial_reset_hands_canonical_rows_to_runtime_and_clears_ctrl() -> None env.init_state() runtime.reset_calls.clear() - env.step(np.full((3, 1), 0.5, dtype=np.float32)) + env.step(torch.full((3, 1), 0.5, dtype=torch.float32)) assert runtime.ctrl_targets[1, 1] == np.float32(1.0) env._state.terminated[:] = torch.tensor([False, True, False]) diff --git a/motrix_env_core/tests/test_torch_manager_actions.py b/motrix_env_core/tests/test_torch_manager_actions.py index f38a47e3..19ed84c0 100644 --- a/motrix_env_core/tests/test_torch_manager_actions.py +++ b/motrix_env_core/tests/test_torch_manager_actions.py @@ -49,8 +49,8 @@ def _exercise_process(process: Callable, *, width: int, delayed: bool) -> None: actions = torch.from_numpy(raw.copy()) queue_before, ptr_before = queue.clone(), ptr.clone() targets, new_queue, new_ptr, current, previous = process(queue, ptr, actions, delays, scales, defaults) - expected = term.process(raw) - torch.testing.assert_close(targets, torch.from_numpy(expected)) + expected = term.process(torch.from_numpy(raw)) + torch.testing.assert_close(targets, expected) torch.testing.assert_close(new_queue, state.action_queue, rtol=0, atol=0) torch.testing.assert_close(new_ptr, state.action_ptr, rtol=0, atol=0) torch.testing.assert_close(current, state.current(), rtol=0, atol=0) @@ -97,7 +97,7 @@ def _exercise_reset(reset: Callable, process: Callable) -> None: state.action_scales, state.default_angles, ) - torch.testing.assert_close(targets, torch.from_numpy(term.process(raw))) + torch.testing.assert_close(targets, term.process(torch.from_numpy(raw))) torch.testing.assert_close(queue, state.action_queue, rtol=0, atol=0) torch.testing.assert_close(current, state.current(), rtol=0, atol=0) torch.testing.assert_close(previous, state.previous(), rtol=0, atol=0) diff --git a/motrix_envs/tests/test_walk_randomization.py b/motrix_envs/tests/test_walk_randomization.py index d2e551aa..22db019f 100644 --- a/motrix_envs/tests/test_walk_randomization.py +++ b/motrix_envs/tests/test_walk_randomization.py @@ -65,20 +65,20 @@ def _tiny_delay_manager(num_envs: int = 2, num_actuators: int = 3, lo: int = 0, def test_action_delay_applies_delayed_action(): env, action_state, controls = _tiny_delay_manager(lo=1, hi=2) action_state.delay_steps[:] = 2 - actions = np.full(controls.shape, 1.0, dtype=np.float32) + actions = torch.full(controls.shape, 1.0, dtype=torch.float32) env.apply_action(actions, env._state) # Only the raw action exists so far; unwritten history remains zero. assert np.all(controls == 0.0) assert bool((action_state.action_queue[:, int(action_state.action_ptr[0])] == 1.0).all()) for value in (2.0, 3.0, 4.0): - env.apply_action(np.full_like(actions, value), env._state) + env.apply_action(torch.full_like(actions, value), env._state) # Two-step delay: the applied action is the one from two steps ago. assert np.all(controls == 2.0) def test_action_delay_disabled_passes_actions_through_and_keeps_history(): env, action_state, controls = _tiny_delay_manager(lo=0, hi=0) - actions = np.full(controls.shape, 0.5, dtype=np.float32) + actions = torch.full(controls.shape, 0.5, dtype=torch.float32) env.apply_action(actions, env._state) assert np.all(controls == 0.5) ptr = int(action_state.action_ptr[0]) @@ -91,7 +91,7 @@ def test_action_queue_wraps_and_selects_delay_per_lane(): action_state.delay_steps[:] = torch.tensor([0, 2]) width = action_state.action_queue.shape[1] for value in range(1, width + 2): - actions = np.full(controls.shape, value, dtype=np.float32) + actions = torch.full(controls.shape, value, dtype=torch.float32) env.apply_action(actions, env._state) assert int(action_state.action_ptr[0]) == value % width assert controls[0, 0] == value @@ -102,7 +102,7 @@ def test_action_queue_wraps_and_selects_delay_per_lane(): def test_action_reset_clears_only_selected_lane_history_and_keeps_pointer(): env, action_state, controls = _tiny_delay_manager(num_envs=2, lo=0, hi=2) for value in (1.0, 2.0, 3.0): - env.apply_action(np.full(controls.shape, value, dtype=np.float32), env._state) + env.apply_action(torch.full(controls.shape, value, dtype=torch.float32), env._state) ptr_before_reset = action_state.action_ptr.clone() other_lane_history = action_state.action_queue[1].clone() env.reset(np.asarray([0], dtype=np.int64)) @@ -121,9 +121,9 @@ def test_joint_position_term_methods_manage_history_and_specific_state(): env, action_state, controls = _tiny_delay_manager(lo=0, hi=2) term = env._action_terms["joint_position"] for value in (1.0, 2.0, 3.0): - env.apply_action(np.full(controls.shape, value, dtype=np.float32), env._state) + env.apply_action(torch.full(controls.shape, value, dtype=torch.float32), env._state) pointer = int(action_state.action_ptr[0]) - term.process(np.full(controls.shape, 4.0, dtype=np.float32)) + term.process(torch.full(controls.shape, 4.0, dtype=torch.float32)) assert int(action_state.action_ptr[0]) == (pointer + 1) % action_state.action_queue.shape[1] torch.testing.assert_close(action_state.current(), torch.full_like(action_state.current(), 4.0)) term.reset(np.asarray([0], dtype=np.int64)) @@ -152,7 +152,7 @@ def test_walk_env_randomized_reset_samples_within_ranges(): cfg = make_g129dof_walk_flat_cfg() env = ManagerEnv(cfg, num_envs=2) try: - env.step(np.zeros((2, env.num_actuators), dtype=np.float32)) + env.step(torch.zeros((2, env.num_actuators), dtype=torch.float32)) env.reset(np.arange(2)) walk = env.command_terms["walk"] # Phase increments invert the configured gait-period interval. diff --git a/motrix_rl_builtin/src/motrix_rl_builtin/fastsac/wrap_np.py b/motrix_rl_builtin/src/motrix_rl_builtin/fastsac/wrap_np.py index 36366973..a1d43819 100644 --- a/motrix_rl_builtin/src/motrix_rl_builtin/fastsac/wrap_np.py +++ b/motrix_rl_builtin/src/motrix_rl_builtin/fastsac/wrap_np.py @@ -4,9 +4,9 @@ """NumPy-frontend env wrapper for FastSAC. Wraps a Motrix ``ArrayEnv`` as the torch-tensor vectorized environment consumed -by FastSAC. Manager environments already expose torch state; the tensors pass -through zero-copy (only a device move when training off-CPU). NumPy-frontend -(Direct) environments convert at this boundary. +by FastSAC. Manager environments already expose torch state and consume torch +actions; both pass through zero-copy (only a device move when training +off-CPU). NumPy-frontend (Direct) environments convert at this boundary. The actor maps its tanh output to the action range declared by the environment; the wrapper only clips to that standard contract before stepping. Done @@ -18,6 +18,7 @@ import torch from motrix_env_core.array.env import ArrayEnv +from motrix_env_core.manager import ManagerEnv from motrix_env_core.renderer import RenderConfig, create_renderer from motrix_rl.utils import env_infos from motrix_rl_builtin.fastsac.wrap import FastSacEnvWrap, FastSacStepResult @@ -29,6 +30,13 @@ class FastSacNpEnvWrap(FastSacEnvWrap): def __init__(self, env: ArrayEnv, device: torch.device, render: RenderConfig | None = None): super().__init__(env, device) self._env: ArrayEnv = env + # Manager environments consume tensor actions: clip with torch ops and + # pass the tensor through. Direct (NumPy-frontend) environments keep + # the numpy clip-and-convert path. + self._torch_actions = isinstance(env, ManagerEnv) + if self._torch_actions: + self._low_t = torch.as_tensor(self._low, dtype=torch.float32) + self._high_t = torch.as_tensor(self._high, dtype=torch.float32) self._renderer = create_renderer(env, render) def _to_torch(self, arr, dtype=torch.float32) -> torch.Tensor: @@ -41,10 +49,13 @@ def reset(self) -> tuple[torch.Tensor, torch.Tensor]: return self._to_torch(state.obs.policy), self._to_torch(state.obs.value_or_policy) def step(self, actions: torch.Tensor) -> FastSacStepResult: - # Zero-copy CPU view of the learner's actions; clipping allocates the - # buffer handed to the numpy-frontend action path. - np_actions = np.clip(actions.detach().cpu().numpy(), self._low, self._high) - state = self._env.step(np_actions) + if self._torch_actions: + env_actions = torch.clamp(actions.detach().cpu(), self._low_t, self._high_t) + else: + # Zero-copy CPU view of the learner's actions; clipping allocates the + # buffer handed to the numpy-frontend action path. + env_actions = np.clip(actions.detach().cpu().numpy(), self._low, self._high) + state = self._env.step(env_actions) return FastSacStepResult( obs=self._to_torch(state.obs.policy), critic_obs=self._to_torch(state.obs.value_or_policy), diff --git a/motrix_rl_rslrl/src/motrix_rl_rslrl/torch/wrap_np.py b/motrix_rl_rslrl/src/motrix_rl_rslrl/torch/wrap_np.py index 2663f465..049a3526 100644 --- a/motrix_rl_rslrl/src/motrix_rl_rslrl/torch/wrap_np.py +++ b/motrix_rl_rslrl/src/motrix_rl_rslrl/torch/wrap_np.py @@ -9,6 +9,7 @@ from tensordict import TensorDict from motrix_env_core.array.env import ArrayEnv, EnvObs +from motrix_env_core.manager import ManagerEnv from motrix_env_core.renderer import RenderConfig, create_renderer @@ -88,11 +89,12 @@ def unwrapped(self) -> "RslrlNpEnvWrap": return self def step(self, actions: torch.Tensor) -> tuple[TensorDict, torch.Tensor, torch.Tensor, dict]: - # Convert torch actions to numpy - actions_np = actions.cpu().numpy() - - # Step the environment - state = self._env.step(actions_np) + if isinstance(self._env, ManagerEnv): + # Manager environments consume tensor actions; pass through. + state = self._env.step(actions.cpu()) + else: + # Convert torch actions to numpy + state = self._env.step(actions.cpu().numpy()) self._state = state # Update episode length buffer diff --git a/motrix_rl_skrl/src/motrix_rl_skrl/torch/wrap_np.py b/motrix_rl_skrl/src/motrix_rl_skrl/torch/wrap_np.py index 3a92f142..abdcf047 100644 --- a/motrix_rl_skrl/src/motrix_rl_skrl/torch/wrap_np.py +++ b/motrix_rl_skrl/src/motrix_rl_skrl/torch/wrap_np.py @@ -9,6 +9,7 @@ from skrl.envs.wrappers.torch import Wrapper as SkrlWrapper from motrix_env_core.array.env import ArrayEnv +from motrix_env_core.manager import ManagerEnv from motrix_env_core.renderer import RenderConfig, VideoRecorder, create_renderer from motrix_env_core.sim.backend import SimRenderer from motrix_rl.device import resolve_default_device @@ -57,8 +58,11 @@ def step( torch.Tensor, Any, ]: - actions = actions.cpu().numpy() - state = self._env.step(actions) + if isinstance(self._env, ManagerEnv): + # Manager environments consume tensor actions; pass through. + state = self._env.step(actions.detach().cpu()) + else: + state = self._env.step(actions.cpu().numpy()) return ( _as_torch(state.obs.policy, torch.float32, self.device), _as_torch(state.reward.reshape(-1, 1), torch.float32, self.device), diff --git a/motrixlab/src/motrixlab/commands/view.py b/motrixlab/src/motrixlab/commands/view.py index 61ab01fd..f8b1ee7b 100644 --- a/motrixlab/src/motrixlab/commands/view.py +++ b/motrixlab/src/motrixlab/commands/view.py @@ -22,6 +22,7 @@ from motrix_env_core.config.scene.base import RobotCfg, SceneCfg, SceneObjsCfg, SceneVisualCfg, SystemCameraCfg from motrix_env_core.config.scene.geometry import FlatTerrainCfg from motrix_env_core.config.scene.light import LightCfg +from motrix_env_core.manager import ManagerEnv from motrix_env_core.renderer import RenderConfig, create_renderer from motrix_env_motrixsim.compiler import build_scene_model @@ -56,11 +57,17 @@ def _sample_random_action_np(env: ABEnv) -> np.ndarray: def _run_np(env: ArrayEnv): + import torch + renderer = create_renderer(env, RenderConfig()) env_dt = env.cfg.ctrl_dt + # Manager environments consume tensor actions; sample the env's currency. + torch_actions = isinstance(env, ManagerEnv) while True: t0 = time.monotonic() actions = _sample_random_action_np(env) + if torch_actions: + actions = torch.from_numpy(actions) env.step(actions) renderer.render() real_dt = time.monotonic() - t0 diff --git a/tests/test_bench_workload.py b/tests/test_bench_workload.py index f6af6771..2a7d11f3 100644 --- a/tests/test_bench_workload.py +++ b/tests/test_bench_workload.py @@ -38,7 +38,7 @@ def test_comparison_reuses_actions_and_closes_incompatible_environment(monkeypat ) environments = iter([reference, compared, incompatible]) registry = SimpleNamespace(make=lambda *args, **kwargs: next(environments)) - monkeypatch.setattr(workload, "load_runtime", lambda: (np, registry)) + monkeypatch.setattr(workload, "load_runtime", lambda: (np, registry, dict)) args = SimpleNamespace(mode="train", sim_backend=None, zero_actions=False, seed=123) env, actions, spec = workload.prepare_env("reference", args, 3) assert env is reference From 8dca412b8e41410d157783ff683e254d2168f71b Mon Sep 17 00:00:00 2001 From: zilch <147668916@qq.com> Date: Fri, 6 Nov 2026 21:27:20 +0800 Subject: [PATCH 33/41] refactor(core)!: implement the state allocation hooks via the array namespace 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. --- .../src/motrix_env_core/array/env.py | 23 ++++++++++++------- .../src/motrix_env_core/manager/env.py | 13 ++--------- 2 files changed, 17 insertions(+), 19 deletions(-) diff --git a/motrix_env_core/src/motrix_env_core/array/env.py b/motrix_env_core/src/motrix_env_core/array/env.py index abe74461..8abe6e7c 100644 --- a/motrix_env_core/src/motrix_env_core/array/env.py +++ b/motrix_env_core/src/motrix_env_core/array/env.py @@ -6,6 +6,7 @@ from dataclasses import dataclass from typing import Generic, TypeVar +import array_api_compat.numpy as array_namespace_numpy import gymnasium as gym import numpy as np import torch @@ -223,17 +224,23 @@ def value_observation_space(self) -> gym.spaces.Box: def has_value_observation(self) -> bool: return _as_obs_space(self.observation_space).value is not None - def _zeros_box(self, space: gym.spaces.Box) -> np.ndarray: - return np.zeros((self._num_envs, *space.shape), dtype=space.dtype) + # Array module for the frontend's native currency (array_api_compat + # namespaces expose zeros/ones/full with module dtypes); numpy frontends + # inherit this default, torch frontends override the attribute only. + _xp = array_namespace_numpy - def _zeros_reward(self) -> np.ndarray: - return np.zeros((self._num_envs,), dtype=np.float32) + def _zeros_box(self, space: gym.spaces.Box) -> T: + dtype = {np.dtype(np.float32): self._xp.float32, np.dtype(np.float64): self._xp.float64}[np.dtype(space.dtype)] + return self._xp.zeros((self._num_envs, *space.shape), dtype=dtype) - def _zeros_terminated(self) -> np.ndarray: - return np.ones((self._num_envs,), dtype=bool) + def _zeros_reward(self) -> T: + return self._xp.zeros((self._num_envs,), dtype=self._xp.float32) - def _zeros_truncated(self) -> np.ndarray: - return np.zeros((self._num_envs,), dtype=bool) + def _zeros_terminated(self) -> T: + return self._xp.ones((self._num_envs,), dtype=self._xp.bool) + + def _zeros_truncated(self) -> T: + return self._xp.zeros((self._num_envs,), dtype=self._xp.bool) def _zeros_obs(self, space: gym.Space | ObsSpace) -> EnvObs[T]: obs_space = _as_obs_space(space) diff --git a/motrix_env_core/src/motrix_env_core/manager/env.py b/motrix_env_core/src/motrix_env_core/manager/env.py index d40b37c3..4ef9f9e7 100644 --- a/motrix_env_core/src/motrix_env_core/manager/env.py +++ b/motrix_env_core/src/motrix_env_core/manager/env.py @@ -9,6 +9,7 @@ from time import perf_counter from typing import TYPE_CHECKING, Any, TypeVar +import array_api_compat.torch as array_namespace_torch import gymnasium as gym import numpy as np import torch @@ -414,17 +415,7 @@ def physics_step(self) -> None: # The manager pipeline's transition storage is torch-native: the compiled # kernels write tensors directly and torch consumers pass them through # without copies. Non-torch consumers take zero-copy ``.numpy()`` views. - def _zeros_box(self, space: gym.spaces.Box) -> torch.Tensor: - return torch.zeros((self.num_envs, *space.shape), dtype=torch.float32) - - def _zeros_reward(self) -> torch.Tensor: - return torch.zeros((self.num_envs,), dtype=torch.float32) - - def _zeros_terminated(self) -> torch.Tensor: - return torch.ones((self.num_envs,), dtype=torch.bool) - - def _zeros_truncated(self) -> torch.Tensor: - return torch.zeros((self.num_envs,), dtype=torch.bool) + _xp = array_namespace_torch def init_state(self) -> ArrayEnvState[torch.Tensor]: startup_started = perf_counter() From 033fecdd857456aa14d02c314655235bcfb3dd10 Mon Sep 17 00:00:00 2001 From: zilch <147668916@qq.com> Date: Fri, 6 Nov 2026 21:40:11 +0800 Subject: [PATCH 34/41] refactor(kernel_data)!: replace the flatten ABI with tree visitors 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. --- .../manager/kernel_data/__init__.py | 13 +-- .../manager/kernel_data/canonical.py | 20 ++--- .../manager/kernel_data/tree.py | 86 +++++++------------ .../src/motrix_env_core/manager/resolver.py | 20 ++--- motrix_env_core/tests/test_kernel_data.py | 60 +++++++------ wiki/design/manager/kernel-abi.md | 7 +- 6 files changed, 97 insertions(+), 109 deletions(-) diff --git a/motrix_env_core/src/motrix_env_core/manager/kernel_data/__init__.py b/motrix_env_core/src/motrix_env_core/manager/kernel_data/__init__.py index 0fa68723..39c8284b 100644 --- a/motrix_env_core/src/motrix_env_core/manager/kernel_data/__init__.py +++ b/motrix_env_core/src/motrix_env_core/manager/kernel_data/__init__.py @@ -1,4 +1,5 @@ # Copyright Motphys Technology Co., Ltd. 2025, 2026 +# SPDX-License-Identifier: Apache-2.0 from motrix_env_core.manager.kernel_data.canonical import canonicalize_kernel_data from motrix_env_core.manager.kernel_data.map import Map @@ -10,12 +11,12 @@ LeafDef, MapEntryDef, TreeClassDef, - flatten_kernel_data, + clone_tree, is_kernel_data, is_kernel_data_type, - iter_tree_leaves, + iter_leaves, leaf_value_type, - unflatten_kernel_data, + tree_def_for, ) __all__ = [ @@ -29,10 +30,10 @@ "ResolvedLeaf", "ResolvedTree", "canonicalize_kernel_data", - "flatten_kernel_data", + "clone_tree", "is_kernel_data", "is_kernel_data_type", - "iter_tree_leaves", + "iter_leaves", "leaf_value_type", - "unflatten_kernel_data", + "tree_def_for", ] diff --git a/motrix_env_core/src/motrix_env_core/manager/kernel_data/canonical.py b/motrix_env_core/src/motrix_env_core/manager/kernel_data/canonical.py index fb1a2e8f..71a2770c 100644 --- a/motrix_env_core/src/motrix_env_core/manager/kernel_data/canonical.py +++ b/motrix_env_core/src/motrix_env_core/manager/kernel_data/canonical.py @@ -1,16 +1,14 @@ # Copyright Motphys Technology Co., Ltd. 2025, 2026 +# SPDX-License-Identifier: Apache-2.0 from typing import TypeVar, cast -import numpy as np -import torch - from motrix_env_core.manager.kernel_data.tree import ( - flatten_kernel_data, + clone_tree, is_kernel_data, - iter_tree_leaves, + iter_leaves, leaf_value_type, - unflatten_kernel_data, + tree_def_for, ) KernelDataType = TypeVar("KernelDataType") @@ -20,17 +18,13 @@ def canonicalize_kernel_data(value: KernelDataType, *, context: str) -> KernelDa """Copy one kernel-data value into environment-owned backing arrays.""" if not is_kernel_data(value): raise TypeError(f"{context} must return a concrete kernel-data dataclass, got {type(value).__name__}.") - leaves, tree_def = flatten_kernel_data(value) - for leaf_def in iter_tree_leaves(tree_def): + tree_def = tree_def_for(value) + for leaf_def, _leaf_value in iter_leaves(value, tree_def): if leaf_value_type(leaf_def.annotation) is None: raise TypeError( f"{context} KernelData {'.'.join(leaf_def.path)} expected a tensor, an ndarray, or a numeric scalar." ) - copied = tuple( - leaf.clone() if isinstance(leaf, torch.Tensor) else leaf.copy() if isinstance(leaf, np.ndarray) else leaf - for leaf in leaves - ) - return cast(KernelDataType, unflatten_kernel_data(tree_def, copied)) + return cast(KernelDataType, clone_tree(value, tree_def)) __all__ = ["canonicalize_kernel_data"] diff --git a/motrix_env_core/src/motrix_env_core/manager/kernel_data/tree.py b/motrix_env_core/src/motrix_env_core/manager/kernel_data/tree.py index b2f95391..49bb2008 100644 --- a/motrix_env_core/src/motrix_env_core/manager/kernel_data/tree.py +++ b/motrix_env_core/src/motrix_env_core/manager/kernel_data/tree.py @@ -4,6 +4,7 @@ from __future__ import annotations import hashlib +from collections.abc import Iterator from dataclasses import dataclass, fields, is_dataclass from typing import Any, get_origin, get_type_hints @@ -84,11 +85,11 @@ class FieldDef(BaseDef): class TreeClassDef(BaseDef): """Root definition of one concrete kernel-data class tree. - Every flattened kernel-data value is represented by a ``TreeClassDef`` root. + Every kernel-data value is described by a ``TreeClassDef`` root. Its fields recursively contain nested class, map, and leaf definitions. Attributes: - logical_type: Frozen dataclass type reconstructed by host unflattening. + logical_type: Concrete dataclass type described by this definition. fields: Field definitions in dataclass declaration order. fingerprint: Deterministic identity of the complete class tree schema. """ @@ -97,14 +98,6 @@ class TreeClassDef(BaseDef): fields: tuple[FieldDef, ...] fingerprint: str - def flatten(self, value: Any) -> tuple[Any, ...]: - """Flatten a matching class value without copying leaves.""" - return _flatten(value, self) - - def unflatten(self, leaves: tuple[Any, ...]) -> Any: - """Rebuild one class value from ordered leaves.""" - return unflatten_kernel_data(self, leaves) - @dataclass(frozen=True) class MapEntryDef(BaseDef): @@ -134,35 +127,22 @@ class KernelMapDef(BaseDef): fingerprint: str -def flatten_kernel_data(value: Any) -> tuple[tuple[Any, ...], TreeClassDef]: - """Return zero-copy leaves and the concrete class tree definition for one kernel-data value.""" +def tree_def_for(value: Any) -> TreeClassDef: + """Build and validate the class tree schema for one kernel-data value.""" if not is_kernel_data(value): raise TypeError(f"KernelData tree root must be a dataclass value, got {type(value).__name__}.") - tree_def = _build_class_def(type(value), value, (), ()) - return _flatten(value, tree_def), tree_def + return _build_class_def(type(value), value, (), ()) -def unflatten_kernel_data(tree_def: TreeClassDef, leaves: tuple[Any, ...]) -> Any: - """Rebuild a logical host value from a tree definition and ordered leaves.""" - value, next_index = _unflatten(tree_def, leaves, 0) - if next_index != len(leaves): - raise ValueError("KernelData unflatten received extra leaves.") - return value +def iter_leaves(value: Any, tree_def: BaseDef) -> Iterator[tuple[LeafDef, Any]]: + """Yield (leaf definition, leaf value) pairs in declaration order. - -def iter_tree_leaves(tree_def: BaseDef) -> tuple[LeafDef, ...]: - """Return logical leaves in stable declaration order.""" - if isinstance(tree_def, LeafDef): - return (tree_def,) - if isinstance(tree_def, TreeClassDef): - return tuple(leaf for field in tree_def.fields for leaf in iter_tree_leaves(field.tree_def)) - assert isinstance(tree_def, KernelMapDef) - return tuple(leaf for entry in tree_def.entries for leaf in iter_tree_leaves(entry.tree_def)) - - -def _flatten(value: Any, tree_def: BaseDef) -> tuple[Any, ...]: + The value is walked alongside the schema; structural mismatches (wrong + class type or changed Map keys) raise ``TypeError``. + """ if isinstance(tree_def, LeafDef): - return (value,) + yield tree_def, value + return if isinstance(tree_def, KernelMapDef): if not isinstance(value, Map): raise TypeError(f"KernelData value at {'.'.join(tree_def.path)!r} must be Map.") @@ -171,34 +151,34 @@ def _flatten(value: Any, tree_def: BaseDef) -> tuple[Any, ...]: f"Map at {'.'.join(tree_def.path)!r} has keys {value.keys_tuple!r}, expected " f"{tuple(entry.key for entry in tree_def.entries)!r}." ) - return tuple(leaf for entry in tree_def.entries for leaf in _flatten(value[entry.key], entry.tree_def)) + for entry in tree_def.entries: + yield from iter_leaves(value[entry.key], entry.tree_def) + return assert isinstance(tree_def, TreeClassDef) if type(value) is not tree_def.logical_type: raise TypeError(f"KernelData value must be {tree_def.logical_type.__name__}, got {type(value).__name__}.") - return tuple(leaf for field in tree_def.fields for leaf in _flatten(getattr(value, field.name), field.tree_def)) + for field in tree_def.fields: + yield from iter_leaves(getattr(value, field.name), field.tree_def) + +def clone_tree(value: Any, tree_def: BaseDef) -> Any: + """Deep-clone one kernel-data value following its tree schema. -def _unflatten( - tree_def: BaseDef, - leaves: tuple[Any, ...], - index: int, -) -> tuple[Any, int]: + Tensor leaves are cloned, ndarray leaves are copied, and scalar leaves are + reused as-is; dataclass fields and Map entries are rebuilt around the + cloned children. + """ if isinstance(tree_def, LeafDef): - if index >= len(leaves): - raise ValueError("KernelData unflatten received too few leaves.") - return leaves[index], index + 1 + if isinstance(value, torch.Tensor): + return value.clone() + if isinstance(value, np.ndarray): + return value.copy() + return value if isinstance(tree_def, KernelMapDef): - entries: dict[str, Any] = {} - for entry in tree_def.entries: - value, index = _unflatten(entry.tree_def, leaves, index) - entries[entry.key] = value - return Map(entries), index + return Map({entry.key: clone_tree(value[entry.key], entry.tree_def) for entry in tree_def.entries}) assert isinstance(tree_def, TreeClassDef) - values = [] - for field in tree_def.fields: - value, index = _unflatten(field.tree_def, leaves, index) - values.append(value) - return tree_def.logical_type(*values), index + cloned = tuple(clone_tree(getattr(value, field.name), field.tree_def) for field in tree_def.fields) + return tree_def.logical_type(*cloned) def _build_class_def( diff --git a/motrix_env_core/src/motrix_env_core/manager/resolver.py b/motrix_env_core/src/motrix_env_core/manager/resolver.py index f7e3137f..2085b009 100644 --- a/motrix_env_core/src/motrix_env_core/manager/resolver.py +++ b/motrix_env_core/src/motrix_env_core/manager/resolver.py @@ -36,10 +36,10 @@ ResolvedLeaf, ResolvedTree, TreeClassDef, - flatten_kernel_data, is_kernel_data, - iter_tree_leaves, + iter_leaves, leaf_value_type, + tree_def_for, ) from motrix_env_core.manager.kernel_data.methods import dispatch_methods from motrix_env_core.manager.kernel_data.tree import BaseDef, KernelMapDef, LeafDef @@ -555,7 +555,7 @@ def _resolve_reset( def _resolve_runtime_terms(self) -> None: for name, action_term in self._env.action_terms.items(): - _, tree_def = flatten_kernel_data(action_term.state) + tree_def = tree_def_for(action_term.state) self._layout_identity.append( ( "action", @@ -569,7 +569,7 @@ def _resolve_runtime_terms(self) -> None: ) ) for name, command_term in self._env.command_terms.items(): - _, tree_def = flatten_kernel_data(command_term) + tree_def = tree_def_for(command_term) self._layout_identity.append( ( "command", @@ -805,9 +805,9 @@ def _register_tree( by their owning runtime objects and scalar/shared leaves are reused unchanged. """ - values, tree_def = flatten_kernel_data(value) + tree_def = tree_def_for(value) leaves: dict[tuple[str, ...], ResolvedLeaf] = {} - for leaf_def, field_value in zip(iter_tree_leaves(tree_def), values, strict=True): + for leaf_def, leaf_value in iter_leaves(value, tree_def): value_type = leaf_value_type(leaf_def.annotation) if value_type is None: raise TypeError( @@ -815,14 +815,14 @@ def _register_tree( "scalar." ) if value_type in (np.ndarray, torch.Tensor): - if not isinstance(field_value, value_type): + if not isinstance(leaf_value, value_type): raise TypeError(f"{source_name} field {'.'.join(leaf_def.path)!r} must be {value_type.__name__}.") - elif type(field_value) is not value_type: + elif type(leaf_value) is not value_type: raise TypeError( f"{source_name} field {'.'.join(leaf_def.path)!r} must return " - f"{value_type.__name__}, got {type(field_value).__name__}." + f"{value_type.__name__}, got {type(leaf_value).__name__}." ) - leaves[leaf_def.path] = ResolvedLeaf(field_value) + leaves[leaf_def.path] = ResolvedLeaf(leaf_value) self._prepared_trees.append(ResolvedTree(value, leaves)) return len(self._prepared_trees) - 1, tree_def diff --git a/motrix_env_core/tests/test_kernel_data.py b/motrix_env_core/tests/test_kernel_data.py index 84aa6ff1..f0c2a89e 100644 --- a/motrix_env_core/tests/test_kernel_data.py +++ b/motrix_env_core/tests/test_kernel_data.py @@ -1,4 +1,5 @@ # Copyright Motphys Technology Co., Ltd. 2025, 2026 +# SPDX-License-Identifier: Apache-2.0 from dataclasses import FrozenInstanceError, dataclass @@ -9,11 +10,11 @@ from motrix_env_core.manager.kernel_data import ( Map, canonicalize_kernel_data, - flatten_kernel_data, + clone_tree, is_kernel_data, - iter_tree_leaves, + iter_leaves, leaf_value_type, - unflatten_kernel_data, + tree_def_for, ) @@ -52,42 +53,47 @@ def test_kernel_data_is_frozen_and_compiles_nested_inherited_leaves() -> None: with pytest.raises(FrozenInstanceError): value.phase = np.float32(1.0) # type: ignore[misc] - _, tree_def = flatten_kernel_data(value) - leaves = iter_tree_leaves(tree_def) - assert [leaf.path for leaf in leaves] == [ + leaves = [(leaf.path, leaf_value) for leaf, leaf_value in iter_leaves(value, tree_def_for(value))] + assert [path for path, _ in leaves] == [ ("robot", "velocity"), ("robot", "pose", "position"), ("robot", "pose", "orientation"), ("phase",), ] - assert leaf_value_type(leaves[0].annotation) is np.ndarray - assert leaf_value_type(leaves[-1].annotation) is np.float32 + annotations = [leaf.annotation for leaf, _ in iter_leaves(value, tree_def_for(value))] + assert leaf_value_type(annotations[0]) is np.ndarray + assert leaf_value_type(annotations[-1]) is np.float32 -def test_kernel_data_flatten_is_zero_copy_and_rebuilds_host_value() -> None: +def test_kernel_data_clone_copies_leaves_and_rebuilds_host_value() -> None: velocity = np.asarray([[1.0]], dtype=np.float32) position = np.asarray([[2.0]], dtype=np.float32) orientation = np.asarray([3.0], dtype=np.float32) value = _Tracking(_Robot(velocity, _Pose(position, orientation)), np.float32(0.25)) - leaves, tree_def = flatten_kernel_data(value) - assert leaves[:3] == (velocity, position, orientation) - rebuilt = unflatten_kernel_data(tree_def, leaves) - canonicalize_kernel_data(value, context="test") - assert rebuilt.robot.pose.position is position - assert rebuilt.robot.pose.orientation is orientation - assert rebuilt.phase == np.float32(0.25) + cloned = clone_tree(value, tree_def_for(value)) + assert cloned.robot.velocity is not velocity + assert not np.shares_memory(cloned.robot.velocity, velocity) + assert np.array_equal(cloned.robot.velocity, velocity) + assert cloned.robot.pose.position is not position + assert np.array_equal(cloned.robot.pose.position, position) + assert cloned.phase is value.phase + assert cloned == value -def test_kernel_data_accepts_tensor_leaves() -> None: +def test_kernel_data_accepts_tensor_leaves_and_clones_them() -> None: @dataclass(frozen=True, slots=True) class _Tensors: buffer: torch.Tensor buffer = torch.zeros((2, 3), dtype=torch.float32) value = _Tensors(buffer) - leaves, tree_def = flatten_kernel_data(value) - assert leaves == (buffer,) - assert leaf_value_type(iter_tree_leaves(tree_def)[0].annotation) is torch.Tensor + leaves = list(iter_leaves(value, tree_def_for(value))) + assert [leaf_value for _, leaf_value in leaves] == [buffer] + assert leaf_value_type(leaves[0][0].annotation) is torch.Tensor + cloned = clone_tree(value, tree_def_for(value)) + assert cloned.buffer is not buffer + assert cloned.buffer.data_ptr() != buffer.data_ptr() + assert torch.equal(cloned.buffer, buffer) def test_kernel_data_schema_fingerprint_ignores_values() -> None: @@ -99,8 +105,8 @@ def test_kernel_data_schema_fingerprint_ignores_values() -> None: _Robot(np.asarray([[4.0]], dtype=np.float32), _Pose(np.asarray([[5.0]]), np.asarray([6.0]))), np.float32(0.75), ) - _, first = flatten_kernel_data(first_value) - _, second = flatten_kernel_data(second_value) + first = tree_def_for(first_value) + second = tree_def_for(second_value) assert first.fingerprint == second.fingerprint assert first.logical_type.__module__ == _Tracking.__module__ @@ -108,14 +114,20 @@ def test_kernel_data_schema_fingerprint_ignores_values() -> None: def test_map_schema_changes_tree_identity() -> None: first = _MappedValues(Map({"value": np.float32(1.0)})) second = _MappedValues(Map({"value": np.asarray([[1.0]], dtype=np.float32)})) - _, first_tree = flatten_kernel_data(first) - _, second_tree = flatten_kernel_data(second) + first_tree = tree_def_for(first) + second_tree = tree_def_for(second) assert first_tree.fingerprint != second_tree.fingerprint first_map = first_tree.fields[0].tree_def second_map = second_tree.fields[0].tree_def assert first_map.fingerprint != second_map.fingerprint +def test_iter_leaves_rejects_changed_map_keys() -> None: + schema = tree_def_for(_MappedValues(Map({"a": np.float32(1.0), "b": np.float32(2.0)}))) + with pytest.raises(TypeError, match="has keys"): + list(iter_leaves(_MappedValues(Map({"a": np.float32(1.0)})), schema)) + + def test_kernel_data_rejects_unsupported_leaf() -> None: @dataclass(frozen=True, slots=True) class _Invalid: diff --git a/wiki/design/manager/kernel-abi.md b/wiki/design/manager/kernel-abi.md index 229244de..dac51c34 100644 --- a/wiki/design/manager/kernel-abi.md +++ b/wiki/design/manager/kernel-abi.md @@ -67,11 +67,12 @@ dof_pos = ctx.sim["robot_dof_pos"] ## 4. Logical tree 与 resolved tree Kernel-data record 是 plain frozen dataclass(`@dataclass(frozen=True, slots=True)`, -不再有专用装饰器),支持有序、可逆的 record tree: +不再有专用装饰器),支持按声明序遍历与深拷贝的 record tree: ```python -leaves, tree_def = flatten_kernel_data(value) -rebuilt = unflatten_kernel_data(tree_def, leaves) +tree_def = tree_def_for(value) +leaves = iter_leaves(value, tree_def) # (LeafDef, value) 声明序 +cloned = clone_tree(value, tree_def) ``` Record 可以嵌套另一个 kernel-data record;不支持递归结构、dict、list、字符串、object array 或 backend handle From e20b3155378816ae2e22c0df4af8dd60ad162109 Mon Sep 17 00:00:00 2001 From: zilch <147668916@qq.com> Date: Fri, 6 Nov 2026 22:50:28 +0800 Subject: [PATCH 35/41] refactor(kernel_data)!: consume cfg-provided carriers by reference 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. --- .../manager/kernel_data/__init__.py | 2 -- .../manager/kernel_data/canonical.py | 12 ++++--- .../manager/kernel_data/tree.py | 20 ----------- motrix_env_core/tests/test_kernel_data.py | 22 +----------- motrix_env_core/tests/test_manager_env.py | 35 +++++++++++++++++++ wiki/design/manager/kernel-abi.md | 6 ++-- 6 files changed, 48 insertions(+), 49 deletions(-) diff --git a/motrix_env_core/src/motrix_env_core/manager/kernel_data/__init__.py b/motrix_env_core/src/motrix_env_core/manager/kernel_data/__init__.py index 39c8284b..53c13b3b 100644 --- a/motrix_env_core/src/motrix_env_core/manager/kernel_data/__init__.py +++ b/motrix_env_core/src/motrix_env_core/manager/kernel_data/__init__.py @@ -11,7 +11,6 @@ LeafDef, MapEntryDef, TreeClassDef, - clone_tree, is_kernel_data, is_kernel_data_type, iter_leaves, @@ -30,7 +29,6 @@ "ResolvedLeaf", "ResolvedTree", "canonicalize_kernel_data", - "clone_tree", "is_kernel_data", "is_kernel_data_type", "iter_leaves", diff --git a/motrix_env_core/src/motrix_env_core/manager/kernel_data/canonical.py b/motrix_env_core/src/motrix_env_core/manager/kernel_data/canonical.py index 71a2770c..68c619ba 100644 --- a/motrix_env_core/src/motrix_env_core/manager/kernel_data/canonical.py +++ b/motrix_env_core/src/motrix_env_core/manager/kernel_data/canonical.py @@ -1,10 +1,9 @@ # Copyright Motphys Technology Co., Ltd. 2025, 2026 # SPDX-License-Identifier: Apache-2.0 -from typing import TypeVar, cast +from typing import TypeVar from motrix_env_core.manager.kernel_data.tree import ( - clone_tree, is_kernel_data, iter_leaves, leaf_value_type, @@ -15,7 +14,12 @@ def canonicalize_kernel_data(value: KernelDataType, *, context: str) -> KernelDataType: - """Copy one kernel-data value into environment-owned backing arrays.""" + """Validate one kernel-data value against the cfg→term read-only contract. + + Cfg-produced values are consumed as-is: per-env mutable state must be + freshly allocated inside the producing ``cfg.__call__``; only the leaf + schema is checked here. + """ if not is_kernel_data(value): raise TypeError(f"{context} must return a concrete kernel-data dataclass, got {type(value).__name__}.") tree_def = tree_def_for(value) @@ -24,7 +28,7 @@ def canonicalize_kernel_data(value: KernelDataType, *, context: str) -> KernelDa raise TypeError( f"{context} KernelData {'.'.join(leaf_def.path)} expected a tensor, an ndarray, or a numeric scalar." ) - return cast(KernelDataType, clone_tree(value, tree_def)) + return value __all__ = ["canonicalize_kernel_data"] diff --git a/motrix_env_core/src/motrix_env_core/manager/kernel_data/tree.py b/motrix_env_core/src/motrix_env_core/manager/kernel_data/tree.py index 49bb2008..5ea12e09 100644 --- a/motrix_env_core/src/motrix_env_core/manager/kernel_data/tree.py +++ b/motrix_env_core/src/motrix_env_core/manager/kernel_data/tree.py @@ -161,26 +161,6 @@ class type or changed Map keys) raise ``TypeError``. yield from iter_leaves(getattr(value, field.name), field.tree_def) -def clone_tree(value: Any, tree_def: BaseDef) -> Any: - """Deep-clone one kernel-data value following its tree schema. - - Tensor leaves are cloned, ndarray leaves are copied, and scalar leaves are - reused as-is; dataclass fields and Map entries are rebuilt around the - cloned children. - """ - if isinstance(tree_def, LeafDef): - if isinstance(value, torch.Tensor): - return value.clone() - if isinstance(value, np.ndarray): - return value.copy() - return value - if isinstance(tree_def, KernelMapDef): - return Map({entry.key: clone_tree(value[entry.key], entry.tree_def) for entry in tree_def.entries}) - assert isinstance(tree_def, TreeClassDef) - cloned = tuple(clone_tree(getattr(value, field.name), field.tree_def) for field in tree_def.fields) - return tree_def.logical_type(*cloned) - - def _build_class_def( logical_type: type, value: Any | None, diff --git a/motrix_env_core/tests/test_kernel_data.py b/motrix_env_core/tests/test_kernel_data.py index f0c2a89e..555e787e 100644 --- a/motrix_env_core/tests/test_kernel_data.py +++ b/motrix_env_core/tests/test_kernel_data.py @@ -10,7 +10,6 @@ from motrix_env_core.manager.kernel_data import ( Map, canonicalize_kernel_data, - clone_tree, is_kernel_data, iter_leaves, leaf_value_type, @@ -65,22 +64,7 @@ def test_kernel_data_is_frozen_and_compiles_nested_inherited_leaves() -> None: assert leaf_value_type(annotations[-1]) is np.float32 -def test_kernel_data_clone_copies_leaves_and_rebuilds_host_value() -> None: - velocity = np.asarray([[1.0]], dtype=np.float32) - position = np.asarray([[2.0]], dtype=np.float32) - orientation = np.asarray([3.0], dtype=np.float32) - value = _Tracking(_Robot(velocity, _Pose(position, orientation)), np.float32(0.25)) - cloned = clone_tree(value, tree_def_for(value)) - assert cloned.robot.velocity is not velocity - assert not np.shares_memory(cloned.robot.velocity, velocity) - assert np.array_equal(cloned.robot.velocity, velocity) - assert cloned.robot.pose.position is not position - assert np.array_equal(cloned.robot.pose.position, position) - assert cloned.phase is value.phase - assert cloned == value - - -def test_kernel_data_accepts_tensor_leaves_and_clones_them() -> None: +def test_kernel_data_accepts_tensor_leaves() -> None: @dataclass(frozen=True, slots=True) class _Tensors: buffer: torch.Tensor @@ -90,10 +74,6 @@ class _Tensors: leaves = list(iter_leaves(value, tree_def_for(value))) assert [leaf_value for _, leaf_value in leaves] == [buffer] assert leaf_value_type(leaves[0][0].annotation) is torch.Tensor - cloned = clone_tree(value, tree_def_for(value)) - assert cloned.buffer is not buffer - assert cloned.buffer.data_ptr() != buffer.data_ptr() - assert torch.equal(cloned.buffer, buffer) def test_kernel_data_schema_fingerprint_ignores_values() -> None: diff --git a/motrix_env_core/tests/test_manager_env.py b/motrix_env_core/tests/test_manager_env.py index 545be58f..ee81f9df 100644 --- a/motrix_env_core/tests/test_manager_env.py +++ b/motrix_env_core/tests/test_manager_env.py @@ -334,6 +334,41 @@ def test_action_and_command_terms_are_environment_owned() -> None: assert _counter_command(env) is env.command_terms["counter"] +def test_same_cfg_class_instances_keep_independent_runtime_state() -> None: + """Cfg-produced state must be freshly allocated per environment instance. + + The framework consumes cfg values by reference (no defensive clone), so + two environments built from the same cfg class must not share any mutable + per-env buffer: their action queues and command state stay distinct + objects and interleaved stepping only reflects each environment's own + actions and resets. + """ + first = _ManagerEnv(num_envs=3) + second = _ManagerEnv(num_envs=3) + first.init_state() + second.init_state() + + first_queue = first.action_terms["test"].state.action_queue + second_queue = second.action_terms["test"].state.action_queue + assert first_queue.data_ptr() != second_queue.data_ptr() + assert _counter_command(first).command.data_ptr() != _counter_command(second).command.data_ptr() + assert _counter_command(first).double.data_ptr() != _counter_command(second).double.data_ptr() + + dt = float(first.cfg.ctrl_dt) + first_state = first.step(torch.full((3, 1), 0.1, dtype=torch.float32)) + second_state = second.step(torch.full((3, 1), -0.1, dtype=torch.float32)) + first.step(torch.full((3, 1), 0.2, dtype=torch.float32)) + second.step(torch.full((3, 1), -0.2, dtype=torch.float32)) + del first_state, second_state + first.reset(np.asarray([0])) + second.reset(np.asarray([0])) + + np.testing.assert_allclose(_counter_command(first).double.numpy()[:, 0], [0.4, 0.4, 0.4]) + np.testing.assert_allclose(_counter_command(second).double.numpy()[:, 0], [-0.4, -0.4, -0.4]) + np.testing.assert_allclose(_counter_command(first).command.numpy()[:, 0], [-1.0, -1.0 + 2 * dt, -1.0 + 2 * dt]) + np.testing.assert_allclose(_counter_command(second).command.numpy()[:, 0], [-1.0, -1.0 + 2 * dt, -1.0 + 2 * dt]) + + def test_command_rematerialization_reuses_transition_kernel_inputs(monkeypatch: pytest.MonkeyPatch) -> None: env = _ManagerEnv(num_envs=3) state = env.init_state() diff --git a/wiki/design/manager/kernel-abi.md b/wiki/design/manager/kernel-abi.md index dac51c34..02008f1f 100644 --- a/wiki/design/manager/kernel-abi.md +++ b/wiki/design/manager/kernel-abi.md @@ -67,14 +67,16 @@ dof_pos = ctx.sim["robot_dof_pos"] ## 4. Logical tree 与 resolved tree Kernel-data record 是 plain frozen dataclass(`@dataclass(frozen=True, slots=True)`, -不再有专用装饰器),支持按声明序遍历与深拷贝的 record tree: +不再有专用装饰器),支持按声明序遍历的 record tree: ```python tree_def = tree_def_for(value) leaves = iter_leaves(value, tree_def) # (LeafDef, value) 声明序 -cloned = clone_tree(value, tree_def) ``` +Cfg 产物按只读契约由框架按引用消费:`canonicalize_kernel_data` 只做 leaf 类型校验, +per-env 可变 state 必须在 `cfg.__call__` 内按 `num_envs` 新分配。 + Record 可以嵌套另一个 kernel-data record;不支持递归结构、dict、list、字符串、object array 或 backend handle (SimDataQuery 是 dataclass 但通过 read program 解析,不属于 kernel-data tree)。 From 51e5f85e996c509600546d7ae2e2bea4409f5870 Mon Sep 17 00:00:00 2001 From: zilch <147668916@qq.com> Date: Fri, 6 Nov 2026 23:07:00 +0800 Subject: [PATCH 36/41] refactor(manager)!: remove the unused warmup execution path 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. --- .../src/motrix_env_core/manager/__init__.py | 2 - .../src/motrix_env_core/manager/env.py | 80 ++++++------- .../src/motrix_env_core/manager/kernel.py | 109 ++++-------------- motrix_env_core/tests/test_kernel_clone.py | 47 -------- motrix_env_core/tests/test_manager_env.py | 21 +++- 5 files changed, 77 insertions(+), 182 deletions(-) delete mode 100644 motrix_env_core/tests/test_kernel_clone.py diff --git a/motrix_env_core/src/motrix_env_core/manager/__init__.py b/motrix_env_core/src/motrix_env_core/manager/__init__.py index 0bf870fa..d63b2661 100644 --- a/motrix_env_core/src/motrix_env_core/manager/__init__.py +++ b/motrix_env_core/src/motrix_env_core/manager/__init__.py @@ -14,7 +14,6 @@ from motrix_env_core.manager.context import BuildContext, ManagerContext from motrix_env_core.manager.dispatch import dispatch from motrix_env_core.manager.env import ManagerBasedEnvCfg, ManagerEnv -from motrix_env_core.manager.kernel import ManagerWarmupResult from motrix_env_core.manager.metrics import metric from motrix_env_core.manager.observations import ( ManagerObservationGroupCfg, @@ -51,7 +50,6 @@ "ManagerRewardsCfg", "ManagerTerminationsCfg", "SimQueriesCfg", - "ManagerWarmupResult", "ObsTerm", "ObservationTermCfg", "BaseTerm", diff --git a/motrix_env_core/src/motrix_env_core/manager/env.py b/motrix_env_core/src/motrix_env_core/manager/env.py index 4ef9f9e7..a961dd16 100644 --- a/motrix_env_core/src/motrix_env_core/manager/env.py +++ b/motrix_env_core/src/motrix_env_core/manager/env.py @@ -22,9 +22,6 @@ from motrix_env_core.manager.commands import CommandCfg, CommandTerm, ManagerCommandsCfg, ResetContext from motrix_env_core.manager.kernel import ( KernelOutputs, - ManagerWarmupResult, - clone_kernel_value, - clone_resolved_trees, validate_kernel_context, ) from motrix_env_core.manager.kernel_data import ( @@ -68,6 +65,25 @@ logger = logging.getLogger(__name__) +def _env_ids_to_numpy(env_ids: torch.Tensor, label: str) -> np.ndarray: + """Convert a public reset id tensor once to the internal numpy plane. + + The manager frontend's public API is uniformly torch; passing a numpy + array here is an API misuse, not a supported input. + """ + if not isinstance(env_ids, torch.Tensor): + raise TypeError( + f"ManagerEnv.reset {label} must be a torch integer tensor, got {type(env_ids).__name__}; " + "use torch.tensor([...], dtype=torch.long)." + ) + if env_ids.dtype not in (torch.int8, torch.int16, torch.int32, torch.int64, torch.uint8): + raise TypeError( + f"ManagerEnv.reset {label} must have an integer dtype, got {env_ids.dtype}; " + "use torch.tensor([...], dtype=torch.long)." + ) + return env_ids.detach().cpu().numpy() + + def _merge_query_declarations( declared: dict[str, _Q], required: Iterable[tuple[str, str, _Q]], @@ -490,8 +506,10 @@ def _capture_final_obs(self, env_ids: np.ndarray) -> FinalObs[torch.Tensor]: slot = self._final_obs_slot self._task_program.observe_kernel(self._resolved_inputs, slot.kernel_outputs) obs = slot.observation + # ``env_ids`` is user-facing output of the torch frontend: hand over + # an owned tensor while internal indexing stays on the numpy copy. return FinalObs( - env_ids=env_ids, + env_ids=torch.from_numpy(env_ids.copy()), obs=EnvObs( policy=obs.policy[env_ids], value=None if obs.value is None else obs.value[env_ids], @@ -644,11 +662,15 @@ def _reset_done_envs(self) -> None: with self.perf.scope("select_done"): np.putmask(state.episode_steps, done, 0) with self.perf.scope("reset_envs"): - self.reset(env_ids, sim_reset_ids) + self._reset_rows(env_ids, sim_reset_ids) - def reset(self, env_ids: np.ndarray, sim_reset_ids: np.ndarray | None = None) -> None: + def reset(self, env_ids: torch.Tensor, sim_reset_ids: torch.Tensor | None = None) -> None: """Reset episode lanes fully and sim-reset the rest in one pass. + ``env_ids`` and ``sim_reset_ids`` are integer torch tensors (the + manager frontend's public currency); they are converted once here to + the internal numpy control plane. + Episode lanes (``env_ids``) run the host lifecycle resets (command ``reset`` with termination statistics, action resets); sim-reset lanes (``ctx.sim_reset_requested`` set mid-transition) skip them — with no @@ -657,6 +679,14 @@ def reset(self, env_ids: np.ndarray, sim_reset_ids: np.ndarray | None = None) -> configured sim reset terms rewrite simulator rows, and kernel inputs are reread once for the union. """ + self._reset_rows( + _env_ids_to_numpy(env_ids, "env_ids"), + None if sim_reset_ids is None else _env_ids_to_numpy(sim_reset_ids, "sim_reset_ids"), + ) + + def _reset_rows(self, env_ids: np.ndarray, sim_reset_ids: np.ndarray | None = None) -> None: + """Numpy-control-plane reset shared by the public wrapper and the + internal auto-reset flow.""" assert self._state is not None if env_ids.size: with self.perf.scope("command_reset"): @@ -799,41 +829,3 @@ def _validate_kernel_context(self, inputs: tuple[ResolvedTree, ...]) -> None: message = f"Invalid kernel context for {type(self).__name__} task kernel: {error}" raise type(error)(message) from error - def warmup(self) -> ManagerWarmupResult: - """Run one isolated evaluate/observe pass to exercise the executor graphs. - - Inputs and buffers are cloned so the warmup cannot perturb live state; - the Torch executor compiles lazily on first real use, so only the - first-execution time is reported. - """ - if self._state is None: - self.init_state() - assert self._state is not None - assert self._task_program is not None - assert self._compiled_manager_program is not None - assert self._kernel_buffers is not None - assert self._kernel_outputs is not None - - self._refresh_sim_reads() - inputs = self._resolved_inputs - self._validate_kernel_context(inputs) - import torch - - started = perf_counter() - warmup_inputs = clone_resolved_trees(inputs) - self._task_program.evaluate_kernel( - warmup_inputs, - self._task_program.reward_weights, - clone_kernel_value(self._kernel_buffers), - clone_kernel_value(self._kernel_outputs), - ) - self._task_program.observe_kernel(warmup_inputs, clone_kernel_value(self._kernel_outputs)) - return ManagerWarmupResult( - compile_seconds=0.0, - signatures=("torch eager",), - threading_layer="openmp", - num_threads=torch.get_num_threads(), - term_compile_seconds=0.0, - kernel_compile_seconds=0.0, - first_execution_seconds=perf_counter() - started, - ) diff --git a/motrix_env_core/src/motrix_env_core/manager/kernel.py b/motrix_env_core/src/motrix_env_core/manager/kernel.py index cb75e8a6..7170320b 100644 --- a/motrix_env_core/src/motrix_env_core/manager/kernel.py +++ b/motrix_env_core/src/motrix_env_core/manager/kernel.py @@ -1,14 +1,11 @@ # Copyright Motphys Technology Co., Ltd. 2025, 2026 # SPDX-License-Identifier: Apache-2.0 -from dataclasses import dataclass -from typing import Any, NamedTuple, TypeVar, cast +from typing import Any, NamedTuple, TypeVar import numpy as np import torch -from motrix_env_core.manager.kernel_data.resolved import ResolvedLeaf, ResolvedTree - T = TypeVar("T") @@ -16,90 +13,11 @@ class KernelOutputs(NamedTuple): """Preallocated tensors written by one manager task kernel.""" policy_obs: torch.Tensor - value_obs: torch.Tensor + value_obs: torch.Tensor | None reward: torch.Tensor terminated: torch.Tensor -@dataclass(frozen=True) -class ManagerWarmupResult: - """Compilation and first-execution diagnostics produced by :meth:`ManagerEnv.warmup`.""" - - compile_seconds: float - signatures: tuple[str, ...] - threading_layer: str | None - num_threads: int - term_compile_seconds: float = 0.0 - kernel_compile_seconds: float = 0.0 - first_execution_seconds: float = 0.0 - - -def validate_kernel_context( - trees: Any, - reward_weights: torch.Tensor, - buffers: Any, - outputs: KernelOutputs, - num_envs: int, -) -> None: - """Validate the numeric kernel contract once when a context is bound. - - ``trees`` is the resolved program input: one :class:`ResolvedTree` per - prepared kernel-data source. Leaf values are checked for dtype and type - sanity without any batch-dimension assumptions. - """ - if not isinstance(trees, tuple): - raise TypeError(f"inputs must be a tuple of resolved trees, got {type(trees).__name__}.") - for tree in trees: - for path, leaf in tree.leaves.items(): - _validate_value(leaf.value, f"inputs.{'.'.join(path)}", num_envs=num_envs, batched=False, writable=False) - _validate_value(reward_weights, "reward_weights", num_envs=num_envs, batched=False, writable=False) - _validate_value(buffers, "buffers", num_envs=num_envs, batched=True, writable=True) - _validate_outputs(outputs, num_envs) - - -def clone_resolved_trees(trees: tuple[ResolvedTree, ...]) -> tuple[ResolvedTree, ...]: - """Clone every leaf of the resolved trees for isolated warmup execution. - - Leaves are cloned with :func:`clone_kernel_value` semantics (dtype, - strides, and read-only flag preserved); the host root objects stay shared - exactly as in real execution, where namespace fallback reads immutable - config attributes from the original records. - """ - return tuple( - ResolvedTree( - tree.value, - {path: ResolvedLeaf(clone_kernel_value(leaf.value)) for path, leaf in tree.leaves.items()}, - ) - for tree in trees - ) - - -def clone_kernel_value(value: T) -> T: - """Recursively clone arrays and tensors while preserving task-specific tuple types.""" - if isinstance(value, torch.Tensor): - return cast(T, value.clone()) - if isinstance(value, np.ndarray): - # Warmup must use the same Numba array type as real execution: copy() - # normalizes strided views to C layout and drops the readonly flag. - lower = ( - sum(min(0, (size - 1) * stride) for size, stride in zip(value.shape, value.strides)) if value.size else 0 - ) - upper = ( - sum(max(0, (size - 1) * stride) for size, stride in zip(value.shape, value.strides)) if value.size else 0 - ) - storage = np.empty(max(value.itemsize, upper - lower + value.itemsize), dtype=np.uint8) - copied = np.ndarray(value.shape, dtype=value.dtype, buffer=storage, offset=-lower, strides=value.strides) - copied[...] = value - copied.flags.writeable = value.flags.writeable - return cast(T, copied) - if isinstance(value, tuple): - items = [clone_kernel_value(item) for item in value] - if hasattr(type(value), "_fields"): - return cast(T, type(value)(*items)) - return cast(T, tuple(items)) - return value - - def _validate_value(value: Any, path: str, *, num_envs: int, batched: bool, writable: bool) -> None: if isinstance(value, torch.Tensor): if batched and (value.dim() == 0 or value.shape[0] != num_envs): @@ -146,3 +64,26 @@ def _validate_outputs(outputs: KernelOutputs, num_envs: int) -> None: raise ValueError(f"outputs.terminated must have shape ({num_envs},), got {tuple(outputs.terminated.shape)}.") if outputs.terminated.dtype is not torch.bool: raise TypeError(f"outputs.terminated must use bool dtype, got {outputs.terminated.dtype}.") + + +def validate_kernel_context( + trees: Any, + reward_weights: torch.Tensor, + buffers: Any, + outputs: KernelOutputs, + num_envs: int, +) -> None: + """Validate the numeric kernel contract once when a context is bound. + + ``trees`` is the resolved program input: one :class:`ResolvedTree` per + prepared kernel-data source. Leaf values are checked for dtype and type + sanity without any batch-dimension assumptions. + """ + if not isinstance(trees, tuple): + raise TypeError(f"inputs must be a tuple of resolved trees, got {type(trees).__name__}.") + for tree in trees: + for path, leaf in tree.leaves.items(): + _validate_value(leaf.value, f"inputs.{'.'.join(path)}", num_envs=num_envs, batched=False, writable=False) + _validate_value(reward_weights, "reward_weights", num_envs=num_envs, batched=False, writable=False) + _validate_value(buffers, "buffers", num_envs=num_envs, batched=True, writable=True) + _validate_outputs(outputs, num_envs) diff --git a/motrix_env_core/tests/test_kernel_clone.py b/motrix_env_core/tests/test_kernel_clone.py deleted file mode 100644 index 1d6f502a..00000000 --- a/motrix_env_core/tests/test_kernel_clone.py +++ /dev/null @@ -1,47 +0,0 @@ -# Copyright Motphys Technology Co., Ltd. 2025, 2026 -# SPDX-License-Identifier: Apache-2.0 - -from collections import namedtuple - -import numpy as np -import pytest - -from motrix_env_core.manager.kernel import clone_kernel_value - - -@pytest.mark.parametrize("layout", ["c", "f", "strided", "reversed", "broadcast", "empty", "empty_strided"]) -@pytest.mark.parametrize("readonly", [False, True]) -def test_kernel_clone_preserves_array_type_and_isolates_storage(layout, readonly): - source = np.arange(48, dtype=np.float32).reshape(8, 6) - if layout == "f": - source = np.asfortranarray(source) - elif layout == "strided": - source = source[::2, ::2] - elif layout == "reversed": - source = source[::-1, ::-2] - elif layout == "broadcast": - source = np.broadcast_to(source[:1], (8, 6)) - elif layout == "empty": - source = source[:0] - elif layout == "empty_strided": - source = np.ndarray((0, 2), dtype=np.float32, buffer=bytearray(4), strides=(0, 2**62)) - source.flags.writeable = not readonly - copied = clone_kernel_value(source) - - assert copied.strides == source.strides - assert copied.dtype == source.dtype - assert not np.shares_memory(copied, source) - np.testing.assert_array_equal(copied, source) - if not readonly: - copied[...] = -1 - assert np.all(source >= 0) - - -def test_kernel_clone_preserves_nested_namedtuple_and_scalar_values(): - outputs = namedtuple("Outputs", ["value", "count"]) - source = (outputs(np.arange(4, dtype=np.float32), 3), np.float32(0.5)) - copied = clone_kernel_value(source) - assert type(copied[0]) is outputs - assert copied[0].count == 3 - assert copied[1] == np.float32(0.5) - assert not np.shares_memory(source[0].value, copied[0].value) diff --git a/motrix_env_core/tests/test_manager_env.py b/motrix_env_core/tests/test_manager_env.py index ee81f9df..03737f98 100644 --- a/motrix_env_core/tests/test_manager_env.py +++ b/motrix_env_core/tests/test_manager_env.py @@ -360,8 +360,8 @@ def test_same_cfg_class_instances_keep_independent_runtime_state() -> None: first.step(torch.full((3, 1), 0.2, dtype=torch.float32)) second.step(torch.full((3, 1), -0.2, dtype=torch.float32)) del first_state, second_state - first.reset(np.asarray([0])) - second.reset(np.asarray([0])) + first.reset(torch.tensor([0], dtype=torch.long)) + second.reset(torch.tensor([0], dtype=torch.long)) np.testing.assert_allclose(_counter_command(first).double.numpy()[:, 0], [0.4, 0.4, 0.4]) np.testing.assert_allclose(_counter_command(second).double.numpy()[:, 0], [-0.4, -0.4, -0.4]) @@ -369,6 +369,19 @@ def test_same_cfg_class_instances_keep_independent_runtime_state() -> None: np.testing.assert_allclose(_counter_command(second).command.numpy()[:, 0], [-1.0, -1.0 + 2 * dt, -1.0 + 2 * dt]) +def test_reset_rejects_numpy_ids_and_non_integer_dtype() -> None: + """The manager frontend's public reset currency is uniformly torch.""" + env = _ManagerEnv(num_envs=3) + env.init_state() + + with pytest.raises(TypeError, match="torch.tensor"): + env.reset(np.asarray([0])) + with pytest.raises(TypeError, match="torch.tensor"): + env.reset(torch.tensor([0]), sim_reset_ids=np.asarray([1])) + with pytest.raises(TypeError, match="integer dtype"): + env.reset(torch.tensor([0.0])) + + def test_command_rematerialization_reuses_transition_kernel_inputs(monkeypatch: pytest.MonkeyPatch) -> None: env = _ManagerEnv(num_envs=3) state = env.init_state() @@ -590,9 +603,6 @@ def test_manager_context_is_injected_once_and_reused_across_all_term_kinds() -> np.testing.assert_allclose(state.reward_terms["source"], [0.005, 0.015]) assert env._compiled_manager_program is not None - warmup = env.warmup() - assert warmup.signatures - def test_read_plan_refreshes_resolved_leaves_in_place_without_rebuilding() -> None: env = _ManagerEnv(num_envs=3) @@ -687,6 +697,7 @@ def test_final_observations_capture_command_before_reset_and_survive_steps(with_ state = env.step(torch.full((3, 1), 0.25, dtype=torch.float32)) final = state.final_obs assert final is not None + assert isinstance(final.env_ids, torch.Tensor) np.testing.assert_array_equal(final.env_ids, [0, 2]) # Final observations capture the post-transition (post-advance) command. np.testing.assert_allclose(final.obs.policy[:, 1], 2.0 + env.cfg.ctrl_dt) From a73d0c3d7cd1cf97770c3607a3a721f1a3a950a1 Mon Sep 17 00:00:00 2001 From: zilch <147668916@qq.com> Date: Fri, 6 Nov 2026 23:14:49 +0800 Subject: [PATCH 37/41] feat(manager)!: take reset env_ids as torch tensors 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. --- .../en/user_guide/tutorial/building_envs/manager_env.md | 5 ++++- .../user_guide/tutorial/building_envs/manager_env.md | 3 +++ motrix_env_core/src/motrix_env_core/array/env.py | 9 +++++++-- motrix_env_core/src/motrix_env_core/manager/env.py | 1 - motrix_envs/tests/test_walk_randomization.py | 6 +++--- wiki/design/manager/task-authoring.md | 4 ++++ 6 files changed, 21 insertions(+), 7 deletions(-) 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 index 72306540..c5fbd52c 100644 --- a/docs/source/en/user_guide/tutorial/building_envs/manager_env.md +++ b/docs/source/en/user_guide/tutorial/building_envs/manager_env.md @@ -48,7 +48,10 @@ 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 +semantics match DirectEnv. The manager frontend's public currency is uniformly torch: +`step(actions)` and `reset(env_ids)` (including `sim_reset_ids`) take integer +`torch.Tensor`s, e.g. `env.reset(torch.tensor([0, 2], dtype=torch.long))`; passing a +numpy array raises a `TypeError`. 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: 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 index 5a4d44f4..910229c9 100644 --- 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 @@ -44,6 +44,9 @@ registry.env("my-task")(ManagerEnv) ## 生命周期 `ManagerEnv` 同样继承 `ArrayEnv`,step 时序与 auto-reset / truncation 语义与 DirectEnv 一致。 +Manager 前端的公开接口统一使用 torch:`step(actions)` 与 `reset(env_ids)`(含 `sim_reset_ids`) +接收整数 `torch.Tensor`,例如 `env.reset(torch.tensor([0, 2], dtype=torch.long))`; +传入 numpy 数组会抛出 `TypeError`。 下图把各回调整合进流水线:绿色是 manager kernel(`@dispatch`)执行的部分, 紫色是 manager 的主机侧 Python 回调,蓝色是框架——编排各阶段的宿主与其自身的数据通路: diff --git a/motrix_env_core/src/motrix_env_core/array/env.py b/motrix_env_core/src/motrix_env_core/array/env.py index 8abe6e7c..766529b3 100644 --- a/motrix_env_core/src/motrix_env_core/array/env.py +++ b/motrix_env_core/src/motrix_env_core/array/env.py @@ -50,9 +50,11 @@ class FinalObs(Generic[T]): Rows in ``obs`` align with ``env_ids``. Arrays remain valid across steps; consumers must treat them as read-only. ``obs.value`` stays optional. + ``env_ids`` follows the frontend's currency: numpy for NumPy-frontend + environments, a torch integer tensor for the manager pipeline. """ - env_ids: np.ndarray + env_ids: np.ndarray | torch.Tensor obs: EnvObs[T] @@ -277,9 +279,12 @@ def _reset_done_envs(self) -> None: self.reset(env_ids) @abc.abstractmethod - def reset(self, env_ids: np.ndarray) -> None: + def reset(self, env_ids: np.ndarray | torch.Tensor) -> None: """Write reset rows for the selected environments. + ``env_ids`` follows the frontend's currency: numpy arrays for the + NumPy frontend (DirectEnv), torch tensors for the manager pipeline. + Subclasses address their simulator's rows themselves. Observation generation is deferred until :meth:`compute_observation`, after reset writes have completed. diff --git a/motrix_env_core/src/motrix_env_core/manager/env.py b/motrix_env_core/src/motrix_env_core/manager/env.py index a961dd16..83384af4 100644 --- a/motrix_env_core/src/motrix_env_core/manager/env.py +++ b/motrix_env_core/src/motrix_env_core/manager/env.py @@ -828,4 +828,3 @@ def _validate_kernel_context(self, inputs: tuple[ResolvedTree, ...]) -> None: except (TypeError, ValueError) as error: message = f"Invalid kernel context for {type(self).__name__} task kernel: {error}" raise type(error)(message) from error - diff --git a/motrix_envs/tests/test_walk_randomization.py b/motrix_envs/tests/test_walk_randomization.py index 22db019f..ec8d455d 100644 --- a/motrix_envs/tests/test_walk_randomization.py +++ b/motrix_envs/tests/test_walk_randomization.py @@ -105,7 +105,7 @@ def test_action_reset_clears_only_selected_lane_history_and_keeps_pointer(): env.apply_action(torch.full(controls.shape, value, dtype=torch.float32), env._state) ptr_before_reset = action_state.action_ptr.clone() other_lane_history = action_state.action_queue[1].clone() - env.reset(np.asarray([0], dtype=np.int64)) + env.reset(torch.tensor([0], dtype=torch.long)) assert bool((action_state.action_queue[0] == 0.0).all()) torch.testing.assert_close(action_state.action_queue[1], other_lane_history, rtol=0, atol=0) torch.testing.assert_close(action_state.action_ptr, ptr_before_reset, rtol=0, atol=0) @@ -113,7 +113,7 @@ def test_action_reset_clears_only_selected_lane_history_and_keeps_pointer(): def test_action_delay_reset_resamples_within_range(): env, action_state, _ = _tiny_delay_manager(lo=0, hi=2) - env.reset(np.arange(action_state.action_queue.shape[0])) + env.reset(torch.arange(action_state.action_queue.shape[0], dtype=torch.long)) assert bool(((action_state.delay_steps >= 0) & (action_state.delay_steps <= 2)).all()) @@ -153,7 +153,7 @@ def test_walk_env_randomized_reset_samples_within_ranges(): env = ManagerEnv(cfg, num_envs=2) try: env.step(torch.zeros((2, env.num_actuators), dtype=torch.float32)) - env.reset(np.arange(2)) + env.reset(torch.arange(2, dtype=torch.long)) walk = env.command_terms["walk"] # Phase increments invert the configured gait-period interval. command_cfg = cfg.commands.walk diff --git a/wiki/design/manager/task-authoring.md b/wiki/design/manager/task-authoring.md index 1e1d5200..1329a011 100644 --- a/wiki/design/manager/task-authoring.md +++ b/wiki/design/manager/task-authoring.md @@ -175,6 +175,10 @@ Observation、reward、termination term 不应通过共享全局 registry 查找 Reset 配置声明 simulator state 和 term state 的恢复行为。Term 的 `reset(env_ids)` 只修改指定 environment rows;不要在 term 中实现全局 auto-reset、episode step 或 truncation。 这些职责属于 [Manager Runtime](./runtime.md) 所述的 `ArrayEnv` lifecycle 和 backend reset program。 +对外 API 层面,`ManagerEnv.reset(env_ids, sim_reset_ids=None)` 接收整数 `torch.Tensor` +(例如 `torch.tensor([0, 2], dtype=torch.long)`),传入 numpy 数组抛出 `TypeError`; +term 侧的 host 回调(command reset、action reset)仍保持 numpy 契约。 + ## 8. 开发验收 新增 Manager task 至少应验证: From 3b0ee910f8e5b7a3a2c7900e09cddb0790f6393a Mon Sep 17 00:00:00 2001 From: zilch <147668916@qq.com> Date: Fri, 6 Nov 2026 23:28:30 +0800 Subject: [PATCH 38/41] refactor(manager)!: delete the program fingerprint chain 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). --- .../motrix_env_core/manager/fingerprint.py | 84 --------- .../src/motrix_env_core/manager/resolver.py | 172 +----------------- .../tests/test_term_arg_fingerprint.py | 77 -------- 3 files changed, 10 insertions(+), 323 deletions(-) delete mode 100644 motrix_env_core/src/motrix_env_core/manager/fingerprint.py delete mode 100644 motrix_env_core/tests/test_term_arg_fingerprint.py diff --git a/motrix_env_core/src/motrix_env_core/manager/fingerprint.py b/motrix_env_core/src/motrix_env_core/manager/fingerprint.py deleted file mode 100644 index 49090d42..00000000 --- a/motrix_env_core/src/motrix_env_core/manager/fingerprint.py +++ /dev/null @@ -1,84 +0,0 @@ -# Copyright Motphys Technology Co., Ltd. 2025, 2026 -# SPDX-License-Identifier: Apache-2.0 - -"""Function dependency fingerprints shared by Numba compilation paths.""" - -import hashlib -import inspect -from collections.abc import Callable -from types import CodeType -from typing import Any - -import numpy as np - - -def _stable_constant(value: Any) -> tuple[Any, ...]: - if isinstance(value, CodeType): - return ("code", _code_fingerprint_parts(value)) - if isinstance(value, tuple): - return ("tuple", tuple(_stable_constant(item) for item in value)) - if isinstance(value, frozenset): - items = sorted((_stable_constant(item) for item in value), key=repr) - return ("frozenset", tuple(items)) - if value is None or isinstance(value, (bool, int, float, complex, str, bytes)): - return (type(value).__name__, value) - return (type(value).__module__, type(value).__qualname__, repr(value)) - - -def _code_fingerprint_parts(code: CodeType) -> tuple[Any, ...]: - # marshal includes string interning/reference flags that can change when - # Numba inspects the same code object. Hash immutable code fields instead. - return ( - code.co_argcount, - code.co_posonlyargcount, - code.co_kwonlyargcount, - code.co_nlocals, - code.co_stacksize, - code.co_flags, - code.co_code, - tuple(_stable_constant(value) for value in code.co_consts), - code.co_names, - code.co_varnames, - code.co_freevars, - code.co_cellvars, - code.co_filename, - code.co_name, - code.co_firstlineno, - code.co_lnotab, - ) - - -def function_fingerprint(function: Callable[..., Any]) -> str: - """Hash function code, defaults, closures, and referenced global helpers. - - Dispatch entries and kernel-data methods use the same dependency traversal. - Numba dispatchers are unwrapped to their Python functions. Code objects include - constants and nested code, even when source is unavailable. Cycles are visited - once; cheaply representable global/closure constants are hashed by value. - Callers remain responsible for discovering member-method dependencies and - incorporating their fingerprints into the appropriate compilation cache key. - """ - hasher = hashlib.sha256() - seen: set[int] = set() - queue = [function] - while queue: - current = queue.pop() - current = getattr(current, "py_func", current) - if not inspect.isfunction(current) or id(current) in seen: - continue - seen.add(id(current)) - code = current.__code__ - hasher.update(f"{current.__module__}.{current.__qualname__}\0".encode()) - hasher.update(repr(_code_fingerprint_parts(code)).encode()) - references = [(name, current.__globals__[name]) for name in code.co_names if name in current.__globals__] - references.extend((f"default[{i}]", value) for i, value in enumerate(current.__defaults__ or ())) - references.extend(sorted((current.__kwdefaults__ or {}).items())) - if current.__closure__: - references.extend(zip(code.co_freevars, (cell.cell_contents for cell in current.__closure__))) - for name, referenced in references: - helper = getattr(referenced, "py_func", referenced) - if inspect.isfunction(helper): - queue.append(helper) - elif isinstance(referenced, (int, float, bool, str, complex, np.generic, tuple, type(None))): - hasher.update(f"{name}={referenced!r}\0".encode()) - return hasher.hexdigest() diff --git a/motrix_env_core/src/motrix_env_core/manager/resolver.py b/motrix_env_core/src/motrix_env_core/manager/resolver.py index 2085b009..f3c1f47a 100644 --- a/motrix_env_core/src/motrix_env_core/manager/resolver.py +++ b/motrix_env_core/src/motrix_env_core/manager/resolver.py @@ -15,12 +15,11 @@ from __future__ import annotations -import hashlib import inspect import logging import warnings from collections.abc import Callable, Mapping -from dataclasses import dataclass, fields +from dataclasses import dataclass from time import perf_counter from types import MappingProxyType from typing import Any, get_args, get_origin, get_type_hints @@ -30,7 +29,6 @@ from motrix_env_core.manager.context import ManagerContext from motrix_env_core.manager.env import ManagerEnv, ReadPlan -from motrix_env_core.manager.fingerprint import function_fingerprint from motrix_env_core.manager.kernel_data import ( Map, ResolvedLeaf, @@ -41,8 +39,6 @@ leaf_value_type, tree_def_for, ) -from motrix_env_core.manager.kernel_data.methods import dispatch_methods -from motrix_env_core.manager.kernel_data.tree import BaseDef, KernelMapDef, LeafDef from motrix_env_core.manager.observations import ObservationGroupEntry, ObsTerm from motrix_env_core.manager.rewards import RewardTerm from motrix_env_core.manager.terminations import TerminationTerm @@ -50,8 +46,6 @@ logger = logging.getLogger(__name__) -_PROGRAM_FINGERPRINT_VERSION = 4 - # Simulator-backed read-only arena views are wrapped as tensors at resolution # time; suppress PyTorch's non-writable-array warning once here instead of # copying every view (mirrors the filter in the Torch executor). @@ -262,7 +256,6 @@ class ResolvedManagerProgram: observations: dict[str, tuple[ObservationTermCall, ...]] prepared_trees: tuple[ResolvedTree, ...] context: ResolvedManagerContext - fingerprint: str __all__ = [ @@ -303,101 +296,9 @@ class ManagerSimInput: query: SimDataQuery -def _manager_sim_inputs(sim_data: PhysicsReadProgram) -> tuple[Mapping[str, ManagerSimInput], tuple[Any, ...]]: +def _manager_sim_inputs(sim_data: PhysicsReadProgram) -> Mapping[str, ManagerSimInput]: inputs = {key: ManagerSimInput(sim_data.view(key), sim_data.query(key)) for key in sim_data.keys} - fingerprint = tuple( - (key, repr(binding.query), binding.value.shape[1:], binding.value.strides) for key, binding in inputs.items() - ) - return MappingProxyType(inputs), fingerprint - - -def _tree_leaf_fingerprint(tree: ResolvedTree) -> tuple[Any, ...]: - """Per-leaf content fingerprint: array/tensor shape/dtype identity.""" - return tuple( - ( - path, - None - if not isinstance(leaf.value, (np.ndarray, torch.Tensor)) - else (tuple(leaf.value.shape), str(leaf.value.dtype), getattr(leaf.value, "strides", None)), - ) - for path, leaf in tree.leaves.items() - ) - - -def _kernel_data_program_fingerprint(tree_def: TreeClassDef) -> str: - """Deterministic kernel-data identity for resolved-program invalidation. - - Combines the structural tree fingerprint with per-leaf value types and - the code fingerprints of ``@dispatch`` methods on every nested record, so - schema, leaf-type, or method-code changes each invalidate compiled - program caches. - """ - raw = repr((_PROGRAM_FINGERPRINT_VERSION, _program_fingerprint_part(tree_def))) - return hashlib.sha256(raw.encode()).hexdigest() - - -def _program_fingerprint_part(tree_def: BaseDef) -> Any: - if isinstance(tree_def, LeafDef): - value_type = leaf_value_type(tree_def.annotation) - return ( - "leaf", - tree_def.path, - None if value_type is None else f"{value_type.__module__}.{value_type.__qualname__}", - ) - if isinstance(tree_def, KernelMapDef): - return ( - "map", - tree_def.fingerprint, - tuple((entry.key, _program_fingerprint_part(entry.tree_def)) for entry in tree_def.entries), - ) - assert isinstance(tree_def, TreeClassDef) - return ( - "record", - tree_def.fingerprint, - tuple( - (name, function_fingerprint(method)) - for name, method in sorted(dispatch_methods(tree_def.logical_type).items()) - ), - tuple((field.name, _program_fingerprint_part(field.tree_def)) for field in tree_def.fields), - ) - - -def _carrier_fingerprint(value: Any) -> Any: - """Structural + content fingerprint of a constant term argument. - - Recurses over dataclass fields recording field names, concrete types, - tensor/ndarray dtype+shape+content hashes, and Python scalar values. - Carriers previously resolved as kernel-data trees contributed no content - to the program fingerprint at all, and array arguments previously lived - as shared leaves outside the fingerprint; this makes every content - change (a grid array value, a scalar) invalidate compiled-program caches. - """ - if isinstance(value, torch.Tensor): - return ( - "tensor", - str(value.dtype), - tuple(value.shape), - hashlib.sha256(value.detach().contiguous().numpy().tobytes()).hexdigest(), - ) - if isinstance(value, np.ndarray): - return ( - "ndarray", - str(value.dtype), - value.shape, - hashlib.sha256(np.ascontiguousarray(value).tobytes()).hexdigest(), - ) - if is_kernel_data(value): - return ( - "dataclass", - type(value).__qualname__, - tuple((field.name, _carrier_fingerprint(getattr(value, field.name))) for field in fields(value)), - ) - if isinstance(value, (bool, int, float, str)): - return (type(value).__name__, value) - raise TypeError( - f"Term argument of unsupported type {type(value).__name__}; " - "allowed: torch.Tensor, ndarray, Python scalars, str, carrier dataclasses." - ) + return MappingProxyType(inputs) class ManagerProgramResolver: @@ -405,14 +306,11 @@ class ManagerProgramResolver: def __init__(self, env: ManagerEnv): self._env = env - self._sim_inputs: Mapping[str, ManagerSimInput] - self._sim_inputs, sim_input_fingerprint = _manager_sim_inputs(env.sim_data) + self._sim_inputs = _manager_sim_inputs(env.sim_data) self._prepared_trees: list[ResolvedTree] = [] - self._layout_identity: list[Any] = [sim_input_fingerprint] def build(self) -> ResolvedManagerProgram: build_started = perf_counter() - self._resolve_runtime_terms() observation_groups = self._env.observation_groups reward_terms = self._env._reward_terms termination_terms = self._env.termination_manager.terms @@ -491,7 +389,6 @@ def build(self) -> ResolvedManagerProgram: observations=observation_terms, prepared_trees=tuple(self._prepared_trees), context=context, - fingerprint=hashlib.sha256(repr(tuple(self._layout_identity)).encode()).hexdigest(), ) # -- resets ------------------------------------------------------------- @@ -541,9 +438,7 @@ def _resolve_reset( "Map[torch.Tensor] (conventionally named sim_writes)." ) dispatcher = function - args = tuple( - self._resolve_term_arg(f"sim_reset.{name}.args[{index}]", value) for index, value in enumerate(term.args) - ) + args = tuple(self._resolve_term_arg(value) for index, value in enumerate(term.args)) return ResolvedSimReset( SimResetTermCall(dispatcher, args), f"sim_reset_{name}_writes", @@ -553,32 +448,6 @@ def _resolve_reset( # -- layouts ------------------------------------------------------------ - def _resolve_runtime_terms(self) -> None: - for name, action_term in self._env.action_terms.items(): - tree_def = tree_def_for(action_term.state) - self._layout_identity.append( - ( - "action", - name, - type(action_term).__qualname__, - tree_def.fingerprint, - self._env.action_slices[name], - None - if self._env.action_actuators[name] is None - else tuple(spec.name for spec in self._env.action_actuators[name]), - ) - ) - for name, command_term in self._env.command_terms.items(): - tree_def = tree_def_for(command_term) - self._layout_identity.append( - ( - "command", - name, - type(command_term).__qualname__, - tree_def.fingerprint, - ) - ) - def _resolve_manager_context( self, observation_groups: dict[str, ObservationGroupEntry], @@ -604,17 +473,7 @@ def _resolve_manager_context( dt=np.float32(self._env.cfg.ctrl_dt), sim_reset_requested=torch.from_numpy(self._env._sim_reset_requested), ) - tree_index, tree_def = self._register_tree("manager_context", value) - tree = self._prepared_trees[tree_index] - self._layout_identity.append( - ("manager_context", _kernel_data_program_fingerprint(tree_def), _tree_leaf_fingerprint(tree)) - ) - self._layout_identity.append( - ( - "metrics", - tuple((name, value.shape, value.dtype.str) for name, value in self._env.metrics.items()), - ) - ) + tree_index = self._register_tree("manager_context", value)[0] return ResolvedManagerContext(tree_index) def _resolve_command_hooks(self, function_name: str) -> tuple[tuple[str, Callable], ...]: @@ -702,10 +561,7 @@ def _resolve_observation(self, group: str, name: str, term: ObsTerm) -> Observat f"Manager observation {group}.{name} dispatch must annotate its first two parameters as " "ManagerContext and torch.Tensor (conventionally named ctx and out)." ) - args = tuple( - self._resolve_term_arg(f"observation.{group}.{name}.args[{index}]", value) - for index, value in enumerate(term.args) - ) + args = tuple(self._resolve_term_arg(value) for index, value in enumerate(term.args)) return ObservationTermCall(group, name, function, args) # -- rewards and terminations -------------------------------------------- @@ -760,30 +616,22 @@ def _resolve_dispatch_args(self, group: str, name: str, term: Any) -> tuple[Prep f"Manager {group} {name} dispatch must annotate its first parameter as ManagerContext " f"(conventionally named ctx)." ) - return tuple( - self._resolve_term_arg(f"{group}.{name}.args[{index}]", value) for index, value in enumerate(term.args) - ) + return tuple(self._resolve_term_arg(value) for index, value in enumerate(term.args)) # -- argument resolution --------------------------------------------------- - def _resolve_term_arg(self, source_name: str, value: Any) -> PreparedArg: + def _resolve_term_arg(self, value: Any) -> PreparedArg: """Resolve one manager term argument into a prepared descriptor. Carrier dataclass and tensor/ndarray arguments pass through as constants: Dynamo traces their attribute access and tensor ops - directly, and a structural fingerprint over their raw fields (dtype, - shape, content hash) feeds the program fingerprint for cache - invalidation. Float-like numpy scalars normalize to Python floats; + directly. Float-like numpy scalars normalize to Python floats; simulator queries resolve to their read-program batched array; every other value is a literal constant. """ if is_kernel_data(value): - self._layout_identity.append( - ("term_arg", source_name, type(value).__qualname__, _carrier_fingerprint(value)) - ) return ConstantArg(value) if isinstance(value, (torch.Tensor, np.ndarray)): - self._layout_identity.append(("term_arg", source_name, _carrier_fingerprint(value))) return ConstantArg(value) if isinstance(value, np.floating): return ConstantArg(float(value)) diff --git a/motrix_env_core/tests/test_term_arg_fingerprint.py b/motrix_env_core/tests/test_term_arg_fingerprint.py deleted file mode 100644 index 12300122..00000000 --- a/motrix_env_core/tests/test_term_arg_fingerprint.py +++ /dev/null @@ -1,77 +0,0 @@ -# Copyright Motphys Technology Co., Ltd. 2025, 2026 -# SPDX-License-Identifier: Apache-2.0 - -"""Constant term-argument fingerprint contract. - -Carrier dataclasses and raw tensors/ndarrays passed as manager term -arguments resolve as constants; their structural fingerprint must keep -compiled-program caches invalidated by any content change (tensor bytes or -scalar values). -""" - -from dataclasses import dataclass - -import numpy as np -import pytest -import torch - -from motrix_env_core.manager.resolver import _carrier_fingerprint - - -@dataclass(frozen=True, slots=True) -class _Nested: - flag: bool - - -@dataclass(frozen=True, slots=True) -class _Carrier: - grid: torch.Tensor - label: str - offset: float - nested: _Nested - - -def test_fingerprint_changes_on_tensor_content_and_scalar() -> None: - carrier = _Carrier(torch.arange(6, dtype=torch.float32).reshape(2, 3), "a", 1.5, _Nested(True)) - base = _carrier_fingerprint(carrier) - edited_grid = _Carrier(carrier.grid.clone(), "a", 1.5, _Nested(True)) - edited_grid.grid[0, 0] += 1.0 - assert _carrier_fingerprint(edited_grid) != base - assert _carrier_fingerprint(_Carrier(carrier.grid, "b", 1.5, _Nested(True))) != base - assert _carrier_fingerprint(_Carrier(carrier.grid, "a", 2.0, _Nested(True))) != base - assert _carrier_fingerprint(_Carrier(carrier.grid, "a", 1.5, _Nested(False))) != base - - -def test_fingerprint_stable_for_equal_content() -> None: - first = _Carrier(torch.ones(3), "a", 1.5, _Nested(True)) - second = _Carrier(torch.ones(3), "a", 1.5, _Nested(True)) - assert _carrier_fingerprint(first) == _carrier_fingerprint(second) - - -def test_raw_tensor_arg_fingerprint_flips_on_content_edit() -> None: - arg = torch.arange(4, dtype=torch.float32) - base = _carrier_fingerprint(arg) - assert _carrier_fingerprint(torch.arange(4, dtype=torch.float32)) == base - assert _carrier_fingerprint(torch.arange(4, dtype=torch.float64)) != base - assert _carrier_fingerprint(torch.arange(5, dtype=torch.float32)) != base - edited = arg.clone() - edited[0] += 1.0 - assert _carrier_fingerprint(edited) != base - - -def test_raw_ndarray_arg_fingerprint_flips_on_content_edit() -> None: - arg = np.arange(4, dtype=np.float32) - base = _carrier_fingerprint(arg) - assert _carrier_fingerprint(np.arange(4, dtype=np.float32)) == base - edited = arg.copy() - edited[0] += 1.0 - assert _carrier_fingerprint(edited) != base - - -def test_fingerprint_rejects_unsupported_field_type() -> None: - @dataclass(frozen=True, slots=True) - class _Bad: - leak: np.float32 - - with pytest.raises(TypeError, match="unsupported type"): - _carrier_fingerprint(_Bad(np.float32(1.0))) From 82ed40c19fb8ec163392f89790d7291ab9e1a8fc Mon Sep 17 00:00:00 2001 From: zilch <147668916@qq.com> Date: Fri, 6 Nov 2026 23:33:27 +0800 Subject: [PATCH 39/41] feat(torch_manager): optional barrier-removal pass for Inductor CPU kernels MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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). --- .../torch_manager/merged_omp.py | 188 ++++++++++++++++++ .../motrix_env_core/torch_manager/runtime.py | 4 + wiki/design/manager/torch-executor.md | 8 + 3 files changed, 200 insertions(+) create mode 100644 motrix_env_core/src/motrix_env_core/torch_manager/merged_omp.py diff --git a/motrix_env_core/src/motrix_env_core/torch_manager/merged_omp.py b/motrix_env_core/src/motrix_env_core/torch_manager/merged_omp.py new file mode 100644 index 00000000..62647ce2 --- /dev/null +++ b/motrix_env_core/src/motrix_env_core/torch_manager/merged_omp.py @@ -0,0 +1,188 @@ +# Copyright Motphys Technology Co., Ltd. 2025, 2026 +# SPDX-License-Identifier: Apache-2.0 + +"""Experiment: drop needless OpenMP barriers in Inductor-generated CPU kernels. + +Inductor (torch 2.7) already emits ONE ``#pragma omp parallel`` region per +compiled graph containing a sequence of ``#pragma omp for`` loops. Every one +of those loops carries an implicit barrier, so a graph with 30 loops pays 30 +full-team barriers per invocation — the fork/join cost is shared, but the +barrier cost is not, which is the dominant overhead when thread counts grow. + +This module post-processes the generated C++ source before Inductor compiles +it: it converts an ``#pragma omp for`` to ``#pragma omp for nowait`` when no +later loop in the same parallel region reads (or rewrites) a buffer that an +earlier loop wrote since the previous surviving barrier. Barriers are kept +conservatively whenever a write/read or write/write pair touches the same +buffer across a loop boundary. + +Enabled only through ``MOTRIX_TORCH_MERGED_OMP=1`` (set ``2`` for a per-graph +report on stderr). Integration is a monkeypatch of +``torch._inductor.codecache.CppCodeCache.load_async`` — the single choke point +every CPU kernel source passes through — applied before the first +``torch.compile`` call (async-compile workers fork afterwards and inherit it), +and it is idempotent. +""" + +from __future__ import annotations + +import os +import re +import sys + +_PTR = r"(?:in|out|in_out)_ptr\d+" +_ALIAS = re.compile(rf"^\s*auto\s+({_PTR})\s*=\s*({_PTR})\s*;") +_WRITE = re.compile(rf"^\s*({_PTR})\s*\[[^]]*\]\s*=(?!=)", re.M) +_READ = re.compile(rf"\b({_PTR})\b") +_LOOP_PRAGMA = re.compile(r"^\s*#pragma omp for\b(?P.*)$") +_PARALLEL = re.compile(r"^\s*#pragma omp parallel\b") + + +def _collect_aliases(lines: list[str]) -> dict[str, str]: + aliases: dict[str, str] = {} + for line in lines: + m = _ALIAS.match(line) + if m: + aliases[m.group(1)] = aliases.get(m.group(2), m.group(2)) + return aliases + + +def _region_loops(lines: list[str], start: int) -> list[dict]: + """Describe the ``#pragma omp for`` loop blocks of one parallel region. + + ``start`` indexes the line holding the region's opening brace. Returns one + record per loop: its ``#pragma omp for`` line index, end line, and the + canonical read/write pointer sets. The final record carries the region's + exclusive end line. + """ + aliases = _collect_aliases(lines) + loops: list[dict] = [] + depth = 0 + i = start + n = len(lines) + while i < n: + line = lines[i] + pragma = _LOOP_PRAGMA.match(line) + if pragma and depth >= 1: + # Loop body extends from the pragma to the closing brace that + # returns to the depth held just before the enclosing block. + loop_depth = depth + j = i + 1 + while j < n: + for ch in lines[j]: + if ch == "{": + loop_depth += 1 + elif ch == "}": + loop_depth -= 1 + if loop_depth < depth: + break + j += 1 + body = "\n".join(lines[i : j + 1]) + writes = {_ALIAS_TO(aliases, w) for w in _WRITE.findall(body)} + reads = {_ALIAS_TO(aliases, r) for r in _READ.findall(body)} + loops.append({"pragma": i, "end": j, "reads": reads, "writes": writes}) + i = j + continue + depth += line.count("{") - line.count("}") + if depth == 0: + loops.append({"region_end": i}) + return loops + i += 1 + return loops + + +def _ALIAS_TO(aliases: dict[str, str], name: str) -> str: + return aliases.get(name, name) + + +def transform_source(source: str) -> tuple[str, dict[str, int]]: + """Rewrite ``#pragma omp for`` loops to ``nowait`` where barrier-free. + + Returns the (possibly unchanged) source and a stats dict. Any parse doubt + leaves the source untouched. + """ + if "#pragma omp parallel" not in source: + return source, {"regions": 0, "loops": 0, "nowait": 0, "barriers_kept": 0} + lines = source.splitlines(keepends=True) + out = list(lines) + stats = {"regions": 0, "loops": 0, "nowait": 0, "barriers_kept": 0} + i = 0 + n = len(lines) + while i < n: + if not _PARALLEL.match(lines[i]): + i += 1 + continue + # Find the region's opening brace. + j = i + 1 + while j < n and "{" not in lines[j]: + if ";" in lines[j]: # pragma with clause then body on next lines + break + j += 1 + if j >= n or "{" not in lines[j]: + i += 1 + continue + loops = _region_loops(lines, j) + region_end = loops.pop()["region_end"] if loops and "region_end" in loops[-1] else None + if region_end is None or not loops: + i = region_end + 1 if region_end is not None else j + 1 + continue + stats["regions"] += 1 + stats["loops"] += len(loops) + # Forward scan: decide, for each loop, whether the barrier AFTER it + # can be dropped. A barrier must stay when any later loop before the + # next surviving barrier reads or writes a buffer written since the + # last surviving barrier. + pending_writes: set[str] = set() + keep_barrier = [False] * len(loops) + for idx, loop in enumerate(loops): + if (loop["reads"] | loop["writes"]) & pending_writes: + # This loop consumes buffers written since the last surviving + # barrier (idx is > 0 here because pending starts empty), so + # the barrier after loop idx-1 must stay. + keep_barrier[idx - 1] = True + pending_writes = set() + pending_writes |= loop["writes"] + for idx, loop in enumerate(loops): + if idx == len(loops) - 1 or keep_barrier[idx]: + stats["barriers_kept"] += 1 + continue + m = _LOOP_PRAGMA.match(out[loop["pragma"]]) + out[loop["pragma"]] = f"#pragma omp for nowait{m.group('rest')}\n" + stats["nowait"] += 1 + i = region_end + 1 + return "".join(out), stats + + +_installed = False + + +def install() -> None: + """Patch ``CppCodeCache.load_async`` once; no-op when already installed.""" + global _installed + if _installed: + return + from torch._inductor.codecache import CppCodeCache + + verbose = os.environ.get("MOTRIX_TORCH_MERGED_OMP", "0") == "2" + original = CppCodeCache.load_async.__func__ + + def load_async(cls, source_code, *args, **kwargs): + source_code, stats = transform_source(source_code) + if verbose and stats.get("loops"): + print( + f"[merged-omp] regions={stats['regions']} loops={stats['loops']} " + f"nowait={stats['nowait']} barriers_kept={stats['barriers_kept']}", + file=sys.stderr, + ) + return original(cls, source_code, *args, **kwargs) + + CppCodeCache.load_async = classmethod(load_async) + CppCodeCache.load = classmethod( + lambda cls, source_code, device_type="cpu": cls.load_async(source_code, device_type)() + ) + _installed = True + + +def maybe_install() -> None: + if os.environ.get("MOTRIX_TORCH_MERGED_OMP", "0") not in ("0", ""): + install() diff --git a/motrix_env_core/src/motrix_env_core/torch_manager/runtime.py b/motrix_env_core/src/motrix_env_core/torch_manager/runtime.py index 60d4e9b3..8c5d5182 100644 --- a/motrix_env_core/src/motrix_env_core/torch_manager/runtime.py +++ b/motrix_env_core/src/motrix_env_core/torch_manager/runtime.py @@ -45,6 +45,7 @@ SimResetTermCall, TerminationTermCall, ) +from motrix_env_core.torch_manager import merged_omp from motrix_env_core.torch_manager.rand import draw_uniform if TYPE_CHECKING: @@ -75,6 +76,9 @@ def _compile_options() -> dict[str, Any]: # ``#pragma omp single``). Manager term loops are exactly that size, so # lower the threshold to let them parallelize. chunk = int(os.environ.get("MOTRIX_TORCH_MIN_CHUNK", "256")) + # Experiment: barrier-free OpenMP loops in the generated C++ (no-op + # unless MOTRIX_TORCH_MERGED_OMP is set). + merged_omp.maybe_install() return {"cpp.threads": threads, "cpp.min_chunk_size": chunk} diff --git a/wiki/design/manager/torch-executor.md b/wiki/design/manager/torch-executor.md index 234c1720..21dfd458 100644 --- a/wiki/design/manager/torch-executor.md +++ b/wiki/design/manager/torch-executor.md @@ -115,3 +115,11 @@ ActionTerm 的 `process/_reset` 也保持 numpy host 路径不变。 0.59ms(Numba 因单小并行区天然免疫)。 - **已知残余**。EPYC 上相对 Numba 约 1ms 残余 = 冷缓存惩罚 + reset 静态全批布尔掩码设计 (RNG 状态 `where(mask, ...)` 门控)带来的约 3x lane 计算结构性成本;本机持平 16.8–17.0ms。 +- **merged-OMP 实验(默认关闭,`MOTRIX_TORCH_MERGED_OMP=1`)**。Inductor 2.7 生成的 CPU 图 + 本就是单 `#pragma omp parallel` 区,无 per-loop fork/join;成本在每个 `#pragma omp for` 的 + 隐式 barrier(四个编译 kernel 分别含 30/23/13/11 个循环)。该实验在 `CppCodeCache.load_async` + 处后处理生成 C++:按循环的指针读写集(含 `auto out_ptrN = in_out_ptrM` 别名归一)做依赖 + 分析,仅保留跨循环 RAW/WAR 确需的 barrier(25/30、21/23、11/13、7/11 个循环获得 + `nowait`)。obs 与 RNG 流 on/off 逐位一致;EPYC 默认 8 线程图内上限下约 5–8% 提速, + `MOTRIX_TORCH_COMPILE_THREADS=$T` 不封顶时 128t 崩塌 27.6→20.4ms(-26%)、192t + 33.0→29.8ms(-10%);拉平 192t 仍需跨图 region 融合或持久 lane kernel(future work)。 From a2dfea8736882e3f830436c908b9f2434b2e775e Mon Sep 17 00:00:00 2001 From: zilch <147668916@qq.com> Date: Fri, 6 Nov 2026 23:46:41 +0800 Subject: [PATCH 40/41] refactor(kernel_data)!: replace the schema tree with a direct leaf walk MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. --- .../manager/kernel_data/__init__.py | 15 -- .../manager/kernel_data/canonical.py | 9 +- .../manager/kernel_data/tree.py | 198 +++--------------- .../src/motrix_env_core/manager/resolver.py | 23 +- motrix_env_core/tests/test_kernel_data.py | 78 +++---- 5 files changed, 87 insertions(+), 236 deletions(-) diff --git a/motrix_env_core/src/motrix_env_core/manager/kernel_data/__init__.py b/motrix_env_core/src/motrix_env_core/manager/kernel_data/__init__.py index 53c13b3b..fde4bf60 100644 --- a/motrix_env_core/src/motrix_env_core/manager/kernel_data/__init__.py +++ b/motrix_env_core/src/motrix_env_core/manager/kernel_data/__init__.py @@ -1,31 +1,17 @@ # Copyright Motphys Technology Co., Ltd. 2025, 2026 -# SPDX-License-Identifier: Apache-2.0 from motrix_env_core.manager.kernel_data.canonical import canonicalize_kernel_data from motrix_env_core.manager.kernel_data.map import Map from motrix_env_core.manager.kernel_data.resolved import ResolvedLeaf, ResolvedTree from motrix_env_core.manager.kernel_data.tree import ( - BaseDef, - FieldDef, - KernelMapDef, - LeafDef, - MapEntryDef, - TreeClassDef, is_kernel_data, is_kernel_data_type, iter_leaves, leaf_value_type, - tree_def_for, ) __all__ = [ - "BaseDef", "Map", - "KernelMapDef", - "FieldDef", - "LeafDef", - "MapEntryDef", - "TreeClassDef", "ResolvedLeaf", "ResolvedTree", "canonicalize_kernel_data", @@ -33,5 +19,4 @@ "is_kernel_data_type", "iter_leaves", "leaf_value_type", - "tree_def_for", ] diff --git a/motrix_env_core/src/motrix_env_core/manager/kernel_data/canonical.py b/motrix_env_core/src/motrix_env_core/manager/kernel_data/canonical.py index 68c619ba..d985abcd 100644 --- a/motrix_env_core/src/motrix_env_core/manager/kernel_data/canonical.py +++ b/motrix_env_core/src/motrix_env_core/manager/kernel_data/canonical.py @@ -1,5 +1,4 @@ # Copyright Motphys Technology Co., Ltd. 2025, 2026 -# SPDX-License-Identifier: Apache-2.0 from typing import TypeVar @@ -7,7 +6,6 @@ is_kernel_data, iter_leaves, leaf_value_type, - tree_def_for, ) KernelDataType = TypeVar("KernelDataType") @@ -22,11 +20,10 @@ def canonicalize_kernel_data(value: KernelDataType, *, context: str) -> KernelDa """ if not is_kernel_data(value): raise TypeError(f"{context} must return a concrete kernel-data dataclass, got {type(value).__name__}.") - tree_def = tree_def_for(value) - for leaf_def, _leaf_value in iter_leaves(value, tree_def): - if leaf_value_type(leaf_def.annotation) is None: + for leaf_path, annotation, _leaf_value in iter_leaves(value): + if leaf_value_type(annotation) is None: raise TypeError( - f"{context} KernelData {'.'.join(leaf_def.path)} expected a tensor, an ndarray, or a numeric scalar." + f"{context} KernelData {'.'.join(leaf_path)} expected a tensor, an ndarray, or a numeric scalar." ) return value diff --git a/motrix_env_core/src/motrix_env_core/manager/kernel_data/tree.py b/motrix_env_core/src/motrix_env_core/manager/kernel_data/tree.py index 5ea12e09..da65aa1e 100644 --- a/motrix_env_core/src/motrix_env_core/manager/kernel_data/tree.py +++ b/motrix_env_core/src/motrix_env_core/manager/kernel_data/tree.py @@ -3,9 +3,9 @@ from __future__ import annotations -import hashlib from collections.abc import Iterator -from dataclasses import dataclass, fields, is_dataclass +from dataclasses import fields, is_dataclass +from functools import cache from typing import Any, get_origin, get_type_hints import numpy as np @@ -14,8 +14,6 @@ from motrix_env_core.manager.kernel_data.map import Map from motrix_env_core.sim.read import SimDataQuery -_TREE_SCHEMA_VERSION = 3 - def is_kernel_data_type(value_type: object) -> bool: """Return whether a type is a concrete kernel-data dataclass.""" @@ -51,191 +49,63 @@ def leaf_value_type(annotation: Any) -> type | None: return None -class BaseDef: - """Base class for nodes contained in a ``TreeClassDef`` schema tree.""" - - -@dataclass(frozen=True) -class LeafDef(BaseDef): - """Definition of one terminal value in a logical KernelData tree. - - Attributes: - path: Field and map-key components from the tree root to this leaf. - annotation: Declared leaf annotation. - """ - - path: tuple[str, ...] - annotation: Any - - -@dataclass(frozen=True) -class FieldDef(BaseDef): - """Named child of a ``TreeClassDef``. +@cache +def _type_hints(cls: type) -> dict[str, Any]: + return get_type_hints(cls, include_extras=True) - Attributes: - name: Python dataclass field name on the logical kernel-data class. - tree_def: Logical tree definition for the field value. - """ - - name: str - tree_def: BaseDef - - -@dataclass(frozen=True) -class TreeClassDef(BaseDef): - """Root definition of one concrete kernel-data class tree. - Every kernel-data value is described by a ``TreeClassDef`` root. - Its fields recursively contain nested class, map, and leaf definitions. +def iter_leaves(value: Any) -> Iterator[tuple[tuple[str, ...], Any, Any]]: + """Yield (path, annotation, leaf value) triples for one kernel-data value. - Attributes: - logical_type: Concrete dataclass type described by this definition. - fields: Field definitions in dataclass declaration order. - fingerprint: Deterministic identity of the complete class tree schema. + The tree is walked live along dataclass field declaration order and Map + insertion order. Structural mismatches raise ``TypeError``: a field whose + value is not its declared kernel-data class, a ``Map``-annotated field + without a concrete ``Map`` value, and class cycles. """ - - logical_type: type - fields: tuple[FieldDef, ...] - fingerprint: str - - -@dataclass(frozen=True) -class MapEntryDef(BaseDef): - """Named child of a value-dependent ``KernelMapDef``. - - Attributes: - key: Immutable string key exposed by ``Map.__getitem__``. - tree_def: Logical tree definition for the value associated with ``key``. - """ - - key: str - tree_def: BaseDef - - -@dataclass(frozen=True) -class KernelMapDef(BaseDef): - """Value-derived definition of one immutable heterogeneous ``Map``. - - Attributes: - path: Field and map-key components from the tree root to this map. - entries: Map children in stable insertion order. - fingerprint: Deterministic identity including keys and child schemas. - """ - - path: tuple[str, ...] - entries: tuple[MapEntryDef, ...] - fingerprint: str - - -def tree_def_for(value: Any) -> TreeClassDef: - """Build and validate the class tree schema for one kernel-data value.""" if not is_kernel_data(value): raise TypeError(f"KernelData tree root must be a dataclass value, got {type(value).__name__}.") - return _build_class_def(type(value), value, (), ()) - + yield from _iter_class_leaves(type(value), value, (), ()) -def iter_leaves(value: Any, tree_def: BaseDef) -> Iterator[tuple[LeafDef, Any]]: - """Yield (leaf definition, leaf value) pairs in declaration order. - The value is walked alongside the schema; structural mismatches (wrong - class type or changed Map keys) raise ``TypeError``. - """ - if isinstance(tree_def, LeafDef): - yield tree_def, value - return - if isinstance(tree_def, KernelMapDef): - if not isinstance(value, Map): - raise TypeError(f"KernelData value at {'.'.join(tree_def.path)!r} must be Map.") - if value.keys_tuple != tuple(entry.key for entry in tree_def.entries): - raise TypeError( - f"Map at {'.'.join(tree_def.path)!r} has keys {value.keys_tuple!r}, expected " - f"{tuple(entry.key for entry in tree_def.entries)!r}." - ) - for entry in tree_def.entries: - yield from iter_leaves(value[entry.key], entry.tree_def) - return - assert isinstance(tree_def, TreeClassDef) - if type(value) is not tree_def.logical_type: - raise TypeError(f"KernelData value must be {tree_def.logical_type.__name__}, got {type(value).__name__}.") - for field in tree_def.fields: - yield from iter_leaves(getattr(value, field.name), field.tree_def) - - -def _build_class_def( - logical_type: type, - value: Any | None, +def _iter_class_leaves( + declared_type: type, + value: Any, path: tuple[str, ...], stack: tuple[type, ...], -) -> TreeClassDef: - if logical_type in stack: - chain = " -> ".join(item.__name__ for item in (*stack, logical_type)) +) -> Iterator[tuple[tuple[str, ...], Any, Any]]: + if declared_type in stack: + chain = " -> ".join(item.__name__ for item in (*stack, declared_type)) raise TypeError(f"KernelData cycle detected: {chain}.") - annotations = get_type_hints(logical_type, include_extras=True) - field_defs: list[FieldDef] = [] - fingerprint_fields: list[Any] = [] - for data_field in fields(logical_type): - annotation = annotations[data_field.name] + if type(value) is not declared_type: + raise TypeError(f"KernelData value must be {declared_type.__name__}, got {type(value).__name__}.") + stack = (*stack, declared_type) + hints = _type_hints(declared_type) + for data_field in fields(declared_type): + annotation = hints[data_field.name] field_path = (*path, data_field.name) - field_value = getattr(value, data_field.name) if value is not None else None - child: BaseDef + field_value = getattr(value, data_field.name) if annotation is Map or get_origin(annotation) is Map: if not isinstance(field_value, Map): raise TypeError( - f"KernelData {logical_type.__name__}.{'.'.join(field_path)} requires a concrete Map value." + f"KernelData {declared_type.__name__}.{'.'.join(field_path)} requires a concrete Map value." ) - child = _build_map_def(field_value, field_path, (*stack, logical_type)) + yield from _iter_map_leaves(field_value, field_path, stack) elif is_kernel_data_type(annotation): - child = _build_class_def(annotation, field_value, field_path, (*stack, logical_type)) + yield from _iter_class_leaves(annotation, field_value, field_path, stack) else: - child = LeafDef(field_path, annotation) - field_defs.append(FieldDef(data_field.name, child)) - fingerprint_fields.append((data_field.name, _tree_fingerprint_part(child))) - identity = f"{logical_type.__module__}.{logical_type.__qualname__}" - raw_fingerprint = repr((_TREE_SCHEMA_VERSION, identity, tuple(fingerprint_fields))) - fingerprint = hashlib.sha256(raw_fingerprint.encode()).hexdigest() - return TreeClassDef(logical_type, tuple(field_defs), fingerprint) + yield field_path, annotation, field_value -def _build_map_def( +def _iter_map_leaves( value: Map, path: tuple[str, ...], stack: tuple[type, ...], -) -> KernelMapDef: - entries: list[MapEntryDef] = [] - fingerprint_entries: list[Any] = [] +) -> Iterator[tuple[tuple[str, ...], Any, Any]]: for key, entry_value in zip(value.keys_tuple, value.values_tuple, strict=True): entry_path = (*path, key) if isinstance(entry_value, Map): - child = _build_map_def(entry_value, entry_path, stack) + yield from _iter_map_leaves(entry_value, entry_path, stack) elif is_kernel_data(entry_value): - child = _build_class_def(type(entry_value), entry_value, entry_path, stack) + yield from _iter_class_leaves(type(entry_value), entry_value, entry_path, stack) else: - child = LeafDef(entry_path, type(entry_value)) - entries.append(MapEntryDef(key, child)) - fingerprint_entries.append((key, _tree_fingerprint_part(child))) - raw_fingerprint = repr((_TREE_SCHEMA_VERSION, "map", path, tuple(fingerprint_entries))) - fingerprint = hashlib.sha256(raw_fingerprint.encode()).hexdigest() - return KernelMapDef(path, tuple(entries), fingerprint) - - -def _tree_fingerprint_part(tree_def: BaseDef) -> Any: - if isinstance(tree_def, LeafDef): - return "leaf", tree_def.path, _annotation_identity(tree_def.annotation) - if isinstance(tree_def, KernelMapDef): - return ( - "map", - tree_def.path, - tuple((entry.key, _tree_fingerprint_part(entry.tree_def)) for entry in tree_def.entries), - ) - assert isinstance(tree_def, TreeClassDef) - return ( - "class", - f"{tree_def.logical_type.__module__}.{tree_def.logical_type.__qualname__}", - tuple((field.name, _tree_fingerprint_part(field.tree_def)) for field in tree_def.fields), - ) - - -def _annotation_identity(annotation: Any) -> Any: - if isinstance(annotation, type): - return f"{annotation.__module__}.{annotation.__qualname__}" - return repr(annotation) + yield entry_path, type(entry_value), entry_value diff --git a/motrix_env_core/src/motrix_env_core/manager/resolver.py b/motrix_env_core/src/motrix_env_core/manager/resolver.py index f3c1f47a..c250a187 100644 --- a/motrix_env_core/src/motrix_env_core/manager/resolver.py +++ b/motrix_env_core/src/motrix_env_core/manager/resolver.py @@ -33,11 +33,9 @@ Map, ResolvedLeaf, ResolvedTree, - TreeClassDef, is_kernel_data, iter_leaves, leaf_value_type, - tree_def_for, ) from motrix_env_core.manager.observations import ObservationGroupEntry, ObsTerm from motrix_env_core.manager.rewards import RewardTerm @@ -473,7 +471,7 @@ def _resolve_manager_context( dt=np.float32(self._env.cfg.ctrl_dt), sim_reset_requested=torch.from_numpy(self._env._sim_reset_requested), ) - tree_index = self._register_tree("manager_context", value)[0] + tree_index = self._register_tree("manager_context", value) return ResolvedManagerContext(tree_index) def _resolve_command_hooks(self, function_name: str) -> tuple[tuple[str, Callable], ...]: @@ -646,33 +644,32 @@ def _register_tree( self, source_name: str, value: Any, - ) -> tuple[int, TreeClassDef]: - """Validate one kernel-data value against its tree schema and record its tree. + ) -> int: + """Validate one kernel-data value's leaves and record its tree. Leaves are captured by reference: array leaves are refreshed in place by their owning runtime objects and scalar/shared leaves are reused unchanged. """ - tree_def = tree_def_for(value) leaves: dict[tuple[str, ...], ResolvedLeaf] = {} - for leaf_def, leaf_value in iter_leaves(value, tree_def): - value_type = leaf_value_type(leaf_def.annotation) + for leaf_path, annotation, leaf_value in iter_leaves(value): + value_type = leaf_value_type(annotation) if value_type is None: raise TypeError( - f"{source_name} KernelData {'.'.join(leaf_def.path)} expected a tensor, an ndarray, or a numeric " + f"{source_name} KernelData {'.'.join(leaf_path)} expected a tensor, an ndarray, or a numeric " "scalar." ) if value_type in (np.ndarray, torch.Tensor): if not isinstance(leaf_value, value_type): - raise TypeError(f"{source_name} field {'.'.join(leaf_def.path)!r} must be {value_type.__name__}.") + raise TypeError(f"{source_name} field {'.'.join(leaf_path)!r} must be {value_type.__name__}.") elif type(leaf_value) is not value_type: raise TypeError( - f"{source_name} field {'.'.join(leaf_def.path)!r} must return " + f"{source_name} field {'.'.join(leaf_path)!r} must return " f"{value_type.__name__}, got {type(leaf_value).__name__}." ) - leaves[leaf_def.path] = ResolvedLeaf(leaf_value) + leaves[leaf_path] = ResolvedLeaf(leaf_value) self._prepared_trees.append(ResolvedTree(value, leaves)) - return len(self._prepared_trees) - 1, tree_def + return len(self._prepared_trees) - 1 def _validate_observation_spaces(self, groups: dict[str, ObservationGroupEntry]) -> None: expected_policy = self._env.policy_observation_space.shape diff --git a/motrix_env_core/tests/test_kernel_data.py b/motrix_env_core/tests/test_kernel_data.py index 555e787e..729747c1 100644 --- a/motrix_env_core/tests/test_kernel_data.py +++ b/motrix_env_core/tests/test_kernel_data.py @@ -13,7 +13,6 @@ is_kernel_data, iter_leaves, leaf_value_type, - tree_def_for, ) @@ -40,8 +39,8 @@ class _Tracking: @dataclass(frozen=True, slots=True) -class _MappedValues: - values: Map +class _Cyclic: + child: "_Cyclic" def test_kernel_data_is_frozen_and_compiles_nested_inherited_leaves() -> None: @@ -52,14 +51,14 @@ def test_kernel_data_is_frozen_and_compiles_nested_inherited_leaves() -> None: with pytest.raises(FrozenInstanceError): value.phase = np.float32(1.0) # type: ignore[misc] - leaves = [(leaf.path, leaf_value) for leaf, leaf_value in iter_leaves(value, tree_def_for(value))] + leaves = [(path, leaf_value) for path, _, leaf_value in iter_leaves(value)] assert [path for path, _ in leaves] == [ ("robot", "velocity"), ("robot", "pose", "position"), ("robot", "pose", "orientation"), ("phase",), ] - annotations = [leaf.annotation for leaf, _ in iter_leaves(value, tree_def_for(value))] + annotations = [annotation for _, annotation, _ in iter_leaves(value)] assert leaf_value_type(annotations[0]) is np.ndarray assert leaf_value_type(annotations[-1]) is np.float32 @@ -71,41 +70,44 @@ class _Tensors: buffer = torch.zeros((2, 3), dtype=torch.float32) value = _Tensors(buffer) - leaves = list(iter_leaves(value, tree_def_for(value))) - assert [leaf_value for _, leaf_value in leaves] == [buffer] - assert leaf_value_type(leaves[0][0].annotation) is torch.Tensor + leaves = list(iter_leaves(value)) + assert [leaf_value for _, _, leaf_value in leaves] == [buffer] + assert leaf_value_type(leaves[0][1]) is torch.Tensor -def test_kernel_data_schema_fingerprint_ignores_values() -> None: - first_value = _Tracking( - _Robot(np.asarray([[1.0]], dtype=np.float32), _Pose(np.asarray([[2.0]]), np.asarray([3.0]))), - np.float32(0.5), - ) - second_value = _Tracking( - _Robot(np.asarray([[4.0]], dtype=np.float32), _Pose(np.asarray([[5.0]]), np.asarray([6.0]))), - np.float32(0.75), - ) - first = tree_def_for(first_value) - second = tree_def_for(second_value) - assert first.fingerprint == second.fingerprint - assert first.logical_type.__module__ == _Tracking.__module__ - - -def test_map_schema_changes_tree_identity() -> None: - first = _MappedValues(Map({"value": np.float32(1.0)})) - second = _MappedValues(Map({"value": np.asarray([[1.0]], dtype=np.float32)})) - first_tree = tree_def_for(first) - second_tree = tree_def_for(second) - assert first_tree.fingerprint != second_tree.fingerprint - first_map = first_tree.fields[0].tree_def - second_map = second_tree.fields[0].tree_def - assert first_map.fingerprint != second_map.fingerprint - - -def test_iter_leaves_rejects_changed_map_keys() -> None: - schema = tree_def_for(_MappedValues(Map({"a": np.float32(1.0), "b": np.float32(2.0)}))) - with pytest.raises(TypeError, match="has keys"): - list(iter_leaves(_MappedValues(Map({"a": np.float32(1.0)})), schema)) +def test_iter_leaves_rejects_wrong_class_field() -> None: + @dataclass(frozen=True, slots=True) + class _Inner: + x: np.ndarray + + @dataclass(frozen=True, slots=True) + class _Outer: + inner: _Inner + + @dataclass(frozen=True, slots=True) + class _Impostor: + x: np.ndarray + + with pytest.raises(TypeError, match="must be _Inner, got _Impostor"): + list(iter_leaves(_Outer(_Impostor(np.zeros(1))))) # type: ignore[arg-type] + + +def test_iter_leaves_rejects_class_cycle() -> None: + @dataclass(frozen=True, slots=True) + class _Node: + child: _Cyclic + + with pytest.raises(TypeError, match="cycle detected"): + list(iter_leaves(_Node(_Cyclic(None)))) # type: ignore[arg-type] + + +def test_iter_leaves_rejects_map_annotation_without_map_value() -> None: + @dataclass(frozen=True, slots=True) + class _HoldsMap: + values: Map + + with pytest.raises(TypeError, match="requires a concrete Map value"): + list(iter_leaves(_HoldsMap({"a": np.float32(1.0)}))) # type: ignore[arg-type] def test_kernel_data_rejects_unsupported_leaf() -> None: From 3ef701c0a58bb67b5105421e3fce2a45f928ba03 Mon Sep 17 00:00:00 2001 From: zilch <147668916@qq.com> Date: Sat, 7 Nov 2026 02:02:25 +0800 Subject: [PATCH 41/41] fix(kernel_data): drop the cache on type-hint resolution 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. --- motrix_env_core/src/motrix_env_core/manager/kernel_data/tree.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/motrix_env_core/src/motrix_env_core/manager/kernel_data/tree.py b/motrix_env_core/src/motrix_env_core/manager/kernel_data/tree.py index da65aa1e..370385a4 100644 --- a/motrix_env_core/src/motrix_env_core/manager/kernel_data/tree.py +++ b/motrix_env_core/src/motrix_env_core/manager/kernel_data/tree.py @@ -5,7 +5,6 @@ from collections.abc import Iterator from dataclasses import fields, is_dataclass -from functools import cache from typing import Any, get_origin, get_type_hints import numpy as np @@ -49,7 +48,6 @@ def leaf_value_type(annotation: Any) -> type | None: return None -@cache def _type_hints(cls: type) -> dict[str, Any]: return get_type_hints(cls, include_extras=True)