Migrate EmbodiAgent #82

yangchen73 · 2026-01-21T03:14:26Z

Description

This PR migrates EmbodiAgent, a LLM-based agent system for autonomous robot task execution in EmbodiChain. The system enables robots to perceive, plan, code, execute, and validate complex manipulation tasks through a closed-loop control cycle, which is developed by https://github.com/Jasonxu1225 .

Key Features

Three specialized agents working in coordination:

TaskAgent: High-level reasoning and task decomposition from visual observations
CodeAgent: Translates natural language plans into executable Python code using atomic actions
ValidationAgent: Validates task execution and provides feedback for refinement

Key Related Files

Core agent implementation: embodichain/agents/hierarchy/
Agent environments: embodichain/lab/gym/envs/tasks/tableware/
Toolkit interfaces: embodichain/toolkits/interfaces.py
Motion generation: embodichain/lab/gym/motion_generation/
Documentation: embodichain/agents/README.md

Type of change

New feature (non-breaking change which adds functionality)

Screenshots

Pour Water

pour_waterr.mp4

Rearrangement

rearrangement.mp4

Dual Arm Pour Water

pour_water_dual.mp4

Checklist

I have run the black . command to format the code base.
I have made corresponding changes to the documentation
I have added tests that prove my fix is effective or that my feature works
Dependencies have been updated, if applicable.

Copilot

Pull request overview

Copilot reviewed 38 out of 39 changed files in this pull request and generated 51 comments.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Copilot · 2026-01-22T08:41:51Z

embodichain/lab/gym/envs/tasks/tableware/base_agent_env.py

+        if agent_config.get("TaskAgent") is not None:
+            self.task_agent = TaskAgent(
+                task_llm,
+                **agent_config["Agent"],
+                **agent_config["TaskAgent"],
+                task_name=task_name,
+            )
+        self.code_agent = CodeAgent(
+            code_llm,
+            **agent_config["Agent"],
+            **agent_config.get("CodeAgent"),
+            task_name=task_name,
+        )
+        self.validation_agent = ValidationAgent(
+            validation_llm,
+            task_name=task_name,
+            task_description=self.code_agent.prompt_kwargs.get("task_prompt")[
+                "content"
+            ],
+            basic_background=self.code_agent.prompt_kwargs.get("basic_background")[
+                "content"
+            ],
+            atom_actions=self.code_agent.prompt_kwargs.get("atom_actions")["content"],
+        )


The agent initialization in _init_agents is missing error handling when agent_config does not contain expected keys. If agent_config doesn't have 'Agent', 'TaskAgent', or 'CodeAgent' keys, this will raise a KeyError. Consider adding validation or default values to handle missing configuration gracefully.

Copilot · 2026-01-22T08:41:51Z

embodichain/agents/hierarchy/code_agent.py

+            class InjectKwargs(ast.NodeTransformer):
+                def visit_Call(self, node):
+                    self.generic_visit(node)
+                    # Inject **kwargs if not present
+                    has_kwargs = any(
+                        kw.arg is None
+                        and isinstance(kw.value, ast.Name)
+                        and kw.value.id == "kwargs"
+                        for kw in node.keywords
+                    )
+                    if not has_kwargs:
+                        node.keywords.append(
+                            ast.keyword(
+                                arg=None, value=ast.Name(id="kwargs", ctx=ast.Load())
+                            )
+                        )
+                    return node
+
+            # Transform AST to inject kwargs
+            tree = InjectKwargs().visit(tree)
+            ast.fix_missing_locations(tree)


The AST injection of kwargs into function calls (InjectKwargs transformer) modifies all function calls globally within the code, not just drive() calls. This could have unintended side effects if the generated code calls other functions that don't expect **kwargs. Consider making this more targeted to only inject kwargs into specific function calls like drive() and atomic action functions.

Copilot · 2026-01-22T08:41:52Z

embodichain/toolkits/interfaces.py

+def extract_drive_calls(code_str: str) -> list[str]:
+    tree = ast.parse(code_str)
+    lines = code_str.splitlines()
+
+    drive_blocks = []
+
+    for node in tree.body:
+        # Match: drive(...)
+        if (
+            isinstance(node, ast.Expr)
+            and isinstance(node.value, ast.Call)
+            and isinstance(node.value.func, ast.Name)
+            and node.value.func.id == "drive"
+        ):
+            # AST line numbers are 1-based
+            start = node.lineno - 1
+            end = node.end_lineno
+            block = "\n".join(lines[start:end])
+            drive_blocks.append(block)
+
+    return drive_blocks


The function extract_drive_calls uses ast.parse on user-provided code_str without any validation or error handling. If the code string is malformed or contains syntax errors, this will raise an exception. Add try-except block to handle parsing errors gracefully.

Copilot · 2026-01-22T08:41:52Z

embodichain/toolkits/interfaces.py

+def drive(
+    left_arm_action=None,
+    right_arm_action=None,
+    env=None,
+    **kwargs,
+):
+
+    if left_arm_action is not None and right_arm_action is not None:
+        len_left = len(left_arm_action)
+        len_right = len(right_arm_action)
+
+        if len_left < len_right:
+            diff = len_right - len_left
+            padding = np.repeat(left_arm_action[-1:], diff, axis=0)
+            left_arm_action = np.concatenate([left_arm_action, padding], axis=0)
+        elif len_right < len_left:
+            diff = len_left - len_right
+            padding = np.repeat(right_arm_action[-1:], diff, axis=0)
+            right_arm_action = np.concatenate([right_arm_action, padding], axis=0)
+
+        left_arm_index = env.left_arm_joints + env.left_eef_joints
+        right_arm_index = env.right_arm_joints + env.right_eef_joints
+        actions = np.zeros((len(right_arm_action), len(env.init_qpos)))
+        actions[:, left_arm_index] = left_arm_action
+        actions[:, right_arm_index] = right_arm_action
+
+    elif left_arm_action is None and right_arm_action is not None:
+        left_arm_index = env.left_arm_joints + env.left_eef_joints
+        right_arm_index = env.right_arm_joints + env.right_eef_joints
+        left_arm_action = finalize_actions(
+            env.left_arm_current_qpos, env.left_arm_current_gripper_state
+        )
+        left_arm_action = np.repeat(
+            left_arm_action[None, :], len(right_arm_action), axis=0
+        )
+
+        actions = np.zeros(
+            (len(right_arm_action), len(env.robot.get_qpos().squeeze(0))),
+            dtype=np.float32,
+        )
+        actions[:, left_arm_index] = left_arm_action
+        actions[:, right_arm_index] = right_arm_action
+
+    elif right_arm_action is None and left_arm_action is not None:
+        left_arm_index = env.left_arm_joints + env.left_eef_joints
+        right_arm_index = env.right_arm_joints + env.right_eef_joints
+        right_arm_action = finalize_actions(
+            env.right_arm_current_qpos, env.right_arm_current_gripper_state
+        )
+        right_arm_action = np.repeat(
+            right_arm_action[None, :], len(left_arm_action), axis=0
+        )
+
+        actions = np.zeros(
+            (len(left_arm_action), len(env.robot.get_qpos().squeeze(0))),
+            dtype=np.float32,
+        )
+        actions[:, left_arm_index] = left_arm_action
+        actions[:, right_arm_index] = right_arm_action
+
+    else:
+        log_error("At least one arm action should be provided.")
+
+    actions = torch.from_numpy(actions).to(dtype=torch.float32).unsqueeze(1)
+    actions = list(actions.unbind(dim=0))
+    for i in tqdm(range(len(actions))):
+        action = actions[i]
+        obs, reward, terminated, truncated, info = env.step(action)
+    return actions


The drive() function modifies environment state via env.step() in a loop without any error handling. If step() raises an exception (e.g., due to invalid action), the loop will terminate abruptly. Additionally, the return value 'actions' is a list of tensors but the function documentation doesn't specify this return type. Add error handling and document the return type.

Copilot · 2026-01-22T08:41:52Z

embodichain/agents/hierarchy/validation_agent.py

+        response = view_selection_llm.invoke(messages).content.strip()
+
+        if response not in img_dirs:
+            raise ValueError(f"Invalid camera selection from LLM: {response}")
+
+        return response


The validation_llm.invoke() call in select_best_view_dir doesn't have any error handling. If the LLM call fails or times out, this will raise an unhandled exception. Additionally, if the LLM returns an unexpected value (not in img_dirs), a ValueError is raised, but this might be too strict - consider logging and falling back to a default camera instead.

Copilot · 2026-01-22T08:42:04Z

embodichain/agents/hierarchy/task_agent.py

+import time
+import cv2
+import glob
+import json


Import of 'json' is not used.

Suggested change

import json

Copilot · 2026-01-22T08:42:04Z

embodichain/agents/mllm/prompt/task_prompt.py

+# ----------------------------------------------------------------------------
+
+import torch
+import numpy as np


Import of 'np' is not used.

Suggested change

import numpy as np

Copilot · 2026-01-22T08:42:04Z

embodichain/agents/mllm/prompt/task_prompt.py

+    ChatPromptTemplate,
+    HumanMessagePromptTemplate,
+)
+from embodichain.utils.utility import encode_image, encode_image_from_path


Import of 'encode_image_from_path' is not used.

Suggested change

from embodichain.utils.utility import encode_image, encode_image_from_path

from embodichain.utils.utility import encode_image

Copilot · 2026-01-22T08:42:05Z

embodichain/toolkits/toolkits.py

+
+from abc import ABCMeta, abstractmethod
+import os
+import cv2


Import of 'cv2' is not used.

embodichain/agents/hierarchy/validation_agent.py

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

yangchen73 · 2026-01-22T08:49:02Z

@copilot open a new pull request to apply changes based on the comments in this thread

Copilot · 2026-01-22T08:49:14Z

@yangchen73 I've opened a new pull request, #86, to work on those changes. Once the pull request is ready, I'll request review from you.

yuecideng · 2026-01-22T08:33:59Z

configs/gym/agent/pour_water_agent/agent_config.json

@@ -0,0 +1,31 @@
+{    "TaskAgent": {


change the folder into pour_water_agent

Got it. I'v renamed this folder.

yuecideng · 2026-01-22T08:36:14Z

embodichain/agents/README.md

@@ -0,0 +1,97 @@
+# EmbodiAgent System


Add agent related docs into docs/features folder

Okay. I'll rewrite the docs and move it to docs/features.

yuecideng · 2026-01-22T08:41:35Z

embodichain/agents/README.md

+
+```bash
+# Set environment variables
+export AZURE_OPENAI_ENDPOINT="[https://your-endpoint.openai.azure.com/](https://your-endpoint.openai.azure.com/)"


Add docs to show how to change different LLM/VLM API

OK, I'll add it into the docs which discribe agent system.

yuecideng · 2026-01-22T08:44:05Z

configs/gym/agent/pour_water_agent/task_prompt_dual.txt

@@ -0,0 +1,5 @@
+Task:


Move prompt file to configs/gym/agent/{specific_task} would be better

Good suggustion.

yuecideng · 2026-01-22T08:46:26Z

embodichain/agents/prompts/atom_actions.txt

@@ -0,0 +1,136 @@
+### Atom Functions for Robot Arm Control


database folder has been removed. Put these reusable prompts into agent folder

yuecideng · 2026-01-22T08:47:37Z

embodichain/lab/gym/envs/action_bank/utils.py

    return "{}_{}_qpos".format(scope, tag)


+def get_control_part(env, agent_uid):


These teo functions seems useless. Please removed

yuecideng · 2026-01-22T08:51:07Z

embodichain/data/enum.py

+    return dofs > 10
+
+
+class HandQposNormalizer:


Remove useless objects in this file

Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> Co-authored-by: yangchen73 <115123709+yangchen73@users.noreply.github.com>

yhnsu · 2026-01-22T09:35:30Z

embodichain/lab/scripts/run_agent.py

+    )
+
+    # Run main function
+    main(args, env, gym_config)


from .run_env import main instead of rewrite it.

yhnsu · 2026-01-22T09:39:31Z

embodichain/lab/gym/envs/tasks/tableware/base_agent_env.py

+        code_file_path, kwargs, _ = self.generate_code_for_actions(
+            regenerate=regenerate
+        )
+        action_list = self.code_agent.act(code_file_path, **kwargs)


need to define self.action_length = len(action_list) for truncate detection inside the environment

I'v defined it.

yangchen73 added 30 commits January 21, 2026 10:37

Migrate the config of tasks

3cbbc21

Migrate execute script

f4dbddf

Migrate base_agent_env

8286998

Migrate prompt template

9af2785

Migrate the core code of agent

60c168c

Migrate prompt file

ee5a568

Update pyproject and gitignore

b1212af

Migrate generate video script

52a2b45

Migrate data engine

0b65baa

Fix: can not process 'func'

8bfeb7f

Migrate indices and mapping

8dba0da

Migrate environments: pour water and rearrangement

4195d30

Migrate API interfaces

96bb2b9

Migrate object

74eee5e

Migrate object

3421022

Migrate necessory files in toolkits

fa8f0fe

Migrate motion generation part

302608a

migrate lab.sim.utility

5baed9b

Migrate necessory files in toolkits

8411a35

Config LLM

64642ca

Extend enum

c3898b6

Add direction of database

2839548

Update enum

18cfc03

Migrate end effector

2ede9a0

Migrate robot

983fd9f

Update direction of database

974ee5f

Migrate robot interface

40957f1

Migrate two get control information function in action_bank utils

d4a3ae5

Migrate misc

c4eb89f

Fix: change data configs to functor format

189965d

yangchen73 added 4 commits January 22, 2026 15:59

Remove global indices and mapping

8130d75

Add license

b502411

update API interface

381216d

Reformat files

bc5cd49

yuecideng requested a review from Copilot January 22, 2026 08:32

Copilot started reviewing on behalf of yuecideng January 22, 2026 08:33 View session

Copilot AI reviewed Jan 22, 2026

View reviewed changes

Update embodichain/agents/hierarchy/validation_agent.py

a0506a9

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

Copilot AI mentioned this pull request Jan 22, 2026

[WIP] Migrate EmbodiAgent to EmbodiChain #86

Merged

22 tasks

yuecideng requested changes Jan 22, 2026

View reviewed changes

[WIP] Migrate EmbodiAgent to EmbodiChain (#86)

3a54b6f

Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> Co-authored-by: yangchen73 <115123709+yangchen73@users.noreply.github.com>

yhnsu reviewed Jan 22, 2026

View reviewed changes

yangchen73 added 15 commits January 22, 2026 17:40

Rename pour_water_agent

c443789

Move prompts

65236ce

Remove useless objects

123878a

Remove unused functions

218f80d

Remove old prompts

64a3fed

Remove old prompts

26bd1b0

Add path to load prompt

f5d7b82

Remove unused imported func

c2281a0

Formate files

b216e1a

Remove llm generated content

ad5f773

Add Documentation for EmbodiAgent

5450d65

Fix: define self.action_length = len(action_list) for truncate detection

cd922b5

Removed the rewrite main function

db6159f

Reformate file

0c4e785

Fix: package conflict

96d4e0e

	from embodichain.utils.utility import encode_image, encode_image_from_path
	from embodichain.utils.utility import encode_image

		return "{}_{}_qpos".format(scope, tag)


		def get_control_part(env, agent_uid):

Migrate EmbodiAgent #82

Are you sure you want to change the base?

Migrate EmbodiAgent #82

Uh oh!

Conversation

yangchen73 commented Jan 21, 2026

Description

Key Features

Key Related Files

Type of change

Screenshots

Checklist

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Copilot AI Jan 22, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Jan 22, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Jan 22, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Jan 22, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Jan 22, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Jan 22, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Jan 22, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Jan 22, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Jan 22, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

yangchen73 commented Jan 22, 2026

Uh oh!

Copilot AI commented Jan 22, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

yhnsu Jan 22, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

yhnsu Jan 22, 2026 •

edited

Loading