unisim.backend.mujoco.chunk_tuner

Adaptive thread-pool chunk_size selection for the MuJoCo BatchEnvPool.

All probing/benchmarking happens on the cold path (materialize()), never on step/reset. Results are cached on disk keyed by model signature + device fingerprint + (num_envs, nthread, dtype, …) so repeated trainings reuse them.

Functions

benchmark_chunk_sizes(pool, state, nstep, ...)

Min wall-clock of a representative pool.step per candidate.

cache_path()

device_fingerprint()

Coarse, deterministic local-device descriptor for the cache key.

file_lock(lock_path)

Best-effort inter-process exclusive lock; no-op where fcntl is absent.

filter_candidates(candidates, num_envs, *[, ...])

Clamp to [1, num_envs], dedup, sort.

load_cache(path)

make_cache_key(*, backend_type, model_sig, ...)

make_candidates(num_envs, nthread)

Candidate chunk_sizes within [1, upper], densified at the sweet spot.

model_signature(model, n_variants)

Structural determinants of per-step cost (proxy for 'task').

resolve_chunk_size(*, pool, state, model, ...)

select_chunk_size(timings, base, *[, ...])

store_cache(path, key, value)

unisim.backend.mujoco.chunk_tuner.make_candidates(num_envs, nthread)[source]

Candidate chunk_sizes within [1, upper], densified at the sweet spot.

upper = ceil(num_envs / nthread) is a hard ceiling: a chunk_size beyond it produces fewer chunks than threads, leaving threads idle -> always slower. The optimum sits near heur = num_envs / (10 * nthread) (~10 chunks per thread), so the band around it is sampled densely while the always-bad large chunks are never generated.

Parameters:

  • num_envs (int)

  • nthread (int)

Return type:

list[int]

unisim.backend.mujoco.chunk_tuner.filter_candidates(candidates, num_envs, *, max_candidates=16)[source]

Clamp to [1, num_envs], dedup, sort. If over the cap, trim the coarse middle while keeping the low/optimum band (smallest values) plus the coarsest anchor – never log-spaced (which used to drop the optimum).

Parameters:
Return type:

list[int]

unisim.backend.mujoco.chunk_tuner.device_fingerprint()[source]

Coarse, deterministic local-device descriptor for the cache key.

Return type:

dict[str, Any]

unisim.backend.mujoco.chunk_tuner.model_signature(model, n_variants)[source]

Structural determinants of per-step cost (proxy for ‘task’).

Parameters:

  • model (Any)

  • n_variants (int)

Return type:

dict[str, int]

unisim.backend.mujoco.chunk_tuner.make_cache_key(*, backend_type, model_sig, device, num_envs, nthread, dtype, post_step_forward_sensor, bench_nsteps)[source]
Parameters:

  • backend_type (str)

  • model_sig (dict)

  • device (dict)

  • num_envs (int)

  • nthread (int)

  • dtype (Any)

  • post_step_forward_sensor (bool)

  • bench_nsteps (int)

Return type:

str

unisim.backend.mujoco.chunk_tuner.cache_path()[source]
Return type:

Path

unisim.backend.mujoco.chunk_tuner.load_cache(path)[source]
Parameters:

path (Path)

Return type:

dict

unisim.backend.mujoco.chunk_tuner.store_cache(path, key, value)[source]
Parameters:
Return type:

None

unisim.backend.mujoco.chunk_tuner.file_lock(lock_path)[source]

Best-effort inter-process exclusive lock; no-op where fcntl is absent.

Parameters:

lock_path (Path)

unisim.backend.mujoco.chunk_tuner.benchmark_chunk_sizes(pool, state, nstep, candidates, *, control, post_step_forward_sensor, warmup=2, reps=20, time_budget_s=25.0)[source]

Min wall-clock of a representative pool.step per candidate.

Each candidate is sampled reps times and reduced via min (uncontended cost). None (native default) is always measured as the baseline anchor.

Return type:

dict[int | None, float]

unisim.backend.mujoco.chunk_tuner.select_chunk_size(timings, base, *, margin=0.03, tie_tol=0.02)[source]
Parameters:
Return type:

int | None

unisim.backend.mujoco.chunk_tuner.resolve_chunk_size(*, pool, state, model, n_variants, num_envs, nthread, dtype, post_step_forward_sensor, bench_nsteps, manual_chunk_size, adaptive, backend_type='mujoco', model_file=None)[source]
Parameters:

  • n_variants (int)

  • num_envs (int)

  • nthread (int)

  • post_step_forward_sensor (bool)

  • bench_nsteps (int)

  • manual_chunk_size (int | None)

  • adaptive (bool)

  • backend_type (str)

  • model_file (str | None)

Return type:

int | None