From ce6c0ba1b7c01714fe121d26cf6621e2ea3aeee9 Mon Sep 17 00:00:00 2001 From: Pepijn Kooijmans Date: Tue, 7 Apr 2026 13:55:27 +0200 Subject: [PATCH] feat(eval): batch_size=auto + faster env loading - batch_size=0 (default) auto-tunes based on CPU cores, capped by n_episodes and 64. Removes the need for users to guess the right value. The old batch_size > n_episodes error is replaced by silently clamping to n_episodes. - _LazyAsyncVectorEnv accepts pre-computed spaces so only one temp env is created per suite (not per task). For libero_spatial (10 tasks) this avoids 9 redundant LiberoEnv instantiations during env setup. Made-with: Cursor --- src/lerobot/configs/default.py | 24 ++++++++++++-------- src/lerobot/envs/libero.py | 35 +++++++++++++++++++++-------- src/lerobot/scripts/lerobot_eval.py | 2 +- 3 files changed, 42 insertions(+), 19 deletions(-) diff --git a/src/lerobot/configs/default.py b/src/lerobot/configs/default.py index 38039a7bf..d6ad665bf 100644 --- a/src/lerobot/configs/default.py +++ b/src/lerobot/configs/default.py @@ -65,21 +65,27 @@ class WandBConfig: class EvalConfig: n_episodes: int = 50 # `batch_size` specifies the number of environments to use in a gym.vector.VectorEnv. - batch_size: int = 50 + # Set to 0 for auto-tuning based on available CPU cores and n_episodes. + batch_size: int = 0 # `use_async_envs` specifies whether to use asynchronous environments (multiprocessing). # Defaults to True; automatically downgraded to SyncVectorEnv when batch_size=1. use_async_envs: bool = True def __post_init__(self) -> None: + if self.batch_size == 0: + self.batch_size = self._auto_batch_size() if self.batch_size > self.n_episodes: - raise ValueError( - "The eval batch size is greater than the number of eval episodes " - f"({self.batch_size} > {self.n_episodes}). As a result, {self.batch_size} " - f"eval environments will be instantiated, but only {self.n_episodes} will be used. " - "This might significantly slow down evaluation. To fix this, you should update your command " - f"to increase the number of episodes to match the batch size (e.g. `eval.n_episodes={self.batch_size}`), " - f"or lower the batch size (e.g. `eval.batch_size={self.n_episodes}`)." - ) + self.batch_size = self.n_episodes + + def _auto_batch_size(self) -> int: + """Pick batch_size based on CPU cores, capped by n_episodes.""" + import math + import os + + cpu_cores = os.cpu_count() or 4 + # Each async env worker needs ~1 core; leave headroom for main process + inference. + by_cpu = max(1, math.floor(cpu_cores * 0.7)) + return min(by_cpu, self.n_episodes, 64) @dataclass diff --git a/src/lerobot/envs/libero.py b/src/lerobot/envs/libero.py index e731d46ed..3cb28d5b5 100644 --- a/src/lerobot/envs/libero.py +++ b/src/lerobot/envs/libero.py @@ -413,17 +413,25 @@ class _LazyAsyncVectorEnv: AsyncVectorEnv on first reset(), keeping peak process count = n_envs. """ - def __init__(self, env_fns: list[Callable]): + def __init__( + self, + env_fns: list[Callable], + observation_space: spaces.Space | None = None, + action_space: spaces.Space | None = None, + ): self._env_fns = env_fns self._env: gym.vector.AsyncVectorEnv | None = None self.num_envs = len(env_fns) - # Instantiate one env to expose spaces (no GPU — _ensure_env is lazy). - tmp = env_fns[0]() - self.observation_space = tmp.observation_space - self.action_space = tmp.action_space - self.single_observation_space = tmp.observation_space - self.single_action_space = tmp.action_space - tmp.close() + if observation_space is not None and action_space is not None: + self.observation_space = observation_space + self.action_space = action_space + else: + tmp = env_fns[0]() + self.observation_space = tmp.observation_space + self.action_space = tmp.action_space + tmp.close() + self.single_observation_space = self.observation_space + self.single_action_space = self.action_space def _ensure(self): if self._env is None: @@ -507,6 +515,11 @@ def create_libero_envs( if not selected: raise ValueError(f"No tasks selected for suite '{suite_name}' (available: {total}).") + # All tasks in a suite share identical observation/action spaces. + # Probe once and reuse to avoid creating a temp env per task. + cached_obs_space: spaces.Space | None = None + cached_act_space: spaces.Space | None = None + for tid in selected: fns = _make_env_fns( suite=suite, @@ -521,7 +534,11 @@ def create_libero_envs( camera_name_mapping=camera_name_mapping, ) if is_async: - out[suite_name][tid] = _LazyAsyncVectorEnv(fns) + lazy = _LazyAsyncVectorEnv(fns, cached_obs_space, cached_act_space) + if cached_obs_space is None: + cached_obs_space = lazy.observation_space + cached_act_space = lazy.action_space + out[suite_name][tid] = lazy else: out[suite_name][tid] = env_cls(fns) print(f"Built vec env | suite={suite_name} | task_id={tid} | n_envs={n_envs}") diff --git a/src/lerobot/scripts/lerobot_eval.py b/src/lerobot/scripts/lerobot_eval.py index 9739a77fd..04484bf3a 100644 --- a/src/lerobot/scripts/lerobot_eval.py +++ b/src/lerobot/scripts/lerobot_eval.py @@ -526,7 +526,7 @@ def eval_main(cfg: EvalPipelineConfig): logging.info(colored("Output dir:", "yellow", attrs=["bold"]) + f" {cfg.output_dir}") - logging.info("Making environment.") + logging.info(f"Making environment (batch_size={cfg.eval.batch_size}, async={cfg.eval.use_async_envs}).") envs = make_env( cfg.env, n_envs=cfg.eval.batch_size,