feat(collection): unify episode targets across runners - #727
Conversation
Use collection.target_episodes as the committed data count for ordinary and Expansion runs. Separate batch width, recipe selection, retry attempts, and reset accounting while preserving legacy CLI and configuration aliases.
 |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| stats["prepare_reset_count"] += 1 | ||
| env.reset() |
There was a problem hiding this comment.
env.reset() before every batch, which enables recorder and trajectory saving even though generate_function already commits or discards the preceding batch. Empty dataset buffers are skipped, but this still runs save-capable reset paths during preparation and could persist residual recordings. Use a non-saving prepare reset to keep saving at the intended commit boundary.
Prompt To Fix With AI
This is a comment left during a code review.
Path: embodichain/lab/scripts/run_env.py
Line: 1374-1375
Comment:
**Preparation uses saving reset** The ordinary collector calls `env.reset()` before every batch, which enables recorder and trajectory saving even though `generate_function` already commits or discards the preceding batch. Empty dataset buffers are skipped, but this still runs save-capable reset paths during preparation and could persist residual recordings. Use a non-saving prepare reset to keep saving at the intended commit boundary.
---
For each issue above, determine whether it is valid and should be fixed. If so, fix it directly.
| collection_stats = getattr(self, "_collection_stats", None) | ||
| if isinstance(collection_stats, Mapping): | ||
| metadata["collection"] = dict(collection_stats) |
There was a problem hiding this comment.
collection counters while saving the episode, but the runner increments committed_episodes and batch_count only after the commit returns. As a result, the first saved episode reports zero committed episodes and zero batches, and later saved episodes carry counters from before their own commit. This makes the episode metadata misleading.
Prompt To Fix With AI
This is a comment left during a code review.
Path: embodichain/lab/gym/envs/embodied_env.py
Line: 1364-1366
Comment:
**Episode counters lag commits** The recorder copies `collection` counters while saving the episode, but the runner increments `committed_episodes` and `batch_count` only after the commit returns. As a result, the first saved episode reports zero committed episodes and zero batches, and later saved episodes carry counters from before their own commit. This makes the episode metadata misleading.
---
For each issue above, determine whether it is valid and should be fixed. If so, fix it directly.Document collection terms, reset accounting, explicit recipe selection, and manifest fields. Preserve CLI target precedence and exclude startup and shutdown resets from discard counts.
| if cli_target is not None: | ||
| target = cli_target | ||
| elif target_values: | ||
| if len(set(target_values)) != 1: | ||
| raise ValueError("collection.target_episodes is declared more than once") | ||
| target = target_values[0] |
There was a problem hiding this comment.
--max_episodes skips the check that would reject them. The run proceeds with contradictory configuration, making it harder to detect a mistaken target declaration.
| if cli_target is not None: | |
| target = cli_target | |
| elif target_values: | |
| if len(set(target_values)) != 1: | |
| raise ValueError("collection.target_episodes is declared more than once") | |
| target = target_values[0] | |
| if len(set(target_values)) > 1: | |
| raise ValueError("collection.target_episodes is declared more than once") | |
| if cli_target is not None: | |
| target = cli_target | |
| elif target_values: | |
| target = target_values[0] |
Prompt To Fix With AI
This is a comment left during a code review.
Path: embodichain/lab/scripts/run_env.py
Line: 164-169
Comment:
**Conflicting targets go unchecked** When the task and Expansion declaration specify different episode targets, supplying `--max_episodes` skips the check that would reject them. The run proceeds with contradictory configuration, making it harder to detect a mistaken target declaration.
```suggestion
if len(set(target_values)) > 1:
raise ValueError("collection.target_episodes is declared more than once")
if cli_target is not None:
target = cli_target
elif target_values:
target = target_values[0]
```
---
For each issue above, determine whether it is valid and should be fixed. If so, fix it directly.
| "start_line": 1497, | ||
| "end_line": 1502, | ||
| "excerpt": " cfg: EmbodiedEnvCfg = config_to_cfg(\n gym_config,\n manager_modules=get_manager_modules(),\n source_path=gym_config_source_path,\n task_program_path_override=getattr(args, \"task_program\", None),\n )" |
There was a problem hiding this comment.
The snapshot is pinned to revision 3224ac1, but this excerpt and the other updated anchors describe the PR’s current source. Direct validation compares them with the pinned revision and reports stale evidence. The architecture links also no longer point to the code the snapshot claims to document.
Prompt To Fix With AI
This is a comment left during a code review.
Path: docs/architecture/curated.json
Line: 3865-3867
Comment:
**Evidence no longer matches revision**
The snapshot is pinned to revision `3224ac1`, but this excerpt and the other updated anchors describe the PR’s current source. Direct validation compares them with the pinned revision and reports stale evidence. The architecture links also no longer point to the code the snapshot claims to document.
---
For each issue above, determine whether it is valid and should be fixed. If so, fix it directly.
| func = functor_cfg.func | ||
| if not inspect.isclass(func) or func.__name__ not in { | ||
| "LeRobotRecorder", | ||
| "AsyncLeRobotRecorder", | ||
| }: | ||
| continue |
There was a problem hiding this comment.
robot_meta.robot_type, relying on this method to obtain it from the selected embodiment. But the base manager has already replaced each recorder class with an instance before this check runs, so inspect.isclass(func) is false. The recorder has already captured empty robot metadata, and both datasets use the generic robot type "robot" rather than the selected embodiment’s type.
Prompt To Fix With AI
This is a comment left during a code review.
Path: embodichain/lab/gym/envs/managers/dataset_manager.py
Line: 166-171
Comment:
**Recorder robot type is lost** The repeated-pick-place Default and Newton configs now omit `robot_meta.robot_type`, relying on this method to obtain it from the selected embodiment. But the base manager has already replaced each recorder class with an instance before this check runs, so `inspect.isclass(func)` is false. The recorder has already captured empty robot metadata, and both datasets use the generic robot type `"robot"` rather than the selected embodiment’s type.
---
For each issue above, determine whether it is valid and should be fixed. If so, fix it directly.| full_recipe_indices = _pad_expansion_recipe_indices(selected, num_envs=num_envs) | ||
| batch_recipes = tuple(recipes[index] for index in full_recipe_indices) |
There was a problem hiding this comment.
generate_function() nevertheless requires every physical row to succeed. If a padded row fails while all selected rows succeed, the entire attempt is discarded and retried; exhausting retries rejects valid selected episodes. Filtering records afterward does not affect that success check.
Prompt To Fix With AI
This is a comment left during a code review.
Path: embodichain/lab/scripts/run_env.py
Line: 949-950
Comment:
**Unused rows can reject batches** In a partial final Expansion batch, this padding executes duplicate recipes in rows that will not be saved. `generate_function()` nevertheless requires every physical row to succeed. If a padded row fails while all selected rows succeed, the entire attempt is discarded and retried; exhausting retries rejects valid selected episodes. Filtering records afterward does not affect that success check.
---
For each issue above, determine whether it is valid and should be fixed. If so, fix it directly.
1 participant
Description
Unify ordinary task and Expansion data collection around
collection.target_episodes.This change separates the final committed episode count from vector batch width, logical recipe selection, retry attempts, and reset boundaries. It adds sequential and explicit recipe selection, preserves
--max_episodesand legacy environmentmax_episodescompatibility, and rejects conflicting target declarations.Expansion and ordinary
run-envcollection now use the same episode-count loop. Expansion supports partial final batches and explicit recipe IDs without treating recipe IDs as batch starts. Manifests record target, planned, committed, rejected, attempts, batch count, and prepare/commit/discard reset counts. Expansion recipe indices are threaded per physical row so sparse explicit selections remain deterministic.The task-facing expansion configs now declare collection targets and selection policy; duplicate
candidate_indicesbatch-start configuration is removed.Type of change
Checklist
black .command to format the code base.Validation
429 passedacross run-env, Expansion, Task Program, motion expansion, and bridge/environment tests.python docs/scripts/check_api_docs.py:2387/2387exports documented.black --check --diff --color .passed.git diff --checkpassed.