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

finshi press action

Browse Source
This commit is contained in:
hofee
2025-09-24 11:11:24 +08:00
parent 8b89b2e35b
commit 330d903047
16 changed files with 296 additions and 91 deletions
Download Patch File Download Diff File Expand all files Collapse all files

37
data_gen_dependencies/omniagent.py
View File

@@ -133,17 +133,16 @@ class Agent(BaseAgent):
objects['fix_pose'].obj_pose[:3, :3] = rotation_matrix
continue
# TODO(unify part_name and obj_name)
if '/' in obj_id:
obj = objects[obj_id]
if obj.is_articulated:
obj_name = obj_id.split('/')[0]
part_name = obj_id.split('/')[1]
object_joint_state = self.robot.client.get_object_joint('/World/Objects/%s' % obj_name)
for joint_name, joint_position, joint_velocity in zip(object_joint_state.joint_names,
object_joint_state.joint_positions,
object_joint_state.joint_velocities):
if joint_name[-1] == part_name[-1]:
objects[obj_id].joint_position = joint_position
objects[obj_id].joint_velocity = joint_velocity
if joint_name == obj.joint_name:
obj.joint_position = joint_position
obj.joint_velocity = joint_velocity
objects[obj_id].obj_pose = self.robot.get_prim_world_pose('/World/Objects/%s' % obj_id)
if hasattr(objects[obj_id], 'info') and 'simple_place' in objects[obj_id].info and objects[obj_id].info[
@@ -239,14 +238,15 @@ class Agent(BaseAgent):
camera_prim_list=camera_list, fps=fps,
render_semantic=render_semantic) # TODO 录制判断
stage_id = -1
success = False
substages = None
for _stages in split_stages:
stage_id = -1
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)}]')
Log.info(f'Generating action stages')
action_stages = generate_action_stages(objects, _stages, self.robot)
if not len(action_stages):
success = False
@@ -255,12 +255,14 @@ class Agent(BaseAgent):
# Execution
success = True
last_gripper_state = None
for action, substages in action_stages:
stage_id += 1
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']:
init_pose = self.robot.reset_pose[arm]
curr_pose = self.robot.get_ee_pose(ee_type='gripper', id=arm)
@@ -277,7 +279,6 @@ class Agent(BaseAgent):
self.robot.client.set_frame_state(action, substages.step_id, active_id, passive_id,
self.attached_obj_id is not None, arm=arm,
target_pose=target_gripper_pose)
if target_gripper_pose is not None:
self.robot.move_pose(target_gripper_pose, motion_type, arm=arm, block=True)
@@ -286,13 +287,15 @@ class Agent(BaseAgent):
self.robot.client.set_frame_state(action, substages.step_id, active_id, passive_id,
self.attached_obj_id is not None, set_gripper_open,
set_gripper_close, arm=arm, target_pose=target_gripper_pose)
self.robot.set_gripper_action(gripper_action, arm=arm)
if gripper_action == 'open':
time.sleep(1)
if last_gripper_state != gripper_action and gripper_action is not None:
self.robot.set_gripper_action(gripper_action, arm=arm)
last_gripper_state = gripper_action
if gripper_action == 'open':
time.sleep(1)
self.robot.client.set_frame_state(action, substages.step_id, active_id, passive_id,
self.attached_obj_id is not None, arm=arm,
target_pose=target_gripper_pose)
self.robot.client.set_frame_state(action, substages.step_id, active_id, passive_id,
self.attached_obj_id is not None, arm=arm,
target_pose=target_gripper_pose)
# check sub-stage completion
objects['gripper'].obj_pose = self.robot.get_ee_pose(ee_type='gripper', id=arm)
@@ -305,7 +308,7 @@ class Agent(BaseAgent):
if success == False:
# import ipdb;ipdb.set_trace()
self.attached_obj_id = None
Log.error('Failed at sub-stage %d' % substages.step_id)
Log.warning('Failed at sub-stage %d' % substages.step_id)
break
# attach grasped object to gripper # TODO avoid articulated objects