Add agent graph recovery baseline#287
Open
skywhite1024 wants to merge 19 commits into
Open
Conversation
Contributor
There was a problem hiding this comment.
Pull request overview
This PR introduces a graph-based “Task → Monitor → Recovery” execution baseline for generative simulation agents, replacing the legacy code-generation workflow with JSON graph artifacts compiled and executed via a runtime task graph. It also updates atomic actions to support ordered execution and adds environment/config changes to support recovery, success predicates, and environment-based LLM configuration.
Changes:
- Added agent task-graph + monitor/recovery scaffolding (monitor functions, error functions, graph runtime) and updated scripts to support recovery and public atomic-action options.
- Refactored atomic-action execution to preserve configured action order (even with duplicate action names) and added support for MoveAction joint-position targets.
- Switched LLM configuration to environment variables /
.env+ added new task configs and documentation for the new workflow.
Reviewed changes
Copilot reviewed 64 out of 65 changed files in this pull request and generated 5 comments.
Show a summary per file
| File | Description |
|---|---|
| tests/sim/atomic_actions/test_engine.py | Adds regression test ensuring execute_static uses configured action order even with duplicate names. |
| tests/sim/atomic_actions/test_actions.py | Adds tests for MoveAction qpos-target execution and waypoint validation. |
| embodichain/utils/llm_json.py | New helper to extract/normalize JSON from LLM outputs (plain or fenced). |
| embodichain/lab/sim/atomic_actions/engine.py | Preserves action registration order via _action_sequence for execute_static. |
| embodichain/lab/sim/atomic_actions/actions.py | Adds qpos-target support to MoveAction, plus qpos interpolation helpers and stricter tensor normalization. |
| embodichain/lab/sim/atom_actions.py | Removes legacy monolithic atom-action implementation (migrated into agent-oriented structure). |
| embodichain/lab/sim/agent/monitor_utils.py | Adds helpers to compute object/arm/gripper state for monitoring. |
| embodichain/lab/sim/agent/monitor_functions.py | Adds monitor functions (object moved, object held, gripper distance). |
| embodichain/lab/sim/agent/error_functions.py | Adds error/disturbance helpers and interactive error-injection utilities for recovery experiments. |
| embodichain/lab/sim/agent/atom_action_utils.py | Adds agent-specific atomic-action utilities (IK, trajectory planning, state sync, drive-call extraction). |
| embodichain/lab/sim/agent/agent_graph.py | Adds task-graph runtime that executes edges via drive() and switches to recovery branches on monitor triggers. |
| embodichain/lab/scripts/run_env.py | Adds post-execution task-success logging and plumbing for recovery/public-action flags. |
| embodichain/lab/scripts/run_agent.py | Adds CLI flags for recovery + public atomic-action toggles; fixes import path for main. |
| embodichain/lab/gym/utils/misc.py | Adds apply_rotation helper used by failure injection. |
| embodichain/lab/gym/envs/tasks/tableware/single_arm_agent_env.py | New single-arm agent env mixin exposed via right-arm slot. |
| embodichain/lab/gym/envs/tasks/tableware/configurable_success.py | New config-driven task-success predicate evaluator. |
| embodichain/lab/gym/envs/tasks/tableware/base_agent_env.py | Reworks agent initialization to Task/Recovery/Compile agents; adds arm-slot abstraction and graph generation pipeline. |
| embodichain/lab/gym/envs/tasks/tableware/atomic_actions.py | New AtomicActionsAgentEnv environment using configurable success and grasp pose overrides. |
| embodichain/lab/gym/envs/tasks/init.py | Registers and exports AtomicActionsAgentEnv. |
| embodichain/data/dataset.py | Adds repo-root relative resolution as an additional dataset path fallback. |
| embodichain/data/init.py | Changes generated agent prompt dir location (see review comment about writability/config). |
| embodichain/agents/prompts/recovery_rules.txt | Adds RecoveryAgent authoring constraints/spec rules. |
| embodichain/agents/prompts/monitor_functions.txt | Adds monitor-function prompt documentation. |
| embodichain/agents/prompts/error_functions.txt | Adds failure-description prompt context for recovery spec generation. |
| embodichain/agents/prompts/code_prompt.txt | Removes legacy code-generation prompt rules. |
| embodichain/agents/prompts/code_example.txt | Removes legacy code-generation examples. |
| embodichain/agents/prompts/basic_background.txt | Updates background prompt to remove camera-mirroring guidance and standardize axis text. |
| embodichain/agents/prompts/atom_actions.txt | Removes drive description from atomic-actions prompt context. |
| embodichain/agents/mllm/prompt/task_prompt.py | Updates TaskAgent prompting to emit a nominal JSON task graph (no prose/code). |
| embodichain/agents/mllm/prompt/recovery_prompt.py | Adds prompt for generating recovery bindings spec JSON. |
| embodichain/agents/mllm/prompt/compile_prompt.py | Adds compatibility namespace for compile prompt imports (LLM-free compilation). |
| embodichain/agents/mllm/prompt/code_prompt.py | Removes legacy code prompt class. |
| embodichain/agents/mllm/prompt/init.py | Exports new prompt modules and removes CodePrompt export. |
| embodichain/agents/hierarchy/validation_agent.py | Removes legacy image-based validation agent implementation. |
| embodichain/agents/hierarchy/task_agent.py | Updates TaskAgent to generate/store normalized JSON graph output. |
| embodichain/agents/hierarchy/recovery_agent.py | Adds RecoveryAgent to generate/store normalized recovery spec JSON. |
| embodichain/agents/hierarchy/llm.py | Switches from AzureChatOpenAI to ChatOpenAI + .env loading and env-based configuration. |
| embodichain/agents/hierarchy/compile_agent.py | Adds CompileAgent that normalizes/expands recovery spec and executes compiled graph artifacts. |
| embodichain/agents/hierarchy/code_agent.py | Removes legacy code-generation and AST execution agent. |
| embodichain/agents/hierarchy/agent_base.py | Adjusts composed observation assembly (no longer auto-pulls env obs by default). |
| docs/source/features/generative_sim/index.rst | Adds EmbodiAgent + failure-monitor-recovery docs links. |
| docs/source/features/generative_sim/agents.md | Updates EmbodiAgent docs for graph/recovery workflow (needs updates re removed ValidationAgent). |
| docs/source/features/failure_monitor_recovery.md | Adds documentation for monitor/recovery graph semantics and artifacts. |
| configs/gym/pour_water/gym_config.json | Adjusts env config (removes explicit control_parts list under env). |
| configs/gym/pour_water/gym_config_simple.json | Updates joint normalization ids and removes env-level control_parts list. |
| configs/gym/agent/ur10_cup_transfer_agent/task_prompt.txt | New single-arm UR10 cup transfer task prompt. |
| configs/gym/agent/ur10_cup_transfer_agent/fast_gym_config.json | New AtomicActionsAgent-v3 config for UR10 cup transfer with success spec. |
| configs/gym/agent/ur10_cup_transfer_agent/basic_background.txt | New basic background prompt for UR10 cup transfer. |
| configs/gym/agent/ur10_cup_transfer_agent/atom_actions.txt | New task-scoped atomic-action prompt for UR10 cup transfer. |
| configs/gym/agent/ur10_cup_transfer_agent/agent_config.json | New agent config wiring Task/Recovery/Compile agents for UR10 cup transfer. |
| configs/gym/agent/rearrangement_agent/fast_gym_config.json | Migrates rearrangement agent config to AtomicActionsAgent-v3 + adds success spec and camera/sensor updates. |
| configs/gym/agent/rearrangement_agent/atom_actions.txt | New task-scoped atomic-action prompt for rearrangement. |
| configs/gym/agent/rearrangement_agent/agent_config.json | Migrates to Task/Recovery/Compile agents and replaces code prompt/example with recovery prompt inputs. |
| configs/gym/agent/pour_water_agent/fast_gym_config.json | Updates pour-water agent config (episodes/steps, camera setup, observation normalization, dataset metadata). |
| configs/gym/agent/pour_water_agent/agent_config.json | Migrates to Task/Recovery/Compile agents and recovery prompt inputs. |
| configs/gym/agent/pour_water_agent/agent_config_dual.json | Migrates dual pour-water agent config to Task/Recovery/Compile agents. |
| configs/gym/agent/apple_pick_agent/task_prompt.txt | New single-arm apple pick-and-place task prompt. |
| configs/gym/agent/apple_pick_agent/fast_gym_config.json | New AtomicActionsAgent-v3 config for apple pick task with success spec and grasp overrides. |
| configs/gym/agent/apple_pick_agent/basic_background.txt | New basic background prompt for apple pick task. |
| configs/gym/agent/apple_pick_agent/atom_actions.txt | New task-scoped atomic-action prompt for apple pick task. |
| configs/gym/agent/apple_pick_agent/agent_config.json | New agent config wiring Task/Recovery/Compile agents for apple pick task. |
| .gitignore | Ignores generated agent artifacts and additional local/generated assets. |
| .env.example | Adds template for local LLM configuration via environment variables. |
Comments suppressed due to low confidence (1)
docs/source/features/generative_sim/agents.md:95
- The
ValidationAgentsection refers to a module that no longer exists. Either remove this section or rewrite it to describe the current success evaluation entrypoint (env.is_task_success()/ configurable_success).
### ValidationAgent
*Located in:* `embodichain/agents/hierarchy/validation_agent.py`
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
+21
to
+25
| database_dir = EMBODICHAIN_DEFAULT_DATABASE_ROOT | ||
| database_2d_dir = os.path.join(database_dir, "2dasset") | ||
| database_agent_prompt_dir = os.path.join(database_dir, "agent_prompt") | ||
| database_agent_prompt_dir = os.path.abspath( | ||
| os.path.join(os.path.dirname(__file__), "..", "database/agent_generated_content") | ||
| ) |
Comment on lines
+62
to
+69
| self.open_state = torch.as_tensor( | ||
| getattr(self, "gripper_open_state", [0.05]), | ||
| dtype=self.init_qpos.dtype, | ||
| ).flatten() | ||
| self.close_state = torch.as_tensor( | ||
| getattr(self, "gripper_close_state", [0.0]), | ||
| dtype=self.init_qpos.dtype, | ||
| ).flatten() |
Comment on lines
+54
to
+56
| - **Execute**: The graph runtime executes atomic actions and switches to recovery branches when monitors trigger. | ||
| - **Validate**: The `ValidationAgent` analyzes the result images, selects the best camera angle, and judges success. | ||
| - **Refine**: If validation fails, feedback is sent back to the agents to regenerate the plan or code. | ||
| - **Refine**: If validation fails, feedback is sent back to the agents to regenerate the graph artifacts. |
Comment on lines
+96
to
+99
| Closes the loop by verifying if the robot actually achieved what it planned. | ||
|
|
||
| * Uses a specialized LLM call (`select_best_view_dir`) to analyze images from all cameras and pick the single best angle that proves the action's outcome, ignoring irrelevant views. | ||
| * If an error occurs (runtime or logic), it generates a detailed explanation which is fed back to the `TaskAgent` or `CodeAgent` for the next attempt. | ||
| * If an error occurs (runtime or logic), it generates a detailed explanation which is fed back to the `TaskAgent` or `CompileAgent` for the next attempt. |
Comment on lines
+129
to
131
| │ ├── compile_agent.py # Graph compilation & execution interface | ||
| │ ├── validation_agent.py # Visual analysis & view selection | ||
| │ └── llm.py # LLM configuration and instances |
yuecideng
requested changes
Jun 3, 2026
| self._semantic_analyzer = SemanticAnalyzer() | ||
|
|
||
| # Initialize default actions | ||
| self._action_sequence: List[tuple[str, AtomicAction]] = [] |
Contributor
There was a problem hiding this comment.
Why adding a new sequence of AtomicAction?
| @@ -0,0 +1,9 @@ | |||
| # Copy this file to .env and fill in your local credentials. | |||
Contributor
There was a problem hiding this comment.
We should unified MLLM entry point @XuanchaoPENG
| scripts/benchmark/rl/reports/* | ||
|
|
||
| # Local/generated assets | ||
| obj/asset_simready.obj |
Contributor
There was a problem hiding this comment.
We should add general gitignore contents. Specific items should not be added to this file
| arm_dof = self.dof if arm_dof is None else arm_dof | ||
| if start_qpos is None: | ||
| start_qpos = self.robot.get_qpos(name=self.cfg.control_part) | ||
| start_qpos = torch.as_tensor( |
Contributor
There was a problem hiding this comment.
This line should be removed. We only accept torch tensor
| trajectory of shape (n_envs, n_waypoints, dof), | ||
| joint_ids corresponding to trajectory | ||
| """ | ||
| if isinstance(target, torch.Tensor): |
Contributor
There was a problem hiding this comment.
We have Joint space intepolation method for motion generator. Please use the unified module.
skywhite1024
force-pushed
the
ljd/baseline_agentchoard_clean
branch
from
2a16b73 to
8602d5e
Compare
Contributor
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 85 out of 89 changed files in this pull request and generated 2 comments.
Comments suppressed due to low confidence (1)
embodichain/gen_sim/action_agent_pipeline/error_functions.py:320
wrong_affordanceduplicates the (single-column)gripper_state_trajtwice in the returned action. After preserving both gripper DOFs (see comment above), the return should append the fullgripper_state_trajonce to keep the original gripper joint values intact.
Comment on lines
+50
to
+53
| if last_frame_pose is None: | ||
| last_frame_pose = env.obj_info.get(obj_name).get("pose") | ||
| current_pose = _get_object_pose(env, obj_name) | ||
| previous_pose = _as_pose_matrix(last_frame_pose, device=current_pose.device) |
Comment on lines
+272
to
+277
| else: | ||
| select_arm = "left_arm" if is_left else "right_arm" | ||
| start_qpos = torch.as_tensor(action[0, :-2], dtype=torch.float32) | ||
| last_qpos = torch.as_tensor(action[-1, :-2], dtype=torch.float32) | ||
| gripper_state_traj = action[:, -1:].copy() | ||
| last_pose = env.get_arm_fk(qpos=last_qpos, is_left=is_left).clone() |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
This PR adds the baseline agent graph and recovery workflow for generative simulation agents.
Main changes:
embodichain/lab/sim/agent/.run_agent.py/run_env.pyto support the new public atomic action and recovery-related options..env.exampleand make LLM configuration load from environment variables instead of hardcoded credentials.No issue is currently linked.
Fixes #
Type of change
Screenshots
Not applicable.
Checklist
black .command to format the code base.Additional validation performed:
git diff --check origin/main...HEADpython -m black --check .python -m pytest tests/sim/atomic_actions/test_actions.py tests/sim/atomic_actions/test_engine.py -q.env, generated agent content, and temporary export artifacts are not included in the PR diff.