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

finish task_gen

Browse Source
This commit is contained in:
hofee
2025-09-05 11:10:42 +08:00
parent da022d4f83
commit 4e51158215
17 changed files with 82 additions and 50 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
app.py
View File

@@ -4,8 +4,8 @@ from runners.data_generator import DataGenerator
from runners.data_recorder import DataRecorder
from runners.task_templates_divider import TaskTemplatesDivider
@PybootApplication("agent")
class AgentApp:
@PybootApplication("main")
class MainApp:
@staticmethod
def start():
TaskGenerator(config_path="configs/generate_task_config.yaml").run()

6
configs/divide_task_config.yaml
View File

@@ -3,8 +3,8 @@ runner:
name: "divide_task"
root_dir: "workspace"
divide:
task_templates_root_dir: "/home/ubuntu/Projects/tasks_divide/pick&place"
output_task_templates_dir: "/home/ubuntu/Projects/Refactory_isaac_sim/input/templates/task_templates"
output_template_targets_dir: "/home/ubuntu/Projects/Refactory_isaac_sim/input/templates/template_targets"
input_task_templates_root_dir: "/home/ubuntu/Projects/tasks_divide/pick&place"
output_task_templates_dir:
output_template_targets_dir:
divide_num: 100
total_nums: 10000

4
configs/generate_task_config.yaml
View File

@@ -3,4 +3,6 @@ runner:
name: "generate_task"
root_dir: "workspace"
generate:
target_task_templates: "task_template_target_0.json"
input_data_root: "/home/ubuntu/Projects/docker_isaac_sim/input/data"
input_target_task_templates_path: "workspace/divide_task/template_targets/task_template_target_0.json"
output_task_root_dir:

29
runners/task_generator.py
View File

@@ -1,12 +1,33 @@
import json
import os
from pyboot.runner import Runner
from others.task_generate import OldTaskGenerator
from pyboot.utils.log import Log
from task_gen_dependencies.task_generate import OldTaskGenerator
class TaskGenerator(Runner):
def __init__(self, config_path: str):
super().__init__(config_path)
self.generate_config = self.config["generate"]
self.input_target_task_templates_path = self.generate_config["input_target_task_templates_path"]
self.input_data_root = self.generate_config["input_data_root"]
self.output_task_root_dir = self.generate_config.get("output_task_root_dir", None)
if self.output_task_root_dir is None:
self.output_task_root_dir = os.path.join(self.workspace_path, "task_root_dir")
self.target_task_templates = json.load(open(self.input_target_task_templates_path, "r"))
def run(self):
pass
for task_template_name, task_template_path in self.target_task_templates.items():
task_template = json.load(open(task_template_path, "r"))
self.generate_from_template(task_template_name, task_template)
Log.success(f"Generated {task_template['recording_setting']['num_of_episode']} tasks from <{task_template_name}>")
def generate_from_template(self, template: dict):
pass
def generate_from_template(self, template_name: str, template: dict):
task_dir = os.path.join(self.output_task_root_dir, template_name)
task_num = template["recording_setting"]["num_of_episode"]
task_name = template["task"]
old_task_generator = OldTaskGenerator(template, self.input_data_root)
if not os.path.exists(task_dir):
os.makedirs(task_dir)
old_task_generator.generate_tasks(save_path=task_dir, task_num=task_num, task_name=task_name)

18
runners/task_templates_divider.py
View File

@@ -10,9 +10,13 @@ class TaskTemplatesDivider(Runner):
super().__init__(config_path)
self.divide_config = self.config["divide"]
self.task_templates_root_dir = self.divide_config["task_templates_root_dir"]
self.output_task_templates_dir = self.divide_config["output_task_templates_dir"]
self.output_template_targets_dir = self.divide_config["output_template_targets_dir"]
self.input_task_templates_root_dir = self.divide_config["input_task_templates_root_dir"]
self.output_task_templates_dir = self.divide_config.get("output_task_templates_dir", None)
if self.output_task_templates_dir is None:
self.output_task_templates_dir = os.path.join(self.workspace_path, "task_templates")
self.output_template_targets_dir = self.divide_config.get("output_template_targets_dir", None)
if self.output_template_targets_dir is None:
self.output_template_targets_dir = os.path.join(self.workspace_path, "template_targets")
self.divide_num = self.divide_config["divide_num"]
self.total_nums = self.divide_config["total_nums"]
@@ -26,11 +30,11 @@ class TaskTemplatesDivider(Runner):
def load_all_task_templates(self):
task_list = []
for task_template_dir in os.listdir(self.task_templates_root_dir):
if os.path.isdir(os.path.join(self.task_templates_root_dir, task_template_dir)):
for file in os.listdir(os.path.join(self.task_templates_root_dir, task_template_dir)):
for task_template_dir in os.listdir(self.input_task_templates_root_dir):
if os.path.isdir(os.path.join(self.input_task_templates_root_dir, task_template_dir)):
for file in os.listdir(os.path.join(self.input_task_templates_root_dir, task_template_dir)):
if file.endswith('.json'):
task_list.append(os.path.join(self.task_templates_root_dir, task_template_dir, file))
task_list.append(os.path.join(self.input_task_templates_root_dir, task_template_dir, file))
Log.success(f"Loaded {len(task_list)} tasks")
return task_list

4
others/cost_function.py → task_gen_dependencies/cost_function.py
View File

@@ -1,7 +1,7 @@
import numpy as np
from others.sdf import get_distance_with_sdf
from others.transform_utils import transform_points, random_point
from task_gen_dependencies.sdf import get_distance_with_sdf
from task_gen_dependencies.transform_utils import transform_points, random_point
def cost_without_collision(obj_active, obj_passive, threshold=1, **kwargs):

0
others/layout_2d.py → task_gen_dependencies/layout_2d.py
View File

17
others/layout_object.py → task_gen_dependencies/layout_object.py
View File

@@ -5,9 +5,9 @@ import trimesh
from scipy.interpolate import RegularGridInterpolator
from scipy.spatial.transform import Rotation as R
from others.object import OmniObject
from others.sdf import compute_sdf_from_obj_surface
from others.transform_utils import farthest_point_sampling, random_point
from task_gen_dependencies.object import OmniObject
from task_gen_dependencies.sdf import compute_sdf_from_obj_surface
from task_gen_dependencies.transform_utils import farthest_point_sampling, random_point, get_bott_up_point
def load_and_prepare_mesh(obj_path, up_axis, scale=1.0):
@@ -63,9 +63,12 @@ class LayoutObject(OmniObject):
if mesh_points.shape[0] > N_collision_points:
self.collision_points = farthest_point_sampling(mesh_points, N_collision_points) # 碰撞检测点
self.anchor_points = {'top': random_point(self.anchor_points['top'], 3)[np.newaxis, :],
'buttom': random_point(self.anchor_points['buttom'], 3)[np.newaxis, :]}
self.anchor_points = {}
self.anchor_points['top'] = get_bott_up_point(mesh_points, 1.5, descending=False)
self.anchor_points['buttom'] = get_bott_up_point(mesh_points, 1.5, descending=True)
self.size = self.mesh.extents.copy()
self.anchor_points['top'] = random_point(self.anchor_points['top'], 3)[np.newaxis, :]
self.anchor_points['buttom'] = random_point(self.anchor_points['buttom'], 3)[np.newaxis, :]
self.size = self.mesh.extents.copy()
self.up_axis = up_axis[0]

0
others/multi_add_util.py → task_gen_dependencies/multi_add_util.py
View File

2
others/object.py → task_gen_dependencies/object.py
View File

@@ -4,7 +4,7 @@ import pickle
import json
import os
from others.transforms import transform_coordinates_3d
from task_gen_dependencies.transforms import transform_coordinates_3d
from grasp_nms import nms_grasp
class OmniObject:

2
others/sdf.py → task_gen_dependencies/sdf.py
View File

@@ -1,6 +1,6 @@
import numpy as np
import open3d as o3d
from others.transform_utils import transform_points
from task_gen_dependencies.transform_utils import transform_points
def compute_sdf_from_obj_surface(mesh, resolution=2): # 2mm
if mesh is None or len(mesh.vertices) == 0:

12
others/solver_2d.py → task_gen_dependencies/solver_2d.py
View File

@@ -1,10 +1,10 @@
from scipy.spatial.transform import Rotation as R
from others.layout_2d import DFS_Solver_Floor
from others.multi_add_util import *
from others.utils import axis_to_quaternion, quaternion_rotate, get_rotation_matrix_from_quaternion, \
from task_gen_dependencies.layout_2d import DFS_Solver_Floor
from task_gen_dependencies.multi_add_util import *
from task_gen_dependencies.utils import axis_to_quaternion, quaternion_rotate, get_rotation_matrix_from_quaternion, \
get_quaternion_from_rotation_matrix, get_xyz_euler_from_quaternion
from pyboot.utils.log import Log
def rotate_along_axis(target_affine, angle_degrees, rot_axis='z', use_local=True):
"""
@@ -354,9 +354,7 @@ class LayoutSolver2D:
# objects = [obj[0] for obj in placed_objects]
# visualize_objects(objects, self.plane_width, self.plane_height,str(obj_id)+label_flag+".jpg")
if len(fail_objs)>0:
print("*******no solution objects************")
print(fail_objs)
print("******************")
Log.warning("failed objects in layout 2d: " + str(fail_objs))
for obj_id in objs_succ:
self.objects[obj_id].obj_pose = rotate_along_axis(self.objects[obj_id].obj_pose, random.uniform(-self.angle_step/2.0, self.angle_step/2.0), rot_axis='z', use_local=False)

4
others/solver_3d.py → task_gen_dependencies/solver_3d.py
View File

@@ -3,8 +3,8 @@ from scipy.interpolate import RegularGridInterpolator
from scipy.optimize import dual_annealing
from scipy.spatial.transform import Rotation as R
from others.cost_function import get_cost_func
from others.transform_utils import unnormalize_vars, \
from task_gen_dependencies.cost_function import get_cost_func
from task_gen_dependencies.transform_utils import unnormalize_vars, \
normalize_vars, pose2mat, euler2quat
np.random.seed(0)

30
others/task_generate.py → task_gen_dependencies/task_generate.py
View File

@@ -3,13 +3,13 @@
import os
import json
import numpy as np
from others.object import OmniObject
from others.layout_object import LayoutObject
from others.solver_2d import LayoutSolver2D
from others.solver_3d import LayoutSolver3D
from task_gen_dependencies.object import OmniObject
from task_gen_dependencies.layout_object import LayoutObject
from task_gen_dependencies.solver_2d import LayoutSolver2D
from task_gen_dependencies.solver_3d import LayoutSolver3D
from utils.pose import PoseUtil
from pyboot.utils.log import Log
def list_to_dict(data: list):
tmp = {}
@@ -54,18 +54,18 @@ class LayoutGenerator():
objs_succ = self.solver_2d(self.key_obj_ids_2d, self.succ_obj_ids, object_extent=30, start_with_edge=True,
key_obj=True, initial_angle=0)
self.update_obj_info(objs_succ)
print('-- 2d layout done --')
#print('-- 2d layout done --')
if len(self.key_obj_ids_3d) > 0:
objs_succ = self.solver_3d(self.key_obj_ids_3d, self.succ_obj_ids, constraint=self.constraint)
self.update_obj_info(objs_succ)
print('-- 3d layout done --')
#print('-- 3d layout done --')
if len(self.extra_obj_ids) > 0:
objs_succ = self.solver_2d(self.extra_obj_ids, self.succ_obj_ids, object_extent=30, start_with_edge=False,
key_obj=False)
self.update_obj_info(objs_succ)
print('-- extra layout done --')
#print('-- extra layout done --')
''' Check completion '''
res_infos = []
@@ -94,8 +94,8 @@ class LayoutGenerator():
class OldTaskGenerator:
def __init__(self, task_template):
self.data_root = os.path.dirname(os.path.dirname(__file__)) + "/assets"
def __init__(self, task_template, data_root):
self.data_root = data_root
self.init_info(task_template)
def _load_json(self, relative_path):
@@ -135,7 +135,7 @@ class OldTaskGenerator:
else:
obj_dir = os.path.join(self.data_root, obj["data_info_dir"])
if "metadata" in obj:
print(obj["metadata"])
#print(obj["metadata"])
info = obj["metadata"]["info"]
info["interaction"] = obj["metadata"]["interaction"]
else:
@@ -215,6 +215,10 @@ class OldTaskGenerator:
output_file = os.path.join(save_path, f'{task_name}_%d.json' % (i))
self.task_template['objects'] = []
self.task_template['objects'] += self.fix_obj_infos
self.task_template['robot'] = {
"init_position" : self.robot_init_pose["position"],
"init_rotation" : self.robot_init_pose["quaternion"],
}
flag_failed = False
for key in self.layouts:
@@ -231,10 +235,10 @@ class OldTaskGenerator:
self.task_template['objects'].append(fix_pose_dict)
break
if flag_failed:
print(f"Failed to place key object, skipping")
Log.warning(f"Failed to place key object, skipping")
continue
print('Saved task json to %s' % output_file)
Log.info('Saved task json to %s' % output_file)
with open(output_file, 'w') as f:
json.dump(self.task_template, f, indent=4)

0
others/transform_utils.py → task_gen_dependencies/transform_utils.py
View File

0
others/transforms.py → task_gen_dependencies/transforms.py
View File

0
others/utils.py → task_gen_dependencies/utils.py
View File