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

16
data_gen_dependencies/manip_solver.py
View File

@@ -2,8 +2,10 @@
import copy
import json
import os
from sqlite3 import dbapi2
import time
from pyboot.utils.log import Log
import trimesh
curPath = os.path.abspath(os.path.dirname(__file__))
@@ -210,6 +212,7 @@ def get_aligned_pose(active_obj, passive_obj, distance=0.01, N=1):
def load_task_solution(task_info):
Log.info('Loading task solution')
# task_info = json.load(open('%s/task_info.json'%task_dir, 'rb'))
stages = task_info['stages']
@@ -217,7 +220,8 @@ def load_task_solution(task_info):
'gripper': OmniObject('gripper')
}
for obj_info in task_info['objects']:
for i, obj_info in enumerate(task_info['objects']):
Log.info(f'Loading object: {obj_info["object_id"]} [Loading task solution: {i+1}/{len(task_info["objects"])}]')
obj_id = obj_info['object_id']
if obj_id == 'fix_pose':
obj = OmniObject('fix_pose')
@@ -245,15 +249,19 @@ def load_task_solution(task_info):
objects[obj_id] = obj
if hasattr(obj, 'part_ids'):
if hasattr(obj, 'part_joint_limits') and obj.part_joint_limits is not None:
obj_parts_joint_limits = obj.part_joint_limits
for part_id in obj.part_ids:
id = obj_id + '/%s' % part_id
objects[id] = copy.deepcopy(obj)
objects[id].name = id
objects[id].part_joint_limit = obj_parts_joint_limits[part_id]
#====== DEBUG START ======
print(f'debug: ')
import ipdb;ipdb.set_trace()
#====== DEBUG END ======
if hasattr(obj, 'part_joint_limits') and obj.part_joint_limits is not None:
objects[id].part_joint_limit = obj.part_joint_limits[part_id]
if len(obj.part_ids):
del objects[obj_id]
Log.success('Finished Loading task solution')
return stages, objects

9
data_gen_dependencies/object.py
View File

@@ -4,6 +4,7 @@ import pickle
import numpy as np
from grasp_nms import nms_grasp
from pyboot.utils.log import Log
from data_gen_dependencies.transforms import transform_coordinates_3d
@@ -122,22 +123,24 @@ class OmniObject:
}
if isinstance(grasp_files, str):
grasp_files = [grasp_files]
if isinstance(grasp_files, dict):
grasp_files = [grasp_files['grasp_pose_file']]
for grasp_file in grasp_files:
grasp_file = '%s/%s' % (obj_dir, grasp_file)
if not os.path.exists(grasp_file):
print('Grasp file %s not found' % grasp_file)
Log.warning('Grasp file %s not found' % grasp_file)
continue
_data = pickle.load(open(grasp_file, 'rb'))
_data['grasp_pose'] = np.array(_data['grasp_pose'])
_data['width'] = np.array(_data['width'])
if _data['grasp_pose'].shape[0] == 0:
print('Empty grasp pose in %s' % grasp_file)
Log.warning('Empty grasp pose in %s' % grasp_file)
continue
grasp_data['grasp_pose'].append(_data['grasp_pose'])
grasp_data['width'].append(_data['width'])
if len(grasp_data['grasp_pose']) == 0:
print('Empty grasp pose in %s' % grasp_file)
Log.warning('Empty grasp pose in %s' % grasp_file)
continue
grasp_data['grasp_pose'] = np.concatenate(grasp_data['grasp_pose'])

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()