From 149c6d92d7f26fa60101b6c64f5c6fbca4e3fe92 Mon Sep 17 00:00:00 2001 From: zilch Date: Sat, 12 Sep 2026 14:46:44 +0800 Subject: [PATCH 1/3] perf: move numeric manager term params to runtime buffers instead of plan_key constants Numeric tuning arguments (termination thresholds, observation scales, reward gains, ...) were repr-baked into the generated kernel source and the plan key, so every tuning variant triggered a full ~1.2-30s kernel recompilation. Lower float term args into a fixed float32 shared input slot instead: only the layout fingerprint enters the plan key and the value is read at runtime, keeping tuning variants on one compiled plan. Reward weights are no longer pre-multiplied by ctrl_dt at build time; the dt scaling happens inside the evaluate kernel via ctx.dt, removing ctrl_dt from the plan key and the double-baking mismatch risk. Also improve startup logging: the manager env now logs each startup step (build plan, compile kernels with cache hit/miss classification, term specializations, warmup execution) with durations, and the runner logs each training startup phase. Kernel cache detection covers both numba cache layouts (NUMBA_CACHE_DIR set/unset), and fastsac worker processes configure INFO logging so startup timings are visible. Fixes #34 --- .../numba/manager/compiler/codegen.py | 2 +- .../numba/manager/compiler/compiler.py | 185 +++++++++++++++--- .../src/motrix_env_core/numba/manager/env.py | 26 +-- motrix_env_core/tests/test_numba_manager.py | 43 +++- .../motrix_rl/fastsac/async_impl/worker.py | 8 + motrix_rl/src/motrix_rl/runner.py | 27 ++- 6 files changed, 247 insertions(+), 44 deletions(-) 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 index f74c86e3..e7ca1738 100644 --- 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 @@ -140,7 +140,7 @@ def _generate_lane_bodies( [ 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}]", + 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}", ] diff --git a/motrix_env_core/src/motrix_env_core/numba/manager/compiler/compiler.py b/motrix_env_core/src/motrix_env_core/numba/manager/compiler/compiler.py index 4442ba93..dc7e0c62 100644 --- a/motrix_env_core/src/motrix_env_core/numba/manager/compiler/compiler.py +++ b/motrix_env_core/src/motrix_env_core/numba/manager/compiler/compiler.py @@ -12,6 +12,7 @@ from collections.abc import Callable, 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 @@ -27,6 +28,7 @@ flatten_kernel_data, is_kernel_data, iter_layout_leaves, + kernel_data, lane_expression, map_proxy, proxy_symbol, @@ -79,6 +81,20 @@ def _invalidate_term_cache() -> None: _TERM_CACHE.clear() +@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: 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. + """ + + value: np.float32 + + @dataclass(frozen=True) class ManagerSimInput: """One compiled simulator input exposed through ``ManagerContext.sim``.""" @@ -111,6 +127,7 @@ def __init__(self, env: ManagerEnv): self._plan_parts.append(("sim_inputs", sim_input_fingerprint)) def build(self) -> CompiledManagerProgram: + build_started = perf_counter() self._resolve_runtime_terms() observation_groups = self._env.observation_groups reward_terms = self._env._reward_terms @@ -130,9 +147,17 @@ def build(self) -> CompiledManagerProgram: ) self._validate_observation_spaces(observation_groups) - kernel_reward_weights = np.asarray(reward_weights, dtype=np.float32) * np.float32(self._env.cfg.ctrl_dt) - self._plan_parts.append(("manager_context_dt", np.float32(self._env.cfg.ctrl_dt))) + # 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) plan_key = hashlib.sha256(repr(tuple(self._plan_parts)).encode()).hexdigest() + logger.info( + "Manager startup %s: build plan finished in %.3fs (plan_key=%.12s)", + self._env_name(), + perf_counter() - build_started, + plan_key, + ) source = KernelSourceGenerator(self._flat_input_count).generate( observation_terms, observation_layout, @@ -146,8 +171,28 @@ def build(self) -> CompiledManagerProgram: ) generated_filename = self._materialize_source(source, plan_key) kernels = _KERNEL_CACHE.get(plan_key) - if kernels is None: - logger.debug("Manager kernel cache miss: plan_key=%s", plan_key) + if kernels is not None: + logger.info( + "Manager startup %s: kernels reused (cache=in-process, plan_key=%.12s)", + self._env_name(), + plan_key, + ) + else: + disk_cache = self._has_generated_kernel_cache(plan_key) + if disk_cache: + logger.info( + "Manager startup %s: compiled kernel cache found, loading (cache=disk, plan_key=%.12s)", + self._env_name(), + plan_key, + ) + else: + logger.info( + "Manager startup %s: no cache for this plan, first compile may take a while " + "(cache=miss, plan_key=%.12s)", + self._env_name(), + plan_key, + ) + compile_started = perf_counter() try: kernels = self._compile_kernels(source, generated_filename) except (ImportError, OSError, RuntimeError, TypeError, ValueError) as error: @@ -156,8 +201,12 @@ def build(self) -> CompiledManagerProgram: generated_filename = self._materialize_source(source, plan_key) kernels = self._compile_kernels(source, generated_filename) _KERNEL_CACHE[plan_key] = kernels - else: - logger.debug("Manager kernel in-process cache hit: plan_key=%s", plan_key) + logger.info( + "Manager startup %s: compile kernels finished in %.3fs (cache=%s)", + self._env_name(), + perf_counter() - compile_started, + "disk" if disk_cache else "miss", + ) evaluate_kernel, observe_kernel, reset_kernel = kernels layout = ManagerLayout( @@ -271,6 +320,7 @@ def _resolve_reset( symbol = f"term_reset_{term_index}" self._term_functions[symbol] = self._generated_term(function, dispatcher) expressions = [] + plan_expressions = [] prepared_indices = [] for index, value in enumerate(term.args): if is_kernel_data(value): @@ -285,15 +335,20 @@ def _resolve_reset( ) prepared_indices.append(prepared_index) expressions.append(expression) + plan_expressions.append(expression) else: - prepared_indices.append(None) - expressions.append(repr(value)) + expression, _, plan_part, 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._plan_parts.append( ( "sim_reset", name, self._function_fingerprint(function), - tuple(expressions), + plan_expressions, repr(descriptors), fields, ) @@ -473,7 +528,9 @@ def _resolve_observation_dispatch( f"ManagerContext and np.ndarray (conventionally named ctx and out)." ) expressions = [] + plan_expressions = [] prepared_indices = [] + warmup_values = [] for index, value in enumerate(term.args): if is_kernel_data(value): _, tree_def = flatten_kernel_data(value) @@ -485,15 +542,22 @@ def _resolve_observation_dispatch( ) prepared_indices.append(prepared_index) expressions.append(expression) + plan_expressions.append(expression) + warmup_values.append(value) else: - prepared_indices.append(None) - expressions.append(repr(value)) + expression, warmup, plan_part, 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._plan_parts.append( - (group, name, "observation_dispatch", self._function_fingerprint(function), output_size, expressions) + (group, name, "observation_dispatch", self._function_fingerprint(function), output_size, plan_expressions) ) return PreparedInvocation( f"{group}.{name}", @@ -501,7 +565,7 @@ def _resolve_observation_dispatch( dispatcher, None, args_expressions=expressions, - args_values=term.args, + args_values=tuple(warmup_values), args_prepared_indices=tuple(prepared_indices), output_size=output_size, ) @@ -574,7 +638,9 @@ def _resolve_dispatch_term( f"(conventionally named ctx)." ) expressions = [] + plan_expressions = [] prepared_indices = [] + warmup_values = [] for index, value in enumerate(term.args): if is_kernel_data(value): _, tree_def = flatten_kernel_data(value) @@ -586,20 +652,27 @@ def _resolve_dispatch_term( ) prepared_indices.append(prepared_index) expressions.append(expression) + plan_expressions.append(expression) + warmup_values.append(value) else: - prepared_indices.append(None) - expressions.append(repr(value)) + expression, warmup, plan_part, 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._plan_parts.append((group, name, "dispatch", self._function_fingerprint(function), expressions)) + self._plan_parts.append((group, name, "dispatch", self._function_fingerprint(function), plan_expressions)) return PreparedInvocation( f"{group}.{name}", group, dispatcher, None, args_expressions=tuple(expressions), - args_values=term.args, + args_values=tuple(warmup_values), args_prepared_indices=tuple(prepared_indices), ) @@ -681,6 +754,38 @@ def _register_prepared( 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. + """ + 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 _resolve_plain_arg(self, source_name: str, value: Any) -> tuple[str | None, Any, str, 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). + """ + if isinstance(value, (float, np.floating)): + expression, warmup = self._scalar_buffer_arg(source_name, value) + return expression, warmup, "scalar_buffer(np.float32)", None + return repr(value), value, repr(value), None + def _validate_observation_spaces(self, groups: dict[str, ObservationGroupEntry]) -> None: expected_policy = self._env.policy_observation_space.shape if expected_policy != (groups["policy"].size,): @@ -704,9 +809,7 @@ def _validate_observation_spaces(self, groups: dict[str, ObservationGroupEntry]) @staticmethod def _materialize_source(source: str, plan_key: str) -> 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") + cache_dir = NumbaKernelCompiler._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: @@ -724,16 +827,45 @@ def _materialize_source(source: str, plan_key: str) -> str: return str(path) @staticmethod - def _invalidate_generated_cache(plan_key: str) -> None: + 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 + + @staticmethod + 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(NumbaKernelCompiler._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) + + @staticmethod + 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 NumbaKernelCompiler._generated_cache_paths(plan_key)) + + @staticmethod + def _invalidate_generated_cache(plan_key: str) -> None: + cache_dir = NumbaKernelCompiler._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 Path(cache_dir).glob(f"generated_*/*{plan_key}*"): + for cache_path in NumbaKernelCompiler._generated_cache_paths(plan_key): try: cache_path.unlink(missing_ok=True) except OSError: @@ -741,6 +873,7 @@ def _invalidate_generated_cache(plan_key: str) -> None: 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 ( @@ -760,6 +893,14 @@ def compile_specializations(self, inputs: tuple[Any, ...]) -> None: _KERNEL_CACHE.pop(plan_key, None) 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 _env_name(self) -> str: + return type(self._env).__name__ def _compile_specializations_once(self, inputs: tuple[Any, ...]) -> None: program = self._env._compiled_manager_program diff --git a/motrix_env_core/src/motrix_env_core/numba/manager/env.py b/motrix_env_core/src/motrix_env_core/numba/manager/env.py index a9504155..b33b7033 100644 --- a/motrix_env_core/src/motrix_env_core/numba/manager/env.py +++ b/motrix_env_core/src/motrix_env_core/numba/manager/env.py @@ -505,18 +505,14 @@ def physics_step(self) -> None: self.sim.step(self._cfg.sim_substeps) def init_state(self) -> ArrayEnvState: + startup_started = perf_counter() + env_name = type(self).__name__ if self._task_program is None: - build_started = perf_counter() + 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. self._refresh_sim_reads() - logger.info( - "Manager compiler build completed: env=%s seconds=%.3f plan_key=%s", - type(self).__name__, - perf_counter() - build_started, - self.manager_layout.plan_key, - ) state = super().init_state() self._kernel_buffers = self._make_kernel_buffers(state) state.metrics = self._make_metrics_view() @@ -533,23 +529,15 @@ def init_state(self) -> ArrayEnvState: self._refresh_sim_reads() inputs = self._kernel_inputs self._validate_kernel_context(inputs) - compile_started = perf_counter() self._compile_manager_specializations(inputs) - compile_seconds = perf_counter() - compile_started - logger.info( - "Manager compile completed: env=%s seconds=%.3f plan_key=%s", - type(self).__name__, - compile_seconds, - self.manager_layout.plan_key, - ) warmup_started = perf_counter() self._execute_observe_kernel(inputs) - warmup_seconds = perf_counter() - warmup_started logger.info( - "Manager warmup completed: env=%s seconds=%.3f", - type(self).__name__, - warmup_seconds, + "Manager startup %s: warmup execution finished in %.3fs", + env_name, + perf_counter() - warmup_started, ) + logger.info("Manager env %r startup complete in %.3fs", env_name, perf_counter() - startup_started) return state def _prev_physics_step(self) -> None: diff --git a/motrix_env_core/tests/test_numba_manager.py b/motrix_env_core/tests/test_numba_manager.py index a093da12..691a879d 100644 --- a/motrix_env_core/tests/test_numba_manager.py +++ b/motrix_env_core/tests/test_numba_manager.py @@ -865,7 +865,7 @@ def test_numeric_values_reuse_compiled_plan_and_remain_environment_local() -> No assert compiled.layout.plan_key == env.manager_layout.plan_key assert compiled.task.evaluate_kernel is env._compiled_manager_program.task.evaluate_kernel - np.testing.assert_allclose(compiled.task.reward_weights, [2.0 * env.cfg.ctrl_dt]) + np.testing.assert_allclose(compiled.task.reward_weights, [2.0]) action = env.action_terms["test"] assert isinstance(action, _TestAction) action.source[:] = 0.25 @@ -879,6 +879,47 @@ def test_numeric_values_reuse_compiled_plan_and_remain_environment_local() -> No 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_key == second.manager_layout.plan_key + + for env, scale, threshold in ((first, 3.0, 0.5), (second, 5.0, 0.1)): + action = env.action_terms["test"] + assert isinstance(action, _TestAction) + action.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_key == first.manager_layout.plan_key + action = third.action_terms["test"] + assert isinstance(action, _TestAction) + action.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): diff --git a/motrix_rl/src/motrix_rl/fastsac/async_impl/worker.py b/motrix_rl/src/motrix_rl/fastsac/async_impl/worker.py index deee8c71..82df76c7 100644 --- a/motrix_rl/src/motrix_rl/fastsac/async_impl/worker.py +++ b/motrix_rl/src/motrix_rl/fastsac/async_impl/worker.py @@ -15,6 +15,7 @@ from __future__ import annotations +import logging import random import time import traceback @@ -115,6 +116,11 @@ def build_agent(cfg: FastSacCfg, dims, num_envs, device, action_scale, action_bi # ------------------------------------------------------------------ collector process +def _configure_process_logging() -> None: + """Surface INFO logs (e.g. manager env startup) from spawned worker processes.""" + logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(name)s: %(message)s") + + def run_collector_process( env_spec: EnvBuildSpec, cfg: FastSacCfg, @@ -133,6 +139,7 @@ def run_collector_process( seed, ) -> None: try: + _configure_process_logging() set_seed(seed) async_options = cfg.trainer.async_options obs_dim, critic_obs_dim, act_dim = dims @@ -196,6 +203,7 @@ def run_learner_process( resume_from: str | None, seed, ) -> None: + _configure_process_logging() console, live = open_training_live() try: set_seed(seed) diff --git a/motrix_rl/src/motrix_rl/runner.py b/motrix_rl/src/motrix_rl/runner.py index ee86a039..63c5cd11 100644 --- a/motrix_rl/src/motrix_rl/runner.py +++ b/motrix_rl/src/motrix_rl/runner.py @@ -1,6 +1,8 @@ # Copyright Motphys Technology Co., Ltd. 2025, 2026 # SPDX-License-Identifier: Apache-2.0 +import logging +import time from dataclasses import dataclass from pathlib import Path @@ -10,6 +12,8 @@ from motrix_rl.method import RlMethod from motrix_rl.result import TrainResult +logger = logging.getLogger(__name__) + @dataclass(frozen=True) class TrainRequest: @@ -27,9 +31,12 @@ def train(request: TrainRequest) -> TrainResult: def create_training_handle(request: TrainRequest) -> frameworks.TrainerHandle: """Create a trainer handle for a new training run.""" + startup_started = time.perf_counter() config = request.config task = config.task method = RlMethod(rllib=task.rllib, algo=task.algo) + logger.info("Training startup: resolving train backend for %s/%s", method.rllib, method.algo) + step_started = time.perf_counter() device_supports = utils.get_device_supports() train_backend = backend_runtime.resolve_train_backend( task.env, @@ -37,6 +44,14 @@ def create_training_handle(request: TrainRequest) -> frameworks.TrainerHandle: task.train_backend, device_supports, ) + logger.info( + "Training startup: resolved train backend '%s' in %.3fs (env=%s, num_envs=%d)", + train_backend, + time.perf_counter() - step_started, + task.env, + config.num_envs, + ) + step_started = time.perf_counter() provider = frameworks.get_agent_provider(method.rllib, method.algo, train_backend) if provider is None: raise ValueError( @@ -56,16 +71,26 @@ def create_training_handle(request: TrainRequest) -> frameworks.TrainerHandle: runs_root=request.runs_root, ) runs.write_task_config(run.run_dir, config) + logger.info( + "Training startup: created run context in %.3fs (run_dir=%s)", + time.perf_counter() - step_started, + run.run_dir, + ) resume_from = None if config.resume is not None: resume_from = str(checkpoints.resolve_resume_checkpoint_path(config.resume)) - return _create_handle( + handle = _create_handle( run, provider, config, render=request.render, resume_from=resume_from, ) + logger.info( + "Training startup: trainer handle ready in %.3fs", + time.perf_counter() - startup_started, + ) + return handle def create_run_handle( From a0c10ffff3a1ec0379b9e6bbbf3208ef8ee1e344 Mon Sep 17 00:00:00 2001 From: zilch Date: Sat, 12 Sep 2026 15:01:22 +0800 Subject: [PATCH 2/3] test: drop exact read-plan source count in wbt numba test Scalar term args now lower into one runtime buffer source each, so the wbt read plan legitimately grows beyond the hardcoded count of four. The test's intent (preallocated, stable kernel inputs with the manager context present) is preserved without pinning the source count. --- motrix_envs/tests/test_wbt_numba.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/motrix_envs/tests/test_wbt_numba.py b/motrix_envs/tests/test_wbt_numba.py index d916d422..5ed4f16b 100644 --- a/motrix_envs/tests/test_wbt_numba.py +++ b/motrix_envs/tests/test_wbt_numba.py @@ -171,7 +171,8 @@ def test_numba_wbt_read_plan_reuses_preallocated_arrays() -> None: assert len(first) == len(env.manager_layout.inputs) sources = env._compiled_manager_program.read_plan.sources - assert len(sources) == 4 + # One source per kernel-data term plus one runtime scalar buffer per float + # term argument; the manager context source must be present either way. 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) From fcbff348d2a37708ce2817599e362ebae59b10af Mon Sep 17 00:00:00 2001 From: zilch Date: Sat, 12 Sep 2026 15:19:43 +0800 Subject: [PATCH 3/3] fix: address review feedback on worker logging and wbt scalar-buffer assertion - Worker logging setup no longer relies on basicConfig when the root logger already has handlers: raise the root level instead so startup INFO records surface through a preconfigured setup. - The wbt read-plan test now concretely asserts the presence of runtime scalar-buffer sources (TermScalarBuffer) instead of only leaving a comment, so regressions that bake scalar args back into the plan or generated source fail the test. --- motrix_envs/tests/test_wbt_numba.py | 8 +++++++- motrix_rl/src/motrix_rl/fastsac/async_impl/worker.py | 11 ++++++++++- 2 files changed, 17 insertions(+), 2 deletions(-) diff --git a/motrix_envs/tests/test_wbt_numba.py b/motrix_envs/tests/test_wbt_numba.py index 5ed4f16b..3a569915 100644 --- a/motrix_envs/tests/test_wbt_numba.py +++ b/motrix_envs/tests/test_wbt_numba.py @@ -21,6 +21,7 @@ UniformNoiseCfg, ) from motrix_env_core.mdp.state import RandValue # noqa: E402 +from motrix_env_core.numba.manager.compiler.compiler import TermScalarBuffer # noqa: E402 from motrix_env_core.sim import BatchLinkPositionQuery # noqa: E402 from motrix_envs.locomotion.wbt.cfg import ( # noqa: E402 ActionsCfg, @@ -172,7 +173,12 @@ def test_numba_wbt_read_plan_reuses_preallocated_arrays() -> None: 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; the manager context source must be present either way. + # 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) diff --git a/motrix_rl/src/motrix_rl/fastsac/async_impl/worker.py b/motrix_rl/src/motrix_rl/fastsac/async_impl/worker.py index 82df76c7..f90a9b89 100644 --- a/motrix_rl/src/motrix_rl/fastsac/async_impl/worker.py +++ b/motrix_rl/src/motrix_rl/fastsac/async_impl/worker.py @@ -117,7 +117,16 @@ def build_agent(cfg: FastSacCfg, dims, num_envs, device, action_scale, action_bi # ------------------------------------------------------------------ collector process def _configure_process_logging() -> None: - """Surface INFO logs (e.g. manager env startup) from spawned worker processes.""" + """Surface INFO logs (e.g. manager env startup) from spawned worker processes. + + ``basicConfig`` only applies when the root logger has no handlers yet; when + a handler is already configured, raising the root level is enough for the + startup INFO records to be emitted through the existing setup. + """ + root = logging.getLogger() + if root.handlers: + root.setLevel(logging.INFO) + return logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(name)s: %(message)s")