feat: internalize bench scripts as the mxlab bench command

bench/README.md

@@ -1,24 +1,24 @@
-# Environment benchmarks
+# Benchmarks
 
-Activate the project environment before running:
+All benchmarks live in the `mxlab` CLI.
 
-```bash
-source .venv/bin/activate
+## Environment `env.step()` benchmark
 
+```bash
 # Whole-machine throughput: 10 seconds, one worker per usable NUMA node
-python bench/bench_env.py --env g1-wbt-dance
+mxlab bench --mode env --env g1-wbt-dance
 
 # Batch sweep with JSONL output
-python bench/bench_env.py --env g1-wbt-dance --duration 10 --num-envs 2048 8192 --json
+mxlab bench --mode env --env g1-wbt-dance --duration 10 --num-envs 2048 8192 --json
 
 # Compare two environments
-python bench/bench_env.py --env cartpole --compare-env cartpole --duration 5
+mxlab bench --mode env --env cartpole --compare-env cartpole --duration 5
 
 # Single process, preserving the caller's CPU/memory placement
-python bench/bench_env.py --duration 5 --numa single
+mxlab bench --mode env --duration 5 --numa single
 
 # Single-process step latency and CPU utilization
-python bench/bench_env.py --steps 1000 --num-envs 2048
+mxlab bench --mode env --steps 1000 --num-envs 2048
 ```
 
 `--num-envs` is the total batch, split across usable NUMA nodes. Workers respect
@@ -31,4 +31,24 @@ finish after the deadline are excluded. This measures `env.step()`, not policy
 inference or training data transfer. JSONL includes per-worker counts and placement.
 
 `--steps` and `--duration` are mutually exclusive. `--breakdown` requires step-count
-mode and cannot be combined with `--json`. See `--help` for all options.
+mode and cannot be combined with `--json`.
+
+## MotrixSim physics-only benchmark
+
+```bash
+# In-process measurement of one robot
+mxlab bench --mode physics --robot g1-29dof --num-envs 2048
+
+# One pinned worker per NUMA node via numactl
+mxlab bench --mode physics --numa shard
+```
+
+Times `model.step_n(data, nstep)` on a robot-on-flat-ground scene with state reset
+before every timed call; no task env, no observation/reward/termination. `--numa`
+selects `single`, `node` (pinned to node 0), or `shard` (one numactl-pinned worker
+per NUMA node; requires numactl on PATH).
+
+See `mxlab bench --help` for all options.
+
+The remaining script in this directory, `bench_live_metrics.py`, is a standalone
+simulator-side tool that is not part of the CLI.

bench/bench_physics.py → motrixlab/src/motrixlab/benchmark/physics.py

@@ -18,10 +18,13 @@
   shard   spawn one worker per NUMA node, each pinned via numactl, each running
           ``--num-envs / num_nodes`` envs concurrently; aggregate throughput is
           the sum across workers. Reproduces the per-node process sharding
-          validated for per-node process sharding on a many-core host.
+          validated on a many-core host.
 
 Reports per-call latency (median + p10/p90), batch steps/s (step_n calls per
 second over all parallel envs in the process), and total env steps/s.
+
+This module imports the simulation runtime at import time and is therefore
+loaded lazily by ``motrixlab.commands.bench`` only for ``--mode physics``.
 """
 
 from __future__ import annotations
@@ -38,12 +41,25 @@
 
 import motrixsim as mtx
 import numpy as np
-from report import render_table
 
 from motrix_env_core import registry
 from motrix_env_core.base import SimCfg
 from motrix_env_motrixsim.compiler import MotrixSimSceneCompiler
 from motrix_envs.config.scene import StandardSceneCfg, StandardSceneObjsCfg
+from motrixlab.benchmark.report import render_table
+
+RESULT_HEADERS = [
+    "mode",
+    "robot",
+    "num_envs",
+    "median ms",
+    "p10 ms",
+    "p90 ms",
+    "batch steps/s",
+    "total env steps/s",
+]
+
+_WARMUP_DEFAULT = 20
 
 
 def _detect_num_nodes() -> int:
@@ -76,12 +92,11 @@ def _current_binding() -> str:
 
 
 def _device_context() -> dict:
-    ctx = {
+    return {
         "num_cpus": os.cpu_count(),
         "num_numa_nodes": _detect_num_nodes(),
         "binding": _current_binding(),
     }
-    return ctx
 
 
 def measure(
@@ -151,18 +166,6 @@ def reset_state() -> None:
     }
 
 
-_RESULT_HEADERS = [
-    "mode",
-    "robot",
-    "num_envs",
-    "median ms",
-    "p10 ms",
-    "p90 ms",
-    "batch steps/s",
-    "total env steps/s",
-]
-
-
 def _result_row(label: str, r: dict) -> list[str]:
     return [
         label,
@@ -176,22 +179,13 @@ def _result_row(label: str, r: dict) -> list[str]:
     ]
 
 
-def run_single(args) -> tuple[list[list[str]], list[str]]:
-    r = measure(
-        robot=args.robot,
-        num_envs=args.num_envs,
-        dt=args.dt,
-        solver_iterations=args.solver_iterations,
-        nstep=args.nstep,
-        warmup=args.warmup,
-        iters=args.iters,
-    )
+def run_single(args: argparse.Namespace) -> tuple[list[list[str]], list[str]]:
+    r = _measure_from_args(args)
     return ([_result_row(f"single-process ({args.numa})", r)], [])
 
 
-def _worker_entry(args) -> None:
-    """Internal: run one measurement and write JSON to --out-json. Used by node/shard drivers."""
-    r = measure(
+def _measure_from_args(args: argparse.Namespace) -> dict:
+    return measure(
         robot=args.robot,
         num_envs=args.num_envs,
         dt=args.dt,
@@ -200,23 +194,31 @@ def _worker_entry(args) -> None:
         warmup=args.warmup,
         iters=args.iters,
     )
+
+
+def _worker_entry(args: argparse.Namespace) -> None:
+    """Internal: run one measurement and write JSON to --out-json. Used by node/shard drivers."""
+    r = _measure_from_args(args)
     r["node"] = args.node
     Path(args.out_json).write_text(json.dumps(r), encoding="utf-8")
 
 
-def _numactl_wrap(node: int, py_args: list[str]) -> list[str]:
+def _worker_command(args: argparse.Namespace, node: int, num_envs: int, out_json: str) -> list[str]:
+    """Build a numactl-wrapped ``python -m`` command for one pinned worker."""
     numactl = shutil.which("numactl")
     if numactl is None:
         raise RuntimeError("numactl not found on PATH; install numactl for --numa node/shard")
-    return [numactl, "--cpunodebind", str(node), "--membind", str(node)] + py_args
-
-
-def _measure_via_subprocess(args, node: int, num_envs: int, out_json: str) -> dict:
-    """Spawn a pinned worker that runs `measure` and writes JSON; block until done."""
-    script = str(Path(sys.argv[0]).resolve())
-    py_args = [
+    return [
+        numactl,
+        "--cpunodebind",
+        str(node),
+        "--membind",
+        str(node),
         sys.executable,
-        script,
+        "-m",
+        "motrixlab.commands.bench",
+        "--mode",
+        "physics",
         "--numa",
         "_worker",
         "--node",
@@ -234,17 +236,13 @@ def _measure_via_subprocess(args, node: int, num_envs: int, out_json: str) -> di
         "--warmup",
         str(args.warmup),
         "--iters",
-        args.iters if isinstance(args.iters, str) else str(args.iters),
+        str(args.iters),
         "--out-json",
         out_json,
     ]
-    cmd = _numactl_wrap(node, py_args)
-    env = os.environ.copy()
-    subprocess.run(cmd, check=True, env=env)
-    return json.loads(Path(out_json).read_text())
 
 
-def run_node(args) -> tuple[list[list[str]], list[str]]:
+def run_node(args: argparse.Namespace) -> tuple[list[list[str]], list[str]]:
     """Re-exec the measurement pinned to NUMA node 0."""
     notes: list[str] = []
     num_nodes = _detect_num_nodes()
@@ -253,13 +251,15 @@ def run_node(args) -> tuple[list[list[str]], list[str]]:
     with tempfile.NamedTemporaryFile(suffix=".json", delete=False) as f:
         out_json = f.name
     try:
-        r = _measure_via_subprocess(args, node=0, num_envs=args.num_envs, out_json=out_json)
+        cmd = _worker_command(args, node=0, num_envs=args.num_envs, out_json=out_json)
+        subprocess.run(cmd, check=True, env=os.environ.copy())
+        r = json.loads(Path(out_json).read_text())
         return ([_result_row("node-pinned (cpunodebind=0, membind=0)", r)], notes)
     finally:
         Path(out_json).unlink(missing_ok=True)
 
 
-def run_shard(args) -> tuple[list[list[str]], list[str]]:
+def run_shard(args: argparse.Namespace) -> tuple[list[list[str]], list[str]]:
     """Spawn one pinned worker per NUMA node, each running num_envs/num_nodes envs."""
     num_nodes = _detect_num_nodes()
     if num_nodes < 2:
@@ -273,32 +273,7 @@ def run_shard(args) -> tuple[list[list[str]], list[str]]:
         fd, tmp_path = tempfile.mkstemp(suffix=f"_node{node}.json")
         os.close(fd)
         tmp_paths.append(tmp_path)
-        script = str(Path(sys.argv[0]).resolve())
-        py_args = [
-            sys.executable,
-            script,
-            "--numa",
-            "_worker",
-            "--node",
-            str(node),
-            "--num-envs",
-            str(per_node),
-            "--robot",
-            args.robot,
-            "--dt",
-            str(args.dt),
-            "--solver-iterations",
-            str(args.solver_iterations),
-            "--nstep",
-            str(args.nstep),
-            "--warmup",
-            str(args.warmup),
-            "--iters",
-            args.iters if isinstance(args.iters, str) else str(args.iters),
-            "--out-json",
-            tmp_path,
-        ]
-        cmd = _numactl_wrap(node, py_args)
+        cmd = _worker_command(args, node=node, num_envs=per_node, out_json=tmp_path)
         procs.append((node, subprocess.Popen(cmd, env=os.environ.copy())))
     for node, p in procs:
         rc = p.wait()
@@ -323,46 +298,24 @@ def run_shard(args) -> tuple[list[list[str]], list[str]]:
     return rows, notes
 
 
-def main() -> None:
-    parser = argparse.ArgumentParser(description=__doc__)
-    parser.add_argument("--robot", default="g1-29dof")
-    parser.add_argument(
-        "--num-envs",
-        type=int,
-        nargs="+",
-        default=[2048],
-        help="Batch sizes to benchmark. Multiple values run a sweep.",
-    )
-    parser.add_argument("--dt", type=float, default=0.005)
-    parser.add_argument("--solver-iterations", type=int, default=3)
-    parser.add_argument("--nstep", type=int, default=1, help="Native substeps per step_n call.")
-    parser.add_argument("--warmup", type=int, default=20)
-    parser.add_argument("--iters", type=int, default=200)
-    parser.add_argument(
-        "--numa",
-        choices=["single", "node", "shard", "_worker"],
-        default="single",
-        help="single=in-process; node=re-exec pinned to node 0; shard=one worker per NUMA node.",
-    )
-    parser.add_argument("--node", type=int, default=0, help="(internal worker) which node.")
-    parser.add_argument("--out-json", type=str, default="", help="(internal worker) JSON output path.")
-    args = parser.parse_args()
+def main(parser: argparse.ArgumentParser, args: argparse.Namespace) -> None:
     if any(num_envs < 1 for num_envs in args.num_envs):
         parser.error("--num-envs values must be positive")
+    if args.warmup is None:
+        args.warmup = _WARMUP_DEFAULT
 
     if args.numa == "_worker":
         if not args.out_json:
-            raise SystemExit("--numa _worker requires --out-json")
+            parser.error("--numa _worker requires --out-json")
         if len(args.num_envs) != 1:
             parser.error("internal worker mode requires exactly one --num-envs value")
         args.num_envs = args.num_envs[0]
         _worker_entry(args)
         return
 
-    num_envs_values = args.num_envs
     rows: list[list[str]] = []
     notes: list[str] = []
-    for num_envs in num_envs_values:
+    for num_envs in args.num_envs:
         args.num_envs = num_envs
         if args.numa == "single":
             batch_rows, batch_notes = run_single(args)
@@ -376,15 +329,11 @@ def main() -> None:
                 notes.append(note)
 
     context = _device_context()
-    print(render_table(_RESULT_HEADERS, rows))
+    print(render_table(RESULT_HEADERS, rows))
     print(
         f"sim: robot={args.robot} dt={args.dt}s solver_iterations={args.solver_iterations} "
         f"nstep={args.nstep} warmup={args.warmup} iters={args.iters}"
     )
     print(f"context: cpus={context['num_cpus']} numa_nodes={context['num_numa_nodes']} binding=[{context['binding']}]")
     for note in notes:
         print(note)
-
-
-if __name__ == "__main__":
-    main()

bench/env_benchmark/runner.py → motrixlab/src/motrixlab/benchmark/runner.py

@@ -87,7 +87,7 @@ def run_steps(
                 results.append(
                     {
                         "env": name,
-                        "mode": args.mode,
+                        "mode": args.env_mode,
                         "num_envs": num_envs,
                         "steps": args.steps,
                         "warmup": args.warmup,
@@ -222,7 +222,7 @@ def run_duration(args: argparse.Namespace, num_envs: int) -> list[dict[str, obje
             results.append(
                 {
                     "env": name,
-                    "mode": args.mode,
+                    "mode": args.env_mode,
                     "num_envs": num_envs,
                     "num_envs_scope": "total",
                     "duration_seconds": args.duration,

bench/env_benchmark/workload.py → motrixlab/src/motrixlab/benchmark/workload.py

@@ -18,9 +18,10 @@ def load_runtime():
     import numpy as np
 
     from motrix_env_core import registry
+    from motrix_env_core.direct.env import DirectEnv
     from motrix_env_core.manager import ManagerEnv
 
-    return np, registry, ManagerEnv
+    return np, registry, (DirectEnv, ManagerEnv)
 
 
 def prepare_env(
@@ -32,7 +33,7 @@ def prepare_env(
 ):
     """Create an environment and compatible fixed actions, checking comparisons."""
     np, registry, manager_env_type = load_runtime()
-    kwargs = {"mode": args.mode, "num_envs": num_envs}
+    kwargs = {"mode": args.env_mode, "num_envs": num_envs}
     if args.sim_backend is not None:
         kwargs["sim_backend"] = args.sim_backend
     env = registry.make(name, **kwargs)
@@ -61,8 +62,8 @@ def prepare_env(
             # 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.
+        # Direct and Manager environments both 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