hofee/gen_data_agent
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

add articulate debug in load_task_solution at part_joint_limits keyerror problem

Browse Source
This commit is contained in:
hofee
2025-09-15 20:01:32 +08:00
parent 3cdeb975b7
commit 8b89b2e35b
3 changed files with 33 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files

18
data_gen_dependencies/omniagent.py
View File

@@ -37,6 +37,7 @@ class Agent(BaseAgent):
)
def generate_layout(self, task_file):
Log.info(f'Generating layout for task: {task_file}')
self.task_file = task_file
with open(task_file, "r") as f:
task_info = json.load(f)
@@ -50,6 +51,7 @@ class Agent(BaseAgent):
task_info['objects'][i]['mass'] = 10
break
Log.info('Adding objects [Generating layout: 1/5]')
self.articulated_objs = []
for object_info in task_info["objects"]:
if object_info['object_id'] == 'fix_pose':
@@ -61,8 +63,8 @@ class Agent(BaseAgent):
self.add_object(object_info)
time.sleep(2)
Log.info('Setting arm [Generating layout: 2/5]')
self.arm = task_info["robot"]["arm"]
''' For A2D: Fix camera rotaton to look at target object '''
task_related_objs = []
for stage in task_info['stages']:
@@ -80,6 +82,7 @@ class Agent(BaseAgent):
target_lookat_point = np.mean(np.stack(target_lookat_point), axis=0)
self.robot.client.SetTargetPoint(target_lookat_point.tolist())
Log.info('Setting material [Generating layout: 3/5]')
''' Set material '''
material_infos = []
if "object_with_material" in task_info:
@@ -89,6 +92,8 @@ class Agent(BaseAgent):
self.robot.client.SetMaterial(material_infos)
time.sleep(0.3)
Log.info('Setting light [Generating layout: 4/5]')
''' Set light '''
light_infos = []
if "lights" in task_info:
@@ -98,6 +103,7 @@ class Agent(BaseAgent):
self.robot.client.SetLight(light_infos)
time.sleep(0.3)
Log.info('Setting camera [Generating layout: 5/5]')
''' Set camera'''
if "cameras" in task_info:
for cam_id in task_info['cameras']:
@@ -108,6 +114,7 @@ class Agent(BaseAgent):
cam_info['focal_length'], cam_info['horizontal_aperture'], cam_info['vertical_aperture'],
cam_info['is_local']
)
Log.success('Finished Generating Layout')
def update_objects(self, objects, arm='right'):
# update gripper pose
@@ -186,6 +193,8 @@ class Agent(BaseAgent):
file = objs_interaction[passive_obj_id]['passive']['grasp'][primitive]
if isinstance(file, list):
file = file[0]
if isinstance(file, dict):
file = file['grasp_pose_file']
grasp_file = os.path.join(data_root, obj_dir, file)
if not os.path.exists(grasp_file):
Log.error('-- Grasp file not exist: %s' % grasp_file)
@@ -202,7 +211,7 @@ class Agent(BaseAgent):
if not self.check_task_file(task_file):
Log.error("Task file bad: %s" % task_file)
continue
Log.info("start task: "+ task_file)
Log.info("Start task: "+ task_file)
self.reset()
self.attached_obj_id = None
@@ -220,6 +229,7 @@ class Agent(BaseAgent):
#print('Reset pose:', self.robot.reset_pose)
task_info = json.load(open(task_file, 'rb'))
stages, objects = load_task_solution(task_info)
objects = self.update_objects(objects)
split_stages = split_grasp_stages(stages)
@@ -236,6 +246,7 @@ class Agent(BaseAgent):
extra_params = _stages[0].get('extra_params', {})
active_id, passive_id = _stages[0]['active']['object_id'], _stages[0]['passive']['object_id']
arm = extra_params.get('arm', 'right')
Log.info(f'Generating action stages [Executing task: {stage_id+1}/{len(split_stages)}]')
action_stages = generate_action_stages(objects, _stages, self.robot)
if not len(action_stages):
success = False
@@ -247,7 +258,7 @@ class Agent(BaseAgent):
for action, substages in action_stages:
stage_id += 1
Log.info(f'start action stage: {action} ({stage_id}/{len(action_stages)})')
Log.info(f'Start action stage: {action} [Executing task: {stage_id+1}/{len(action_stages)}]')
active_id, passive_id = _stages[stage_id]['active']['object_id'], _stages[stage_id]['passive'][
'object_id']
if action in ['reset']:
@@ -410,6 +421,7 @@ class Agent(BaseAgent):
is_right=arm == 'right')
if success == False:
break
Log.success('Finished executing task.')
time.sleep(0.5)
if self.attached_obj_id is None:
self.robot.client.DetachObj()