mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-25 10:46:01 +00:00
fix(eval): address review comments
- Wrap rollout loop in try/finally so finalize() runs on crash/interrupt - Guard push_to_hub with num_episodes > 0 to avoid pushing empty datasets - Hoist loop-invariant multi_env and base_repo_id out of creation loop
This commit is contained in:
@@ -216,10 +216,10 @@ def rollout(
|
||||
features = _env_features_to_dataset_features(env_features)
|
||||
fps = env.unwrapped.metadata.get("render_fps", 30)
|
||||
recording_datasets = []
|
||||
for i in range(env.num_envs):
|
||||
multi_env = env.num_envs > 1
|
||||
root = str(recording_dir / f"env_{i}") if multi_env else str(recording_dir)
|
||||
base_repo_id = recording_repo_id or "eval_recording"
|
||||
for i in range(env.num_envs):
|
||||
root = str(recording_dir / f"env_{i}") if multi_env else str(recording_dir)
|
||||
repo_id = f"{base_repo_id}_env_{i}" if multi_env else base_repo_id
|
||||
recording_datasets.append(
|
||||
LeRobotDataset.create(
|
||||
@@ -253,6 +253,7 @@ def rollout(
|
||||
leave=False,
|
||||
)
|
||||
check_env_attributes_and_types(env)
|
||||
try:
|
||||
while not np.all(done) and step < max_steps:
|
||||
# Numpy array to tensor and changing dictionary keys to LeRobot policy format.
|
||||
observation = preprocess_observation(observation)
|
||||
@@ -303,7 +304,9 @@ def rollout(
|
||||
elif "is_success" in info:
|
||||
is_success = info["is_success"]
|
||||
successes = (
|
||||
is_success.tolist() if hasattr(is_success, "tolist") else [bool(is_success)] * env.num_envs
|
||||
is_success.tolist()
|
||||
if hasattr(is_success, "tolist")
|
||||
else [bool(is_success)] * env.num_envs
|
||||
)
|
||||
else:
|
||||
successes = [False] * env.num_envs
|
||||
@@ -347,6 +350,15 @@ def rollout(
|
||||
)
|
||||
progbar.set_postfix({"running_success_rate": f"{running_success_rate.item() * 100:.1f}%"})
|
||||
progbar.update()
|
||||
finally:
|
||||
if recording_datasets is not None:
|
||||
for ds in recording_datasets:
|
||||
ds.finalize()
|
||||
if recording_repo_id is not None:
|
||||
if ds.num_episodes > 0:
|
||||
ds.push_to_hub(private=recording_private)
|
||||
else:
|
||||
logging.warning("No episodes recorded for %s — skipping push to hub.", ds.repo_id)
|
||||
|
||||
# Track the final observation.
|
||||
if return_observations:
|
||||
@@ -366,12 +378,6 @@ def rollout(
|
||||
stacked_observations[key] = torch.stack([obs[key] for obs in all_observations], dim=1)
|
||||
ret[OBS_STR] = stacked_observations
|
||||
|
||||
if recording_datasets is not None:
|
||||
for ds in recording_datasets:
|
||||
ds.finalize()
|
||||
if recording_repo_id is not None:
|
||||
ds.push_to_hub(private=recording_private)
|
||||
|
||||
if hasattr(policy, "use_original_modules"):
|
||||
policy.use_original_modules()
|
||||
|
||||
|
||||
Reference in New Issue
Block a user