mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-21 17:01:53 +00:00
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
This commit is contained in:
@@ -65,21 +65,27 @@ class WandBConfig:
|
|||||||
class EvalConfig:
|
class EvalConfig:
|
||||||
n_episodes: int = 50
|
n_episodes: int = 50
|
||||||
# `batch_size` specifies the number of environments to use in a gym.vector.VectorEnv.
|
# `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).
|
# `use_async_envs` specifies whether to use asynchronous environments (multiprocessing).
|
||||||
# Defaults to True; automatically downgraded to SyncVectorEnv when batch_size=1.
|
# Defaults to True; automatically downgraded to SyncVectorEnv when batch_size=1.
|
||||||
use_async_envs: bool = True
|
use_async_envs: bool = True
|
||||||
|
|
||||||
def __post_init__(self) -> None:
|
def __post_init__(self) -> None:
|
||||||
|
if self.batch_size == 0:
|
||||||
|
self.batch_size = self._auto_batch_size()
|
||||||
if self.batch_size > self.n_episodes:
|
if self.batch_size > self.n_episodes:
|
||||||
raise ValueError(
|
self.batch_size = self.n_episodes
|
||||||
"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} "
|
def _auto_batch_size(self) -> int:
|
||||||
f"eval environments will be instantiated, but only {self.n_episodes} will be used. "
|
"""Pick batch_size based on CPU cores, capped by n_episodes."""
|
||||||
"This might significantly slow down evaluation. To fix this, you should update your command "
|
import math
|
||||||
f"to increase the number of episodes to match the batch size (e.g. `eval.n_episodes={self.batch_size}`), "
|
import os
|
||||||
f"or lower the batch size (e.g. `eval.batch_size={self.n_episodes}`)."
|
|
||||||
)
|
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
|
@dataclass
|
||||||
|
|||||||
@@ -413,17 +413,25 @@ class _LazyAsyncVectorEnv:
|
|||||||
AsyncVectorEnv on first reset(), keeping peak process count = n_envs.
|
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_fns = env_fns
|
||||||
self._env: gym.vector.AsyncVectorEnv | None = None
|
self._env: gym.vector.AsyncVectorEnv | None = None
|
||||||
self.num_envs = len(env_fns)
|
self.num_envs = len(env_fns)
|
||||||
# Instantiate one env to expose spaces (no GPU — _ensure_env is lazy).
|
if observation_space is not None and action_space is not None:
|
||||||
tmp = env_fns[0]()
|
self.observation_space = observation_space
|
||||||
self.observation_space = tmp.observation_space
|
self.action_space = action_space
|
||||||
self.action_space = tmp.action_space
|
else:
|
||||||
self.single_observation_space = tmp.observation_space
|
tmp = env_fns[0]()
|
||||||
self.single_action_space = tmp.action_space
|
self.observation_space = tmp.observation_space
|
||||||
tmp.close()
|
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):
|
def _ensure(self):
|
||||||
if self._env is None:
|
if self._env is None:
|
||||||
@@ -507,6 +515,11 @@ def create_libero_envs(
|
|||||||
if not selected:
|
if not selected:
|
||||||
raise ValueError(f"No tasks selected for suite '{suite_name}' (available: {total}).")
|
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:
|
for tid in selected:
|
||||||
fns = _make_env_fns(
|
fns = _make_env_fns(
|
||||||
suite=suite,
|
suite=suite,
|
||||||
@@ -521,7 +534,11 @@ def create_libero_envs(
|
|||||||
camera_name_mapping=camera_name_mapping,
|
camera_name_mapping=camera_name_mapping,
|
||||||
)
|
)
|
||||||
if is_async:
|
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:
|
else:
|
||||||
out[suite_name][tid] = env_cls(fns)
|
out[suite_name][tid] = env_cls(fns)
|
||||||
print(f"Built vec env | suite={suite_name} | task_id={tid} | n_envs={n_envs}")
|
print(f"Built vec env | suite={suite_name} | task_id={tid} | n_envs={n_envs}")
|
||||||
|
|||||||
@@ -526,7 +526,7 @@ def eval_main(cfg: EvalPipelineConfig):
|
|||||||
|
|
||||||
logging.info(colored("Output dir:", "yellow", attrs=["bold"]) + f" {cfg.output_dir}")
|
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(
|
envs = make_env(
|
||||||
cfg.env,
|
cfg.env,
|
||||||
n_envs=cfg.eval.batch_size,
|
n_envs=cfg.eval.batch_size,
|
||||||
|
|||||||
Reference in New Issue
Block a user