finish data generate

This commit is contained in:
2025-09-05 19:18:37 +08:00
parent 21fbd5a323
commit 02440ceeb5
21 changed files with 17866 additions and 679 deletions

View File

@@ -1,17 +1,19 @@
# -*- coding: utf-8 -*-
import json
import os
import pickle
import random
import time
import random
import numpy as np
from pyboot.utils.log import Log
from data_gen_dependencies.transforms import calculate_rotation_matrix
from data_gen_dependencies.manip_solver import load_task_solution, generate_action_stages, split_grasp_stages
from data_gen_dependencies.base_agent import BaseAgent
from data_gen_dependencies.manip_solver import load_task_solution, generate_action_stages, split_grasp_stages
from data_gen_dependencies.omni_robot import IsaacSimRpcRobot
from data_gen_dependencies.transforms import calculate_rotation_matrix
class Agent(BaseAgent):
def __init__(self, robot: IsaacSimRpcRobot):
@@ -59,7 +61,7 @@ class Agent(BaseAgent):
self.add_object(object_info)
time.sleep(2)
self.arm = task_info["arm"]
self.arm = task_info["robot"]["arm"]
''' For A2D: Fix camera rotaton to look at target object '''
task_related_objs = []
@@ -195,142 +197,183 @@ class Agent(BaseAgent):
return False
return True
def run(self, task_file, camera_list, use_recording, workspaces, fps=10, render_semantic=False):
if not self.check_task_file(task_file):
Log.error("Task file bad: %s" % task_file)
return False
print("Start Task:", task_file)
self.reset()
self.attached_obj_id = None
def run(self, task_list, camera_list, use_recording, workspaces, fps=10, render_semantic=False):
for index, task_file in enumerate(task_list):
if not self.check_task_file(task_file):
Log.error("Task file bad: %s" % task_file)
continue
Log.info("start task: "+ task_file)
self.reset()
self.attached_obj_id = None
# import ipdb;ipdb.set_trace()
# import ipdb;ipdb.set_trace()
self.generate_layout(task_file)
# import ipdb;ipdb.set_trace()
self.robot.open_gripper(id='right')
self.robot.open_gripper(id='left')
self.generate_layout(task_file)
# import ipdb;ipdb.set_trace()
self.robot.open_gripper(id='right')
self.robot.open_gripper(id='left')
self.robot.reset_pose = {
'right': self.robot.get_ee_pose(ee_type='gripper', id='right'),
'left': self.robot.get_ee_pose(ee_type='gripper', id='left'),
}
print('Reset pose:', self.robot.reset_pose)
self.robot.reset_pose = {
'right': self.robot.get_ee_pose(ee_type='gripper', id='right'),
'left': self.robot.get_ee_pose(ee_type='gripper', id='left'),
}
#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)
# import ipdb; ipdb.set_trace()
if use_recording:
self.start_recording(task_name="[%s]" % (os.path.basename(task_file).split(".")[0]),
camera_prim_list=camera_list, fps=fps,
render_semantic=render_semantic) # TODO 录制判断
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)
# import ipdb; ipdb.set_trace()
if use_recording:
self.start_recording(task_name="[%s]" % (os.path.basename(task_file).split(".")[0]),
camera_prim_list=camera_list, fps=fps,
render_semantic=render_semantic) # TODO 录制判断
stage_id = -1
substages = None
for _stages in split_stages:
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')
action_stages = generate_action_stages(objects, _stages, self.robot)
if not len(action_stages):
success = False
print('No action stage generated.')
break
stage_id = -1
success = False
substages = None
for _stages in split_stages:
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')
action_stages = generate_action_stages(objects, _stages, self.robot)
if not len(action_stages):
success = False
print('No action stage generated.')
break
# Execution
success = True
# Execution
success = True
for action, substages in action_stages:
stage_id += 1
print('>>>> Stage [%d] <<<<' % (stage_id + 1))
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)
print(arm)
print(curr_pose)
interp_pose = init_pose.copy()
interp_pose[:3, 3] = curr_pose[:3, 3] + (init_pose[:3, 3] - curr_pose[:3, 3]) * 0.25
success = self.robot.move_pose(self.robot.reset_pose[arm], type='AvoidObs', arm=arm, block=True)
continue
for action, substages in action_stages:
stage_id += 1
Log.info(f'start action stage: {action} ({stage_id}/{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)
interp_pose = init_pose.copy()
interp_pose[:3, 3] = curr_pose[:3, 3] + (init_pose[:3, 3] - curr_pose[:3, 3]) * 0.25
success = self.robot.move_pose(self.robot.reset_pose[arm], type='AvoidObs', arm=arm, block=True)
continue
# if action in ['grasp', 'pick']:
# obj_id = substages.passive_obj_id
# if obj_id.split('/')[0] not in self.articulated_objs:
# self.robot.target_object = substages.passive_obj_id
while len(substages):
objects = self.update_objects(objects, arm=arm)
while len(substages):
# get next step actionddd
objects = self.update_objects(objects, arm=arm)
target_gripper_pose, motion_type, gripper_action, arm = substages.get_action(objects)
arm = extra_params.get('arm', 'right')
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)
target_gripper_pose, motion_type, gripper_action, arm = substages.get_action(objects)
arm = extra_params.get('arm', 'right')
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)
# execution action
if target_gripper_pose is not None:
self.robot.move_pose(target_gripper_pose, motion_type, arm=arm, block=True)
set_gripper_open = gripper_action == 'open'
set_gripper_close = gripper_action == 'close'
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)
set_gripper_open = gripper_action == 'open'
set_gripper_close = gripper_action == 'close'
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)
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)
objects = self.update_objects(objects, arm=arm)
# check sub-stage completion
objects['gripper'].obj_pose = self.robot.get_ee_pose(ee_type='gripper', id=arm)
objects = self.update_objects(objects, arm=arm)
success = substages.check_completion(objects)
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 success == False:
# import ipdb;ipdb.set_trace()
self.attached_obj_id = None
Log.error('Failed at sub-stage %d' % substages.step_id)
break
success = substages.check_completion(objects)
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 success == False:
# import ipdb;ipdb.set_trace()
self.attached_obj_id = None
print('Failed at sub-stage %d' % substages.step_id)
break
# attach grasped object to gripper # TODO avoid articulated objects
if gripper_action == 'close': # TODO 确定是grasp才行
self.attached_obj_id = substages.passive_obj_id
elif gripper_action == 'open':
self.attached_obj_id = None
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)
# attach grasped object to gripper # TODO avoid articulated objects
if gripper_action == 'close': # TODO 确定是grasp才行
self.attached_obj_id = substages.passive_obj_id
elif gripper_action == 'open':
self.attached_obj_id = None
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)
# change object position
num_against = substages.extra_params.get('against', 0)
against_area = substages.extra_params.get('against_range', [])
against_type = substages.extra_params.get('against_type', None)
if num_against > 0 and gripper_action == 'open' and action == 'pick' and (
against_type is not None):
parts = against_type.split('_')
need_move_objects = [passive_id]
if parts[0] == 'move' and parts[1] == 'with':
for workspace in objects:
if parts[2] in workspace:
need_move_objects.append(workspace)
## 目前就集体向一个方向移动
offset_y = random.uniform(0, 0.2)
poses = []
for move_object in need_move_objects:
response = self.robot.client.get_object_pose(f'/World/Objects/{move_object}')
# change object position
num_against = substages.extra_params.get('against', 0)
against_area = substages.extra_params.get('against_range', [])
against_type = substages.extra_params.get('against_type', None)
if num_against > 0 and gripper_action == 'open' and action == 'pick' and (
against_type is not None):
parts = against_type.split('_')
need_move_objects = [passive_id]
if parts[0] == 'move' and parts[1] == 'with':
for workspace in objects:
if parts[2] in workspace:
need_move_objects.append(workspace)
## 目前就集体向一个方向移动
offset_y = random.uniform(0, 0.2)
poses = []
for move_object in need_move_objects:
response = self.robot.client.get_object_pose(f'/World/Objects/{move_object}')
pos = [response.object_pose.position.x,
response.object_pose.position.y + offset_y,
response.object_pose.position.z]
quat_wxyz = np.array(
[
response.object_pose.rpy.rw,
response.object_pose.rpy.rx,
response.object_pose.rpy.ry,
response.object_pose.rpy.rz,
]
)
object_pose = {}
object_pose["prim_path"] = f'/World/Objects/{move_object}'
object_pose["position"] = pos
object_pose["rotation"] = quat_wxyz
poses.append(object_pose)
print(poses)
self.robot.client.SetObjectPose(poses, [])
self.robot.client.DetachObj()
elif (num_against > 0 and gripper_action == 'open' and action == 'pick') or \
(num_against > 0 and action == 'place' and gripper_action == 'close'):
# import ipdb;ipdb.set_trace()
x_min, x_max = 999.0, -999.0
y_min, y_max = 999.0, -999.0
if against_area:
selected_against_area = random.choice(against_area)
if selected_against_area in workspaces:
position = workspaces[selected_against_area]['position']
size = workspaces[selected_against_area]['size']
x_min, x_max = position[0] - size[0] / 2, position[0] + size[0] / 2
y_min, y_max = position[1] - size[1] / 2, position[1] + size[1] / 2
x_size, y_size, z_size = objects[passive_id].info['size']
up_axis = objects[passive_id].info['upAxis']
axis_mapping = {
'x': (y_size / 2000.0, z_size / 2000.0),
'y': (x_size / 2000.0, z_size / 2000.0),
'z': (x_size / 2000.0, y_size / 2000.0)
}
dimensions = axis_mapping[up_axis[0]]
distance = np.linalg.norm(dimensions)
x_min += distance * 1.5
x_max -= distance * 1.5
y_min += distance * 1.5
y_max -= distance * 1.5
else:
print("against_range not set")
continue
response = self.robot.client.get_object_pose(f'/World/Objects/{passive_id}')
pos = [response.object_pose.position.x,
response.object_pose.position.y + offset_y,
response.object_pose.position.z]
response.object_pose.position.y,
response.object_pose.position.z + 0.02]
quat_wxyz = np.array(
[
response.object_pose.rpy.rw,
@@ -339,121 +382,44 @@ class Agent(BaseAgent):
response.object_pose.rpy.rz,
]
)
# import ipdb;ipdb.set_trace()
# if position is close to the gripper, then random position again
while True:
pos[0] = random.uniform(x_min, x_max)
pos[1] = random.uniform(y_min, y_max)
distance = np.linalg.norm(
[pos[0] - target_gripper_pose[0][3], pos[1] - target_gripper_pose[1][3]])
if distance >= 0.2:
break
poses = []
object_pose = {}
object_pose["prim_path"] = f'/World/Objects/{move_object}'
object_pose["prim_path"] = f'/World/Objects/{passive_id}'
object_pose["position"] = pos
object_pose["rotation"] = quat_wxyz
poses.append(object_pose)
print(poses)
self.robot.client.SetObjectPose(poses, [])
self.robot.client.DetachObj()
elif (num_against > 0 and gripper_action == 'open' and action == 'pick') or \
(num_against > 0 and action == 'place' and gripper_action == 'close'):
# import ipdb;ipdb.set_trace()
x_min, x_max = 999.0, -999.0
y_min, y_max = 999.0, -999.0
if against_area:
selected_against_area = random.choice(against_area)
if selected_against_area in workspaces:
position = workspaces[selected_against_area]['position']
size = workspaces[selected_against_area]['size']
x_min, x_max = position[0] - size[0] / 2, position[0] + size[0] / 2
y_min, y_max = position[1] - size[1] / 2, position[1] + size[1] / 2
x_size, y_size, z_size = objects[passive_id].info['size']
up_axis = objects[passive_id].info['upAxis']
axis_mapping = {
'x': (y_size / 2000.0, z_size / 2000.0),
'y': (x_size / 2000.0, z_size / 2000.0),
'z': (x_size / 2000.0, y_size / 2000.0)
}
dimensions = axis_mapping[up_axis[0]]
distance = np.linalg.norm(dimensions)
x_min += distance * 1.5
x_max -= distance * 1.5
y_min += distance * 1.5
y_max -= distance * 1.5
else:
print("against_range not set")
continue
response = self.robot.client.get_object_pose(f'/World/Objects/{passive_id}')
pos = [response.object_pose.position.x,
response.object_pose.position.y,
response.object_pose.position.z + 0.02]
quat_wxyz = np.array(
[
response.object_pose.rpy.rw,
response.object_pose.rpy.rx,
response.object_pose.rpy.ry,
response.object_pose.rpy.rz,
]
)
# import ipdb;ipdb.set_trace()
# if position is close to the gripper, then random position again
while True:
pos[0] = random.uniform(x_min, x_max)
pos[1] = random.uniform(y_min, y_max)
distance = np.linalg.norm(
[pos[0] - target_gripper_pose[0][3], pos[1] - target_gripper_pose[1][3]])
if distance >= 0.2:
break
poses = []
object_pose = {}
object_pose["prim_path"] = f'/World/Objects/{passive_id}'
object_pose["position"] = pos
object_pose["rotation"] = quat_wxyz
poses.append(object_pose)
print(poses)
self.robot.client.SetObjectPose(poses, [])
self.robot.client.DetachObj()
# if num_against > 0 and action == 'place' and gripper_action == 'close':
# offset = random.uniform(-0.1, 0.1)
# response = self.robot.client.get_object_pose('/World/Objects/%s'%passive_id)
# pos = [response.object_pose.position.x,
# response.object_pose.position.y - offset,
# response.object_pose.position.z]
# quat_wxyz = np.array(
# [
# response.object_pose.rpy.rw,
# response.object_pose.rpy.rx,
# response.object_pose.rpy.ry,
# response.object_pose.rpy.rz,
# ]
# )
# # import ipdb;ipdb.set_trace()
# poses = []
# object_pose ={}
# object_pose["prim_path"] = '/World/Objects/' + passive_id
# object_pose["position"] =pos
# object_pose["rotation"] = quat_wxyz
# poses.append(object_pose)
# self.robot.client.SetObjectPose(poses, [])
# index += 1
# self.robot.client.DetachObj()
print(poses)
self.robot.client.SetObjectPose(poses, [])
self.robot.client.DetachObj()
if success == False:
self.attached_obj_id = None
break
if self.attached_obj_id is not None:
if self.attached_obj_id.split('/')[0] not in self.articulated_objs:
self.robot.client.DetachObj()
self.robot.client.AttachObj(prim_paths=['/World/Objects/' + self.attached_obj_id],
is_right=arm == 'right')
if success == False:
self.attached_obj_id = None
break
if self.attached_obj_id is not None:
if self.attached_obj_id.split('/')[0] not in self.articulated_objs:
self.robot.client.DetachObj()
self.robot.client.AttachObj(prim_paths=['/World/Objects/' + self.attached_obj_id],
is_right=arm == 'right')
if success == False:
break
time.sleep(0.5)
if self.attached_obj_id is None:
self.robot.client.DetachObj()
self.robot.client.stop_recording()
# try:
# step_id = substages.step_id if substages is not None and len(substages) else -1
# except:
# step_id = -1
step_id = -1
fail_stage_step = [stage_id, step_id] if success == False else [-1, -1]
time.sleep(0.5)
if self.attached_obj_id is None:
self.robot.client.DetachObj()
self.robot.client.stop_recording()
step_id = -1
fail_stage_step = [stage_id, step_id] if not success else [-1, -1]
task_info_saved = task_info.copy()
self.robot.client.SendTaskStatus(success, fail_stage_step)
if success:
print(">>>>>>>>>>>>>>>> Success!!!!!!!!!!!!!!!!")
task_info_saved = task_info.copy()
self.robot.client.SendTaskStatus(success, fail_stage_step)
if success:
print(">>>>>>>>>>>>>>>> Success!!!!!!!!!!!!!!!!")
return True