diff --git a/configs/gym/blocks_ranking_rgb/cobot_magic_3cam.json b/configs/gym/blocks_ranking_rgb/cobot_magic_3cam.json new file mode 100644 index 00000000..93245784 --- /dev/null +++ b/configs/gym/blocks_ranking_rgb/cobot_magic_3cam.json @@ -0,0 +1,283 @@ +{ + "id": "BlocksRankingRGB-v1", + "max_episodes": 5, + "env": { + "events": { + "random_light": { + "func": "randomize_light", + "mode": "interval", + "interval_step": 10, + "params": { + "entity_cfg": {"uid": "light_1"}, + "position_range": [[-0.5, -0.5, 2], [0.5, 0.5, 2]], + "color_range": [[0.6, 0.6, 0.6], [1, 1, 1]], + "intensity_range": [50.0, 100.0] + } + }, + "init_block_1_pose": { + "func": "randomize_rigid_object_pose", + "mode": "reset", + "params": { + "entity_cfg": {"uid": "block_1"}, + "position_range": [[0.445, -0.08, 0.86], [1.005, 0.05, 0.86]], + "rotation_range": [[0, 0, -43.0], [0, 0, 43.0]], + "relative_position": false + } + }, + "init_block_2_pose": { + "func": "randomize_rigid_object_pose", + "mode": "reset", + "params": { + "entity_cfg": {"uid": "block_2"}, + "position_range": [[0.445, -0.08, 0.86], [1.005, 0.05, 0.86]], + "rotation_range": [[0, 0, -43.0], [0, 0, 43.0]], + "relative_position": false + } + }, + "init_block_3_pose": { + "func": "randomize_rigid_object_pose", + "mode": "reset", + "params": { + "entity_cfg": {"uid": "block_3"}, + "position_range": [[0.445, -0.08, 0.86], [1.005, 0.05, 0.86]], + "rotation_range": [[0, 0, -43.0], [0, 0, 43.0]], + "relative_position": false + } + }, + "init_block_sizes": { + "func": "randomize_rigid_objects_scale", + "mode": "reset", + "params": { + "entity_cfgs": [ + {"uid": "block_1"}, + {"uid": "block_2"}, + {"uid": "block_3"} + ], + "scale_factor_range": [[0.75], [1.25]], + "same_scale_all_axes": true, + "shared_sample": true + } + }, + "set_block_1_color": { + "func": "set_rigid_object_visual_material", + "mode": "startup", + "params": { + "entity_cfg": {"uid": "block_1"}, + "mat_cfg": { + "uid": "red_block_mat", + "base_color": [1.0, 0.0, 0.0, 1.0], + "metallic": 0.0, + "roughness": 0.5 + } + } + }, + "set_block_2_color": { + "func": "set_rigid_object_visual_material", + "mode": "startup", + "params": { + "entity_cfg": {"uid": "block_2"}, + "mat_cfg": { + "uid": "green_block_mat", + "base_color": [0.0, 1.0, 0.0, 1.0], + "metallic": 0.0, + "roughness": 0.5 + } + } + }, + "set_block_3_color": { + "func": "set_rigid_object_visual_material", + "mode": "startup", + "params": { + "entity_cfg": {"uid": "block_3"}, + "mat_cfg": { + "uid": "blue_block_mat", + "base_color": [0.0, 0.0, 1.0, 1.0], + "metallic": 0.0, + "roughness": 0.5 + } + } + } + }, + "observations": { + "norm_robot_eef_joint": { + "func": "normalize_robot_joint_data", + "mode": "modify", + "name": "robot/qpos", + "params": { + "joint_ids": [12, 13, 14, 15] + } + }, + "block_1_pose": { + "func": "get_rigid_object_pose", + "mode": "add", + "name": "block_1_pose", + "params": { + "entity_cfg": {"uid": "block_1"} + } + }, + "block_2_pose": { + "func": "get_rigid_object_pose", + "mode": "add", + "name": "block_2_pose", + "params": { + "entity_cfg": {"uid": "block_2"} + } + }, + "block_3_pose": { + "func": "get_rigid_object_pose", + "mode": "add", + "name": "block_3_pose", + "params": { + "entity_cfg": {"uid": "block_3"} + } + } + } + }, + "robot": { + "uid": "CobotMagic", + "robot_type": "CobotMagic", + "init_pos": [0.0, 0.0, 0.7775], + "init_qpos": [-0.3,0.3,1.0,1.0,-1.2,-1.2,0.0,0.0,0.6,0.6,0.0,0.0,0.05,0.05,0.05,0.05] + }, + "sensor": [ + { + "sensor_type": "StereoCamera", + "uid": "cam_high", + "width": 960, + "height": 540, + "enable_mask": true, + "enable_depth": true, + "left_to_right_pos": [0.059684025824163614, 0, 0], + "intrinsics": [453.851402686215, 453.8347628855552, 469.827725021235, 258.6656181845155], + "intrinsics_right": [453.4536601653505, 453.3306024582175, 499.13697412367776, 297.7176248477935], + "extrinsics": { + "eye": [0.35368482807598, 0.014695524383058989, 1.4517046071614774], + "target": [0.7186357573287919, -0.054534732904795505, 0.5232553674540066], + "up": [0.9306678549330372, -0.0005600064212467153, 0.3658647703553347] + } + }, + { + "sensor_type": "Camera", + "uid": "cam_right_wrist", + "width": 640, + "height": 480, + "enable_mask": true, + "intrinsics": [488.1665344238281, 488.1665344238281, 322.7323303222656, 213.17434692382812], + "extrinsics": { + "parent": "right_link6", + "pos": [-0.08, 0.0, 0.04], + "quat": [0.15304635, 0.69034543, -0.69034543, -0.15304635] + } + }, + { + "sensor_type": "Camera", + "uid": "cam_left_wrist", + "width": 640, + "height": 480, + "enable_mask": true, + "intrinsics": [488.1665344238281, 488.1665344238281, 322.7323303222656, 213.17434692382812], + "extrinsics": { + "parent": "left_link6", + "pos": [-0.08, 0.0, 0.04], + "quat": [0.15304635, 0.69034543, -0.69034543, -0.15304635] + } + } + ], + "light": { + "direct": [ + { + "uid": "light_1", + "light_type": "point", + "color": [1.0, 1.0, 1.0], + "intensity": 50.0, + "init_pos": [2, 0, 2], + "radius": 10.0 + } + ] + }, + "background": [ + { + "uid": "table", + "shape": { + "shape_type": "Mesh", + "fpath": "CircleTableSimple/circle_table_simple.ply" + }, + "attrs" : { + "mass": 10.0, + "static_friction": 0.95, + "dynamic_friction": 0.9, + "restitution": 0.01 + }, + "body_scale": [1, 1, 1], + "body_type": "kinematic", + "init_pos": [0.725, 0.0, 0.825], + "init_rot": [0, 90, 0] + } + ], + "rigid_object": [ + { + "uid":"block_1", + "shape": { + "shape_type": "Cube", + "size": [0.04, 0.04, 0.04] + }, + "attrs" : { + "mass": 0.05, + "static_friction": 0.5, + "dynamic_friction": 0.5, + "restitution": 0.0, + "contact_offset": 0.003, + "rest_offset": 0.001, + "max_depenetration_velocity": 1e1, + "min_position_iters": 32, + "min_velocity_iters": 8 + }, + "init_pos": [0.725, -0.015, 0.86], + "body_scale":[1, 1, 1], + "max_convex_hull_num": 1 + }, + { + "uid":"block_2", + "shape": { + "shape_type": "Cube", + "size": [0.04, 0.04, 0.04] + }, + "attrs" : { + "mass": 0.05, + "static_friction": 0.5, + "dynamic_friction": 0.5, + "restitution": 0.0, + "contact_offset": 0.003, + "rest_offset": 0.001, + "max_depenetration_velocity": 1e1, + "min_position_iters": 32, + "min_velocity_iters": 8 + }, + "init_pos": [0.725, -0.015, 0.86], + "body_scale":[1, 1, 1], + "max_convex_hull_num": 1 + }, + { + "uid":"block_3", + "shape": { + "shape_type": "Cube", + "size": [0.04, 0.04, 0.04] + }, + "attrs" : { + "mass": 0.05, + "static_friction": 0.5, + "dynamic_friction": 0.5, + "restitution": 0.0, + "contact_offset": 0.003, + "rest_offset": 0.001, + "max_depenetration_velocity": 1e1, + "min_position_iters": 32, + "min_velocity_iters": 8 + }, + "init_pos": [0.725, -0.015, 0.86], + "body_scale":[1, 1, 1], + "max_convex_hull_num": 1 + } + ] +} + diff --git a/configs/gym/blocks_ranking_size/cobot_magic_3cam.json b/configs/gym/blocks_ranking_size/cobot_magic_3cam.json new file mode 100644 index 00000000..dd628c40 --- /dev/null +++ b/configs/gym/blocks_ranking_size/cobot_magic_3cam.json @@ -0,0 +1,257 @@ +{ + "id": "BlocksRankingSize-v1", + "max_episodes": 5, + "env": { + "events": { + "random_light": { + "func": "randomize_light", + "mode": "interval", + "interval_step": 10, + "params": { + "entity_cfg": {"uid": "light_1"}, + "position_range": [[-0.5, -0.5, 2], [0.5, 0.5, 2]], + "color_range": [[0.6, 0.6, 0.6], [1, 1, 1]], + "intensity_range": [50.0, 100.0] + } + }, + "init_block_1_pose": { + "func": "randomize_rigid_object_pose", + "mode": "reset", + "params": { + "entity_cfg": {"uid": "block_1"}, + "position_range": [[0.445, -0.08, 0.86], [1.005, 0.05, 0.86]], + "rotation_range": [[0, 0, -43.0], [0, 0, 43.0]], + "relative_position": false + } + }, + "init_block_2_pose": { + "func": "randomize_rigid_object_pose", + "mode": "reset", + "params": { + "entity_cfg": {"uid": "block_2"}, + "position_range": [[0.445, -0.08, 0.86], [1.005, 0.05, 0.86]], + "rotation_range": [[0, 0, -43.0], [0, 0, 43.0]], + "relative_position": false + } + }, + "init_block_3_pose": { + "func": "randomize_rigid_object_pose", + "mode": "reset", + "params": { + "entity_cfg": {"uid": "block_3"}, + "position_range": [[0.445, -0.08, 0.86], [1.005, 0.05, 0.86]], + "rotation_range": [[0, 0, -43.0], [0, 0, 43.0]], + "relative_position": false + } + }, + "init_block_1_size": { + "func": "randomize_rigid_object_scale", + "mode": "reset", + "params": { + "entity_cfg": {"uid": "block_1"}, + "scale_factor_range": [[0.95238095], [1.04761905]], + "same_scale_all_axes": true + } + }, + "init_block_2_size": { + "func": "randomize_rigid_object_scale", + "mode": "reset", + "params": { + "entity_cfg": {"uid": "block_2"}, + "scale_factor_range": [[0.94117647], [1.05882353]], + "same_scale_all_axes": true + } + }, + "init_block_3_size": { + "func": "randomize_rigid_object_scale", + "mode": "reset", + "params": { + "entity_cfg": {"uid": "block_3"}, + "scale_factor_range": [[0.92307692], [1.07692308]], + "same_scale_all_axes": true + } + } + }, + "observations": { + "norm_robot_eef_joint": { + "func": "normalize_robot_joint_data", + "mode": "modify", + "name": "robot/qpos", + "params": { + "joint_ids": [12, 13, 14, 15] + } + }, + "block_1_pose": { + "func": "get_rigid_object_pose", + "mode": "add", + "name": "block_1_pose", + "params": { + "entity_cfg": {"uid": "block_1"} + } + }, + "block_2_pose": { + "func": "get_rigid_object_pose", + "mode": "add", + "name": "block_2_pose", + "params": { + "entity_cfg": {"uid": "block_2"} + } + }, + "block_3_pose": { + "func": "get_rigid_object_pose", + "mode": "add", + "name": "block_3_pose", + "params": { + "entity_cfg": {"uid": "block_3"} + } + } + } + }, + "robot": { + "uid": "CobotMagic", + "robot_type": "CobotMagic", + "init_pos": [0.0, 0.0, 0.7775], + "init_qpos": [-0.3,0.3,1.0,1.0,-1.2,-1.2,0.0,0.0,0.6,0.6,0.0,0.0,0.05,0.05,0.05,0.05] + }, + "sensor": [ + { + "sensor_type": "StereoCamera", + "uid": "cam_high", + "width": 960, + "height": 540, + "enable_mask": true, + "enable_depth": true, + "left_to_right_pos": [0.059684025824163614, 0, 0], + "intrinsics": [453.851402686215, 453.8347628855552, 469.827725021235, 258.6656181845155], + "intrinsics_right": [453.4536601653505, 453.3306024582175, 499.13697412367776, 297.7176248477935], + "extrinsics": { + "eye": [0.35368482807598, 0.014695524383058989, 1.4517046071614774], + "target": [0.7186357573287919, -0.054534732904795505, 0.5232553674540066], + "up": [0.9306678549330372, -0.0005600064212467153, 0.3658647703553347] + } + }, + { + "sensor_type": "Camera", + "uid": "cam_right_wrist", + "width": 640, + "height": 480, + "enable_mask": true, + "intrinsics": [488.1665344238281, 488.1665344238281, 322.7323303222656, 213.17434692382812], + "extrinsics": { + "parent": "right_link6", + "pos": [-0.08, 0.0, 0.04], + "quat": [0.15304635, 0.69034543, -0.69034543, -0.15304635] + } + }, + { + "sensor_type": "Camera", + "uid": "cam_left_wrist", + "width": 640, + "height": 480, + "enable_mask": true, + "intrinsics": [488.1665344238281, 488.1665344238281, 322.7323303222656, 213.17434692382812], + "extrinsics": { + "parent": "left_link6", + "pos": [-0.08, 0.0, 0.04], + "quat": [0.15304635, 0.69034543, -0.69034543, -0.15304635] + } + } + ], + "light": { + "direct": [ + { + "uid": "light_1", + "light_type": "point", + "color": [1.0, 1.0, 1.0], + "intensity": 50.0, + "init_pos": [2, 0, 2], + "radius": 10.0 + } + ] + }, + "background": [ + { + "uid": "table", + "shape": { + "shape_type": "Mesh", + "fpath": "CircleTableSimple/circle_table_simple.ply" + }, + "attrs" : { + "mass": 10.0, + "static_friction": 0.95, + "dynamic_friction": 0.9, + "restitution": 0.01 + }, + "body_scale": [1, 1, 1], + "body_type": "kinematic", + "init_pos": [0.725, 0.0, 0.825], + "init_rot": [0, 90, 0] + } + ], + "rigid_object": [ + { + "uid":"block_1", + "shape": { + "shape_type": "Cube", + "size": [0.063, 0.063, 0.063] + }, + "attrs" : { + "mass": 0.05, + "static_friction": 0.5, + "dynamic_friction": 0.5, + "restitution": 0.0, + "contact_offset": 0.003, + "rest_offset": 0.001, + "max_depenetration_velocity": 1e1, + "min_position_iters": 32, + "min_velocity_iters": 8 + }, + "init_pos": [0.725, -0.015, 0.86], + "body_scale":[1, 1, 1], + "max_convex_hull_num": 1 + }, + { + "uid":"block_2", + "shape": { + "shape_type": "Cube", + "size": [0.051, 0.051, 0.051] + }, + "attrs" : { + "mass": 0.05, + "static_friction": 0.5, + "dynamic_friction": 0.5, + "restitution": 0.0, + "contact_offset": 0.003, + "rest_offset": 0.001, + "max_depenetration_velocity": 1e1, + "min_position_iters": 32, + "min_velocity_iters": 8 + }, + "init_pos": [0.725, -0.015, 0.86], + "body_scale":[1, 1, 1], + "max_convex_hull_num": 1 + }, + { + "uid":"block_3", + "shape": { + "shape_type": "Cube", + "size": [0.039, 0.039, 0.039] + }, + "attrs" : { + "mass": 0.05, + "static_friction": 0.5, + "dynamic_friction": 0.5, + "restitution": 0.0, + "contact_offset": 0.003, + "rest_offset": 0.001, + "max_depenetration_velocity": 1e1, + "min_position_iters": 32, + "min_velocity_iters": 8 + }, + "init_pos": [0.725, -0.015, 0.86], + "body_scale":[1, 1, 1], + "max_convex_hull_num": 1 + } + ] +} + diff --git a/configs/gym/match_object_container/cobot_magic_3cam.json b/configs/gym/match_object_container/cobot_magic_3cam.json new file mode 100644 index 00000000..9463c70a --- /dev/null +++ b/configs/gym/match_object_container/cobot_magic_3cam.json @@ -0,0 +1,343 @@ +{ + "id": "MatchObjectContainer-v1", + "max_episodes": 5, + "env": { + "events": { + "random_light": { + "func": "randomize_light", + "mode": "interval", + "interval_step": 10, + "params": { + "entity_cfg": {"uid": "light_1"}, + "position_range": [[-0.5, -0.5, 2], [0.5, 0.5, 2]], + "color_range": [[0.6, 0.6, 0.6], [1, 1, 1]], + "intensity_range": [50.0, 100.0] + } + }, + "init_block_cube_1_pose": { + "func": "randomize_rigid_object_pose", + "mode": "reset", + "params": { + "entity_cfg": {"uid": "block_cube_1"}, + "position_range": [[0.55, -0.10, 0.86], [0.58, -0.05, 0.86]], + "rotation_range": [[0, 0, -0.52], [0, 0, 0.52]], + "relative_position": false + } + }, + "init_block_sphere_1_pose": { + "func": "randomize_rigid_object_pose", + "mode": "reset", + "params": { + "entity_cfg": {"uid": "block_sphere_1"}, + "position_range": [[0.62, -0.10, 0.86], [0.65, -0.05, 0.86]], + "rotation_range": [[0, 0, -0.52], [0, 0, 0.52]], + "relative_position": false + } + }, + "init_block_cube_2_pose": { + "func": "randomize_rigid_object_pose", + "mode": "reset", + "params": { + "entity_cfg": {"uid": "block_cube_2"}, + "position_range": [[0.55, 0.05, 0.86], [0.58, 0.10, 0.86]], + "rotation_range": [[0, 0, -0.52], [0, 0, 0.52]], + "relative_position": false + } + }, + "init_block_sphere_2_pose": { + "func": "randomize_rigid_object_pose", + "mode": "reset", + "params": { + "entity_cfg": {"uid": "block_sphere_2"}, + "position_range": [[0.62, 0.05, 0.86], [0.65, 0.10, 0.86]], + "rotation_range": [[0, 0, -0.52], [0, 0, 0.52]], + "relative_position": false + } + } + }, + "observations": { + "norm_robot_eef_joint": { + "func": "normalize_robot_joint_data", + "mode": "modify", + "name": "robot/qpos", + "params": { + "joint_ids": [12, 13, 14, 15] + } + }, + "block_cube_1_pose": { + "func": "get_rigid_object_pose", + "mode": "add", + "name": "block_cube_1_pose", + "params": { + "entity_cfg": {"uid": "block_cube_1"} + } + }, + "block_sphere_1_pose": { + "func": "get_rigid_object_pose", + "mode": "add", + "name": "block_sphere_1_pose", + "params": { + "entity_cfg": {"uid": "block_sphere_1"} + } + }, + "container_cube_pose": { + "func": "get_rigid_object_pose", + "mode": "add", + "name": "container_cube_pose", + "params": { + "entity_cfg": {"uid": "container_cube"} + } + }, + "container_sphere_pose": { + "func": "get_rigid_object_pose", + "mode": "add", + "name": "container_sphere_pose", + "params": { + "entity_cfg": {"uid": "container_sphere"} + } + }, + "block_cube_2_pose": { + "func": "get_rigid_object_pose", + "mode": "add", + "name": "block_cube_2_pose", + "params": { + "entity_cfg": {"uid": "block_cube_2"} + } + }, + "block_sphere_2_pose": { + "func": "get_rigid_object_pose", + "mode": "add", + "name": "block_sphere_2_pose", + "params": { + "entity_cfg": {"uid": "block_sphere_2"} + } + } + } + }, + "robot": { + "uid": "CobotMagic", + "robot_type": "CobotMagic", + "init_pos": [0.0, 0.0, 0.7775], + "init_qpos": [-0.3,0.3,1.0,1.0,-1.2,-1.2,0.0,0.0,0.6,0.6,0.0,0.0,0.05,0.05,0.05,0.05] + }, + "sensor": [ + { + "sensor_type": "StereoCamera", + "uid": "cam_high", + "width": 960, + "height": 540, + "enable_mask": true, + "enable_depth": true, + "left_to_right_pos": [0.059684025824163614, 0, 0], + "intrinsics": [453.851402686215, 453.8347628855552, 469.827725021235, 258.6656181845155], + "intrinsics_right": [453.4536601653505, 453.3306024582175, 499.13697412367776, 297.7176248477935], + "extrinsics": { + "eye": [0.35368482807598, 0.014695524383058989, 1.4517046071614774], + "target": [0.7186357573287919, -0.054534732904795505, 0.5232553674540066], + "up": [0.9306678549330372, -0.0005600064212467153, 0.3658647703553347] + } + }, + { + "sensor_type": "Camera", + "uid": "cam_right_wrist", + "width": 640, + "height": 480, + "enable_mask": true, + "intrinsics": [488.1665344238281, 488.1665344238281, 322.7323303222656, 213.17434692382812], + "extrinsics": { + "parent": "right_link6", + "pos": [-0.08, 0.0, 0.04], + "quat": [0.15304635, 0.69034543, -0.69034543, -0.15304635] + } + }, + { + "sensor_type": "Camera", + "uid": "cam_left_wrist", + "width": 640, + "height": 480, + "enable_mask": true, + "intrinsics": [488.1665344238281, 488.1665344238281, 322.7323303222656, 213.17434692382812], + "extrinsics": { + "parent": "left_link6", + "pos": [-0.08, 0.0, 0.04], + "quat": [0.15304635, 0.69034543, -0.69034543, -0.15304635] + } + } + ], + "light": { + "direct": [ + { + "uid": "light_1", + "light_type": "point", + "color": [1.0, 1.0, 1.0], + "intensity": 50.0, + "init_pos": [2, 0, 2], + "radius": 10.0 + } + ] + }, + "background": [ + { + "uid": "table", + "shape": { + "shape_type": "Mesh", + "fpath": "CircleTableSimple/circle_table_simple.ply" + }, + "attrs" : { + "mass": 10.0, + "static_friction": 0.95, + "dynamic_friction": 0.9, + "restitution": 0.01 + }, + "body_scale": [1, 1, 1], + "body_type": "kinematic", + "init_pos": [0.725, 0.0, 0.825], + "init_rot": [0, 90, 0] + } + ], + "rigid_object": [ + { + "uid":"block_cube_1", + "shape": { + "shape_type": "Cube", + "size": [0.04, 0.04, 0.04] + }, + "attrs" : { + "mass": 0.05, + "static_friction": 0.5, + "dynamic_friction": 0.5, + "restitution": 0.0, + "contact_offset": 0.003, + "rest_offset": 0.001, + "max_depenetration_velocity": 1e1, + "min_position_iters": 32, + "min_velocity_iters": 8 + }, + "init_pos": [0.565, -0.075, 0.86], + "init_rot": [0, 0, 0], + "body_scale":[1, 1, 1], + "max_convex_hull_num": 1 + }, + { + "uid":"block_sphere_1", + "shape": { + "shape_type": "Sphere", + "radius": 0.025 + }, + "attrs" : { + "mass": 0.05, + "static_friction": 0.5, + "dynamic_friction": 0.5, + "restitution": 0.0, + "contact_offset": 0.003, + "rest_offset": 0.001, + "max_depenetration_velocity": 1e1, + "min_position_iters": 32, + "min_velocity_iters": 8 + }, + "init_pos": [0.635, -0.075, 0.86], + "init_rot": [0, 0, 0], + "body_scale":[1, 1, 1], + "max_convex_hull_num": 1 + }, + { + "uid":"container_cube", + "shape": { + "shape_type": "Mesh", + "fpath": "ContainerMetal/container_metal.obj" + }, + "body_type": "dynamic", + "attrs" : { + "mass": 0.5, + "static_friction": 1.0, + "dynamic_friction": 1.0, + "restitution": 0.0, + "contact_offset": 0.003, + "rest_offset": 0.001, + "max_depenetration_velocity": 2.0, + "linear_damping": 2.0, + "angular_damping": 2.0, + "max_linear_velocity": 5.0, + "max_angular_velocity": 10.0, + "min_position_iters": 32, + "min_velocity_iters": 8 + }, + "init_pos": [0.875, -0.25, 0.86], + "init_rot": [0, 0, 0], + "body_scale":[1, 1, 1], + "max_convex_hull_num": 8 + }, + { + "uid":"container_sphere", + "shape": { + "shape_type": "Mesh", + "fpath": "ContainerMetal/container_metal.obj" + }, + "body_type": "dynamic", + "attrs" : { + "mass": 0.5, + "static_friction": 1.0, + "dynamic_friction": 1.0, + "restitution": 0.0, + "contact_offset": 0.003, + "rest_offset": 0.001, + "max_depenetration_velocity": 2.0, + "linear_damping": 2.0, + "angular_damping": 2.0, + "max_linear_velocity": 5.0, + "max_angular_velocity": 10.0, + "min_position_iters": 32, + "min_velocity_iters": 8 + }, + "init_pos": [0.875, 0.25, 0.86], + "init_rot": [0, 0, 0], + "body_scale":[1, 1, 1], + "max_convex_hull_num": 8 + }, + { + "uid":"block_cube_2", + "shape": { + "shape_type": "Cube", + "size": [0.04, 0.04, 0.04] + }, + "attrs" : { + "mass": 0.05, + "static_friction": 0.5, + "dynamic_friction": 0.5, + "restitution": 0.0, + "contact_offset": 0.003, + "rest_offset": 0.001, + "max_depenetration_velocity": 1e1, + "min_position_iters": 32, + "min_velocity_iters": 8 + }, + "init_pos": [0.565, 0.075, 0.86], + "init_rot": [0, 0, 0], + "body_scale":[1, 1, 1], + "max_convex_hull_num": 1 + }, + { + "uid":"block_sphere_2", + "shape": { + "shape_type": "Sphere", + "radius": 0.025 + }, + "attrs" : { + "mass": 0.05, + "static_friction": 0.5, + "dynamic_friction": 0.5, + "restitution": 0.0, + "contact_offset": 0.003, + "rest_offset": 0.001, + "max_depenetration_velocity": 1e1, + "min_position_iters": 32, + "min_velocity_iters": 8 + }, + "init_pos": [0.635, 0.075, 0.86], + "init_rot": [0, 0, 0], + "body_scale":[1, 1, 1], + "max_convex_hull_num": 1 + } + ] +} + diff --git a/configs/gym/place_object_drawer/cobot_magic_3cam.json b/configs/gym/place_object_drawer/cobot_magic_3cam.json new file mode 100644 index 00000000..185e4617 --- /dev/null +++ b/configs/gym/place_object_drawer/cobot_magic_3cam.json @@ -0,0 +1,165 @@ +{ + "id": "PlaceObjectDrawer-v1", + "max_episodes": 5, + "env": { + "events": { + "random_light": { + "func": "randomize_light", + "mode": "interval", + "interval_step": 10, + "params": { + "entity_cfg": {"uid": "light_1"}, + "position_range": [[-0.5, -0.5, 2], [0.5, 0.5, 2]], + "color_range": [[0.6, 0.6, 0.6], [1, 1, 1]], + "intensity_range": [50.0, 100.0] + } + }, + "init_object_pose": { + "func": "randomize_rigid_object_pose", + "mode": "reset", + "params": { + "entity_cfg": {"uid": "object"}, + "position_range": [[0.66, -0.15, 0.86], [0.70, -0.10, 0.86]], + "rotation_range": [[0, 0, -0.52], [0, 0, 0.52]], + "relative_position": false + } + } + }, + "observations": { + "norm_robot_eef_joint": { + "func": "normalize_robot_joint_data", + "mode": "modify", + "name": "robot/qpos", + "params": { + "joint_ids": [12, 13, 14, 15] + } + }, + "object_pose": { + "func": "get_rigid_object_pose", + "mode": "add", + "name": "object_pose", + "params": { + "entity_cfg": {"uid": "object"} + } + } + } + }, + "robot": { + "uid": "CobotMagic", + "robot_type": "CobotMagic", + "init_pos": [0.0, 0.0, 0.7775], + "init_qpos": [-0.3,0.3,1.0,1.0,-1.2,-1.2,0.0,0.0,0.6,0.6,0.0,0.0,0.05,0.05,0.05,0.05] + }, + "sensor": [ + { + "sensor_type": "StereoCamera", + "uid": "cam_high", + "width": 960, + "height": 540, + "enable_mask": true, + "enable_depth": true, + "left_to_right_pos": [0.059684025824163614, 0, 0], + "intrinsics": [453.851402686215, 453.8347628855552, 469.827725021235, 258.6656181845155], + "intrinsics_right": [453.4536601653505, 453.3306024582175, 499.13697412367776, 297.7176248477935], + "extrinsics": { + "eye": [0.35368482807598, 0.014695524383058989, 1.4517046071614774], + "target": [0.7186357573287919, -0.054534732904795505, 0.5232553674540066], + "up": [0.9306678549330372, -0.0005600064212467153, 0.3658647703553347] + } + }, + { + "sensor_type": "Camera", + "uid": "cam_right_wrist", + "width": 640, + "height": 480, + "enable_mask": true, + "intrinsics": [488.1665344238281, 488.1665344238281, 322.7323303222656, 213.17434692382812], + "extrinsics": { + "parent": "right_link6", + "pos": [-0.08, 0.0, 0.04], + "quat": [0.15304635, 0.69034543, -0.69034543, -0.15304635] + } + }, + { + "sensor_type": "Camera", + "uid": "cam_left_wrist", + "width": 640, + "height": 480, + "enable_mask": true, + "intrinsics": [488.1665344238281, 488.1665344238281, 322.7323303222656, 213.17434692382812], + "extrinsics": { + "parent": "left_link6", + "pos": [-0.08, 0.0, 0.04], + "quat": [0.15304635, 0.69034543, -0.69034543, -0.15304635] + } + } + ], + "light": { + "direct": [ + { + "uid": "light_1", + "light_type": "point", + "color": [1.0, 1.0, 1.0], + "intensity": 50.0, + "init_pos": [2, 0, 2], + "radius": 10.0 + } + ] + }, + "background": [ + { + "uid": "table", + "shape": { + "shape_type": "Mesh", + "fpath": "CircleTableSimple/circle_table_simple.ply" + }, + "attrs" : { + "mass": 10.0, + "static_friction": 0.95, + "dynamic_friction": 0.9, + "restitution": 0.01 + }, + "body_scale": [1, 1, 1], + "body_type": "kinematic", + "init_pos": [0.725, 0.0, 0.825], + "init_rot": [0, 90, 0] + } + ], + "rigid_object": [ + { + "uid":"object", + "shape": { + "shape_type": "Mesh", + "fpath": "ToyDuck/toy_duck.glb" + }, + "attrs" : { + "mass": 0.01, + "static_friction": 1.0, + "dynamic_friction": 1.0, + "restitution": 0.0, + "contact_offset": 0.003, + "rest_offset": 0.001, + "max_depenetration_velocity": 2.0, + "linear_damping": 2.0, + "angular_damping": 2.0, + "max_linear_velocity": 5.0, + "max_angular_velocity": 10.0, + "min_position_iters": 32, + "min_velocity_iters": 8 + }, + "init_pos": [0.725, -0.1, 0.86], + "init_rot": [0, 0, 0], + "body_scale":[0.2, 0.2, 0.2], + "max_convex_hull_num": 8 + } + ], + "articulation": [ + { + "uid": "drawer", + "fpath": "SimpleBoxDrawer/simple_box_drawer/simple_box_drawer.urdf", + "init_pos": [0.725, 0.16, 1.025], + "init_rot": [0, 180, 0] + } + ] +} + diff --git a/configs/gym/stack_blocks_two/cobot_magic_3cam.json b/configs/gym/stack_blocks_two/cobot_magic_3cam.json new file mode 100644 index 00000000..6f160b58 --- /dev/null +++ b/configs/gym/stack_blocks_two/cobot_magic_3cam.json @@ -0,0 +1,191 @@ +{ + "id": "StackBlocksTwo-v1", + "max_episodes": 5, + "env": { + "events": { + "random_light": { + "func": "randomize_light", + "mode": "interval", + "interval_step": 10, + "params": { + "entity_cfg": {"uid": "light_1"}, + "position_range": [[-0.5, -0.5, 2], [0.5, 0.5, 2]], + "color_range": [[0.6, 0.6, 0.6], [1, 1, 1]], + "intensity_range": [50.0, 100.0] + } + }, + "init_block_1_pose": { + "func": "randomize_rigid_object_pose", + "mode": "reset", + "params": { + "entity_cfg": {"uid": "block_1"}, + "position_range": [[-0.08, -0.08, 0.0], [0.08, 0.08, 0.0]], + "rotation_range": [[0, 0, -0.75], [0, 0, 0.75]], + "relative_position": true + } + }, + "init_block_2_pose": { + "func": "randomize_rigid_object_pose", + "mode": "reset", + "params": { + "entity_cfg": {"uid": "block_2"}, + "position_range": [[-0.08, -0.08, 0.0], [0.08, 0.08, 0.0]], + "rotation_range": [[0, 0, -0.75], [0, 0, 0.75]], + "relative_position": true + } + } + }, + "observations": { + "norm_robot_eef_joint": { + "func": "normalize_robot_joint_data", + "mode": "modify", + "name": "robot/qpos", + "params": { + "joint_ids": [12, 13, 14, 15] + } + }, + "block_1_pose": { + "func": "get_rigid_object_pose", + "mode": "add", + "name": "block_1_pose", + "params": { + "entity_cfg": {"uid": "block_1"} + } + }, + "block_2_pose": { + "func": "get_rigid_object_pose", + "mode": "add", + "name": "block_2_pose", + "params": { + "entity_cfg": {"uid": "block_2"} + } + } + } + }, + "robot": { + "uid": "CobotMagic", + "robot_type": "CobotMagic", + "init_pos": [0.0, 0.0, 0.7775], + "init_qpos": [-0.3,0.3,1.0,1.0,-1.2,-1.2,0.0,0.0,0.6,0.6,0.0,0.0,0.05,0.05,0.05,0.05] + }, + "sensor": [ + { + "sensor_type": "StereoCamera", + "uid": "cam_high", + "width": 960, + "height": 540, + "enable_mask": true, + "enable_depth": true, + "left_to_right_pos": [0.059684025824163614, 0, 0], + "intrinsics": [453.851402686215, 453.8347628855552, 469.827725021235, 258.6656181845155], + "intrinsics_right": [453.4536601653505, 453.3306024582175, 499.13697412367776, 297.7176248477935], + "extrinsics": { + "eye": [0.35368482807598, 0.014695524383058989, 1.4517046071614774], + "target": [0.7186357573287919, -0.054534732904795505, 0.5232553674540066], + "up": [0.9306678549330372, -0.0005600064212467153, 0.3658647703553347] + } + }, + { + "sensor_type": "Camera", + "uid": "cam_right_wrist", + "width": 640, + "height": 480, + "enable_mask": true, + "intrinsics": [488.1665344238281, 488.1665344238281, 322.7323303222656, 213.17434692382812], + "extrinsics": { + "parent": "right_link6", + "pos": [-0.08, 0.0, 0.04], + "quat": [0.15304635, 0.69034543, -0.69034543, -0.15304635] + } + }, + { + "sensor_type": "Camera", + "uid": "cam_left_wrist", + "width": 640, + "height": 480, + "enable_mask": true, + "intrinsics": [488.1665344238281, 488.1665344238281, 322.7323303222656, 213.17434692382812], + "extrinsics": { + "parent": "left_link6", + "pos": [-0.08, 0.0, 0.04], + "quat": [0.15304635, 0.69034543, -0.69034543, -0.15304635] + } + } + ], + "light": { + "direct": [ + { + "uid": "light_1", + "light_type": "point", + "color": [1.0, 1.0, 1.0], + "intensity": 50.0, + "init_pos": [2, 0, 2], + "radius": 10.0 + } + ] + }, + "background": [ + { + "uid": "table", + "shape": { + "shape_type": "Mesh", + "fpath": "CircleTableSimple/circle_table_simple.ply" + }, + "attrs" : { + "mass": 10.0, + "static_friction": 0.95, + "dynamic_friction": 0.9, + "restitution": 0.01 + }, + "body_scale": [1, 1, 1], + "body_type": "kinematic", + "init_pos": [0.725, 0.0, 0.825], + "init_rot": [0, 90, 0] + } + ], + "rigid_object": [ + { + "uid":"block_1", + "shape": { + "shape_type": "Cube", + "size": [0.05, 0.05, 0.05] + }, + "attrs" : { + "mass": 0.05, + "static_friction": 0.5, + "dynamic_friction": 0.5, + "restitution": 0.0, + "contact_offset": 0.003, + "rest_offset": 0.001, + "max_depenetration_velocity": 1e1, + "min_position_iters": 32, + "min_velocity_iters": 8 + }, + "init_pos": [0.75, -0.1, 0.9], + "body_scale":[1, 1, 1], + "max_convex_hull_num": 1 + }, + { + "uid":"block_2", + "shape": { + "shape_type": "Cube", + "size": [0.05, 0.05, 0.05] + }, + "attrs" : { + "mass": 0.05, + "static_friction": 0.5, + "dynamic_friction": 0.5, + "restitution": 0.0, + "contact_offset": 0.003, + "rest_offset": 0.001, + "max_depenetration_velocity": 1e1, + "min_position_iters": 32, + "min_velocity_iters": 8 + }, + "init_pos": [0.75, 0.1, 0.9], + "body_scale":[1, 1, 1], + "max_convex_hull_num": 1 + } + ] +} + diff --git a/configs/gym/stack_cups/cobot_magic_3cam.json b/configs/gym/stack_cups/cobot_magic_3cam.json new file mode 100644 index 00000000..bd4de01b --- /dev/null +++ b/configs/gym/stack_cups/cobot_magic_3cam.json @@ -0,0 +1,202 @@ +{ + "id": "StackCups-v1", + "max_episodes": 5, + "env": { + "events": { + "random_light": { + "func": "randomize_light", + "mode": "interval", + "interval_step": 10, + "params": { + "entity_cfg": {"uid": "light_1"}, + "position_range": [[-0.5, -0.5, 2], [0.5, 0.5, 2]], + "color_range": [[0.6, 0.6, 0.6], [1, 1, 1]], + "intensity_range": [50.0, 100.0] + } + }, + "init_cup_1_pose": { + "func": "randomize_rigid_object_pose", + "mode": "reset", + "params": { + "entity_cfg": {"uid": "cup_1"}, + "position_range": [[0.66, -0.15, 0.86], [0.74, -0.05, 0.86]], + "rotation_range": [[0, 0, -0.52], [0, 0, 0.52]], + "relative_position": false + } + }, + "init_cup_2_pose": { + "func": "randomize_rigid_object_pose", + "mode": "reset", + "params": { + "entity_cfg": {"uid": "cup_2"}, + "position_range": [[0.76, -0.15, 0.86], [0.84, -0.05, 0.86]], + "rotation_range": [[0, 0, -0.52], [0, 0, 0.52]], + "relative_position": false + } + } + }, + "observations": { + "norm_robot_eef_joint": { + "func": "normalize_robot_joint_data", + "mode": "modify", + "name": "robot/qpos", + "params": { + "joint_ids": [12, 13, 14, 15] + } + }, + "cup_1_pose": { + "func": "get_rigid_object_pose", + "mode": "add", + "name": "cup_1_pose", + "params": { + "entity_cfg": {"uid": "cup_1"} + } + }, + "cup_2_pose": { + "func": "get_rigid_object_pose", + "mode": "add", + "name": "cup_2_pose", + "params": { + "entity_cfg": {"uid": "cup_2"} + } + } + } + }, + "robot": { + "uid": "CobotMagic", + "robot_type": "CobotMagic", + "init_pos": [0.0, 0.0, 0.7775], + "init_qpos": [-0.3,0.3,1.0,1.0,-1.2,-1.2,0.0,0.0,0.6,0.6,0.0,0.0,0.05,0.05,0.05,0.05] + }, + "sensor": [ + { + "sensor_type": "StereoCamera", + "uid": "cam_high", + "width": 960, + "height": 540, + "enable_mask": true, + "enable_depth": true, + "left_to_right_pos": [0.059684025824163614, 0, 0], + "intrinsics": [453.851402686215, 453.8347628855552, 469.827725021235, 258.6656181845155], + "intrinsics_right": [453.4536601653505, 453.3306024582175, 499.13697412367776, 297.7176248477935], + "extrinsics": { + "eye": [0.35368482807598, 0.014695524383058989, 1.4517046071614774], + "target": [0.7186357573287919, -0.054534732904795505, 0.5232553674540066], + "up": [0.9306678549330372, -0.0005600064212467153, 0.3658647703553347] + } + }, + { + "sensor_type": "Camera", + "uid": "cam_right_wrist", + "width": 640, + "height": 480, + "enable_mask": true, + "intrinsics": [488.1665344238281, 488.1665344238281, 322.7323303222656, 213.17434692382812], + "extrinsics": { + "parent": "right_link6", + "pos": [-0.08, 0.0, 0.04], + "quat": [0.15304635, 0.69034543, -0.69034543, -0.15304635] + } + }, + { + "sensor_type": "Camera", + "uid": "cam_left_wrist", + "width": 640, + "height": 480, + "enable_mask": true, + "intrinsics": [488.1665344238281, 488.1665344238281, 322.7323303222656, 213.17434692382812], + "extrinsics": { + "parent": "left_link6", + "pos": [-0.08, 0.0, 0.04], + "quat": [0.15304635, 0.69034543, -0.69034543, -0.15304635] + } + } + ], + "light": { + "direct": [ + { + "uid": "light_1", + "light_type": "point", + "color": [1.0, 1.0, 1.0], + "intensity": 50.0, + "init_pos": [2, 0, 2], + "radius": 10.0 + } + ] + }, + "background": [ + { + "uid": "table", + "shape": { + "shape_type": "Mesh", + "fpath": "CircleTableSimple/circle_table_simple.ply" + }, + "attrs" : { + "mass": 10.0, + "static_friction": 0.95, + "dynamic_friction": 0.9, + "restitution": 0.01 + }, + "body_scale": [1, 1, 1], + "body_type": "kinematic", + "init_pos": [0.725, 0.0, 0.825], + "init_rot": [0, 90, 0] + } + ], + "rigid_object": [ + { + "uid":"cup_1", + "shape": { + "shape_type": "Mesh", + "fpath": "PaperCup/paper_cup.ply" + }, + "attrs" : { + "mass": 0.01, + "static_friction": 1.0, + "dynamic_friction": 1.0, + "restitution": 0.0, + "contact_offset": 0.003, + "rest_offset": 0.001, + "max_depenetration_velocity": 2.0, + "linear_damping": 2.0, + "angular_damping": 2.0, + "max_linear_velocity": 5.0, + "max_angular_velocity": 10.0, + "min_position_iters": 32, + "min_velocity_iters": 8 + }, + "init_pos": [0.70, -0.1, 0.86], + "init_rot": [0, 0, 0], + "body_scale":[0.75, 0.75, 1.0], + "max_convex_hull_num": 8 + }, + { + "uid":"cup_2", + "shape": { + "shape_type": "Mesh", + "fpath": "PaperCup/paper_cup.ply" + }, + "attrs" : { + "mass": 0.01, + "static_friction": 1.0, + "dynamic_friction": 1.0, + "restitution": 0.0, + "contact_offset": 0.003, + "rest_offset": 0.001, + "max_depenetration_velocity": 2.0, + "linear_damping": 2.0, + "angular_damping": 2.0, + "max_linear_velocity": 5.0, + "max_angular_velocity": 10.0, + "min_position_iters": 32, + "min_velocity_iters": 8 + }, + "init_pos": [0.80, -0.1, 0.86], + "init_rot": [0, 0, 0], + "body_scale":[0.75, 0.75, 1.0], + "max_convex_hull_num": 8 + } + ] +} + + diff --git a/embodichain/lab/gym/envs/__init__.py b/embodichain/lab/gym/envs/__init__.py index 88257690..64c48c24 100644 --- a/embodichain/lab/gym/envs/__init__.py +++ b/embodichain/lab/gym/envs/__init__.py @@ -26,6 +26,22 @@ PourWaterEnv, ) from embodichain.lab.gym.envs.tasks.tableware.scoop_ice import ScoopIce +from embodichain.lab.gym.envs.tasks.tableware.stack_blocks_two import StackBlocksTwoEnv +from embodichain.lab.gym.envs.tasks.tableware.blocks_ranking_rgb import ( + BlocksRankingRGBEnv, +) +from embodichain.lab.gym.envs.tasks.tableware.blocks_ranking_size import ( + BlocksRankingSizeEnv, +) +from embodichain.lab.gym.envs.tasks.tableware.place_object_drawer import ( + PlaceObjectDrawerEnv, +) +from embodichain.lab.gym.envs.tasks.tableware.stack_cups import ( + StackCupsEnv, +) +from embodichain.lab.gym.envs.tasks.tableware.match_object_container import ( + MatchObjectContainerEnv, +) # Reinforcement learning environments from embodichain.lab.gym.envs.tasks.rl.push_cube import PushCubeEnv diff --git a/embodichain/lab/gym/envs/managers/dataset_manager.py b/embodichain/lab/gym/envs/managers/dataset_manager.py index a0ca1688..affdff9a 100644 --- a/embodichain/lab/gym/envs/managers/dataset_manager.py +++ b/embodichain/lab/gym/envs/managers/dataset_manager.py @@ -87,8 +87,11 @@ def __init__(self, cfg: object, env: EmbodiedEnv): super().__init__(cfg, env) ## TODO: fix configurable_action.py to avoid getting env.metadata['dataset'] - # Extract robot_meta from first functor and add to env.metadata for backward compatibility + # Extract robot_meta from functor params or plain config and add to env.metadata for backward compatibility # This allows legacy code (like action_bank) to access robot_meta via env.metadata["dataset"]["robot_meta"] + robot_meta_found = False + + # First, try to extract from functor params for mode_cfgs in self._mode_functor_cfgs.values(): for functor_cfg in mode_cfgs: if "robot_meta" in functor_cfg.params: @@ -102,10 +105,30 @@ def __init__(self, cfg: object, env: EmbodiedEnv): logger.log_info( "Added robot_meta to env.metadata for backward compatibility" ) + robot_meta_found = True break + if robot_meta_found: + break + + # If not found in functor params, try to extract from plain config + if not robot_meta_found: + # Check if config is dict or object + if isinstance(self.cfg, dict): + cfg_items = self.cfg.items() else: - continue - break + cfg_items = self.cfg.__dict__.items() + + for config_name, config_value in cfg_items: + if config_name == "robot_meta" and isinstance(config_value, dict): + if not hasattr(env, "metadata"): + env.metadata = {} + if "dataset" not in env.metadata: + env.metadata["dataset"] = {} + env.metadata["dataset"]["robot_meta"] = config_value + logger.log_info( + "Added robot_meta to env.metadata for backward compatibility (from plain config)" + ) + break logger.log_info( f"DatasetManager initialized with {sum(len(v) for v in self._mode_functor_names.values())} functors" @@ -321,6 +344,12 @@ def _prepare_functors(self): if functor_cfg is None: continue + # Skip non-functor configurations (e.g., robot_meta which is a plain dict) + # Functor configurations must have a "func" field + if isinstance(functor_cfg, dict) and "func" not in functor_cfg: + # This is a plain configuration (not a functor), skip it + continue + # Convert dict to DatasetFunctorCfg if needed (for JSON configs) if isinstance(functor_cfg, dict): functor_cfg = DatasetFunctorCfg(**functor_cfg) diff --git a/embodichain/lab/gym/envs/managers/manager_base.py b/embodichain/lab/gym/envs/managers/manager_base.py index d89503bc..2db673be 100644 --- a/embodichain/lab/gym/envs/managers/manager_base.py +++ b/embodichain/lab/gym/envs/managers/manager_base.py @@ -379,13 +379,17 @@ def _process_functor_cfg_at_play(self, functor_name: str, functor_cfg: FunctorCf functor_name: The name of the functor. functor_cfg: The functor configuration. """ - for key, value in functor_cfg.params.items(): + + def _resolve_scene_entities(value, key_path: str): + # Resolve SceneEntityCfg anywhere in nested params (dict/list/tuple). if isinstance(value, SceneEntityCfg): # load the entity try: value.resolve(self._env.sim) except ValueError as e: - raise ValueError(f"Error while parsing '{functor_name}:{key}'. {e}") + raise ValueError( + f"Error while parsing '{functor_name}:{key_path}'. {e}" + ) # log the entity for checking later msg = f"[{functor_cfg.__class__.__name__}:{functor_name}] Found entity '{value.uid}'." if value.joint_ids is not None: @@ -394,8 +398,27 @@ def _process_functor_cfg_at_play(self, functor_name: str, functor_cfg: FunctorCf msg += f"\n\tBody names: {value.body_names} [{value.body_ids}]" # print the information print(f"[INFO]: {msg}") + return value + if isinstance(value, list): # recursively resolve the list + return [ + _resolve_scene_entities(v, f"{key_path}[{i}]") + for i, v in enumerate(value) + ] + if isinstance(value, tuple): # recursively resolve the tuple + return tuple( + _resolve_scene_entities(v, f"{key_path}[{i}]") + for i, v in enumerate(value) + ) + if isinstance(value, dict): # recursively resolve the dict + return { + k: _resolve_scene_entities(v, f"{key_path}.{k}") + for k, v in value.items() + } + return value + + for key, value in list(functor_cfg.params.items()): # store the entity - functor_cfg.params[key] = value + functor_cfg.params[key] = _resolve_scene_entities(value, key) # initialize the functor if it is a class if inspect.isclass(functor_cfg.func): diff --git a/embodichain/lab/gym/envs/managers/randomization/__init__.py b/embodichain/lab/gym/envs/managers/randomization/__init__.py index 7a965376..16520976 100644 --- a/embodichain/lab/gym/envs/managers/randomization/__init__.py +++ b/embodichain/lab/gym/envs/managers/randomization/__init__.py @@ -17,6 +17,7 @@ from .physics import * # noqa: F401, F403 from .visual import * # noqa: F401, F403 from .spatial import * # noqa: F401, F403 +from .scale import * # noqa: F401, F403 """ Randomization are all implemented as Event functors. diff --git a/embodichain/lab/gym/envs/managers/randomization/geometry.py b/embodichain/lab/gym/envs/managers/randomization/geometry.py new file mode 100644 index 00000000..5e9ee874 --- /dev/null +++ b/embodichain/lab/gym/envs/managers/randomization/geometry.py @@ -0,0 +1,182 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2025 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +from typing import TYPE_CHECKING, List, Union + +import torch + +from embodichain.lab.gym.envs.managers.cfg import SceneEntityCfg +from embodichain.lab.sim.objects import RigidObject +from embodichain.utils import logger +from embodichain.utils.math import sample_uniform + +if TYPE_CHECKING: + from embodichain.lab.gym.envs import EmbodiedEnv + + +def _normalize_env_ids( + env: EmbodiedEnv, env_ids: Union[torch.Tensor, None] +) -> torch.Tensor: + # Target all active environments if no specific IDs are provided + if env_ids is None: + return torch.arange(env.num_envs, device=env.device) + return env_ids + + +def _sample_body_scale( + env: EmbodiedEnv, + env_ids: torch.Tensor, + scale_factor_range: tuple[list[float], list[float]], + same_scale_all_axes: bool, +) -> torch.Tensor: + """Sample per-env body scale factors. + + Returns: + torch.Tensor: Shape (num_envs_selected, 3) scale factors for x/y/z. + """ + num_instance = len(env_ids) + if same_scale_all_axes: + low = torch.tensor(scale_factor_range[0][0], device=env.device) + high = torch.tensor(scale_factor_range[1][0], device=env.device) + s = sample_uniform(lower=low, upper=high, size=(num_instance,)) + return torch.stack([s, s, s], dim=1) + low = torch.tensor(scale_factor_range[0], device=env.device) + high = torch.tensor(scale_factor_range[1], device=env.device) + return sample_uniform(lower=low, upper=high, size=(num_instance, 3)) + + +def randomize_rigid_object_scale( + env: EmbodiedEnv, + env_ids: Union[torch.Tensor, None], + entity_cfg: SceneEntityCfg, + scale_factor_range: tuple[list[float], list[float]] | None = None, + same_scale_all_axes: bool = True, +) -> None: + """Randomize a rigid object's *body scale factors* (multiplicative, not absolute size). + + Args: + env: Environment instance. + env_ids: Target env ids. If None, applies to all envs. + entity_cfg: Scene entity config of the rigid object. + scale_factor_range: If same_scale_all_axes is True, should be [[s_min], [s_max]]. + Otherwise [[sx_min, sy_min, sz_min], [sx_max, sy_max, sz_max]]. + same_scale_all_axes: Whether to use same factor on x/y/z. + """ + if scale_factor_range is None: + return + if entity_cfg.uid not in env.sim.get_rigid_object_uid_list(): + return + + env_ids = _normalize_env_ids(env, env_ids) + rigid_object: RigidObject = env.sim.get_rigid_object(entity_cfg.uid) + scale = _sample_body_scale(env, env_ids, scale_factor_range, same_scale_all_axes) + rigid_object.set_body_scale(scale, env_ids=env_ids) + + +def randomize_rigid_objects_scale( + env: EmbodiedEnv, + env_ids: Union[torch.Tensor, None], + entity_cfgs: List[SceneEntityCfg], + scale_factor_range: tuple[list[float], list[float]] | None = None, + same_scale_all_axes: bool = True, + shared_sample: bool = False, +) -> None: + """Randomize body scale factors for multiple rigid objects. + + Args: + env: Environment instance. + env_ids: Target env ids. If None, applies to all envs. + entity_cfgs: List of scene entity configs (rigid objects). + scale_factor_range: Scale factor sampling range. + same_scale_all_axes: Whether to use same factor on x/y/z. + shared_sample: If True, sample one scale per-env and apply to *all* objects (sync). + If False, each object samples its own scales independently. + """ + if scale_factor_range is None: + return + + if not isinstance(entity_cfgs, list) or len(entity_cfgs) == 0: + return + + env_ids = _normalize_env_ids(env, env_ids) + + if shared_sample: + scale = _sample_body_scale( + env, env_ids, scale_factor_range, same_scale_all_axes + ) + for entity_cfg in entity_cfgs: + if entity_cfg.uid not in env.sim.get_rigid_object_uid_list(): + continue + rigid_object: RigidObject = env.sim.get_rigid_object(entity_cfg.uid) + rigid_object.set_body_scale(scale, env_ids=env_ids) + return + + for entity_cfg in entity_cfgs: + randomize_rigid_object_scale( + env=env, + env_ids=env_ids, + entity_cfg=entity_cfg, + scale_factor_range=scale_factor_range, + same_scale_all_axes=same_scale_all_axes, + ) + + +def randomize_rigid_object_body_scale( + env: EmbodiedEnv, + env_ids: Union[torch.Tensor, None], + entity_cfg: SceneEntityCfg, + scale_range: tuple[list[float], list[float]] | None = None, + same_scale_all_axes: bool = True, +) -> None: + """Deprecated. Use `randomize_rigid_object_scale` + `scale_factor_range`.""" + if scale_range is not None: + logger.log_warning( + "`randomize_rigid_object_body_scale` is deprecated. " + "Please migrate to `randomize_rigid_object_scale` with `scale_factor_range`." + ) + return randomize_rigid_object_scale( + env=env, + env_ids=env_ids, + entity_cfg=entity_cfg, + scale_factor_range=scale_range, + same_scale_all_axes=same_scale_all_axes, + ) + + +def randomize_rigid_objects_body_scale_sync( + env: EmbodiedEnv, + env_ids: Union[torch.Tensor, None], + entity_cfgs: List[SceneEntityCfg], + scale_range: tuple[list[float], list[float]] | None = None, + same_scale_all_axes: bool = True, +) -> None: + """Deprecated. Use `randomize_rigid_objects_scale(shared_sample=True)` + `scale_factor_range`.""" + if scale_range is not None: + logger.log_warning( + "`randomize_rigid_objects_body_scale_sync` is deprecated. " + "Please migrate to `randomize_rigid_objects_scale` with `shared_sample=true` " + "and `scale_factor_range`." + ) + return randomize_rigid_objects_scale( + env=env, + env_ids=env_ids, + entity_cfgs=entity_cfgs, + scale_factor_range=scale_range, + same_scale_all_axes=same_scale_all_axes, + shared_sample=True, + ) diff --git a/embodichain/lab/gym/envs/managers/randomization/scale.py b/embodichain/lab/gym/envs/managers/randomization/scale.py new file mode 100644 index 00000000..5e9ee874 --- /dev/null +++ b/embodichain/lab/gym/envs/managers/randomization/scale.py @@ -0,0 +1,182 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2025 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +from typing import TYPE_CHECKING, List, Union + +import torch + +from embodichain.lab.gym.envs.managers.cfg import SceneEntityCfg +from embodichain.lab.sim.objects import RigidObject +from embodichain.utils import logger +from embodichain.utils.math import sample_uniform + +if TYPE_CHECKING: + from embodichain.lab.gym.envs import EmbodiedEnv + + +def _normalize_env_ids( + env: EmbodiedEnv, env_ids: Union[torch.Tensor, None] +) -> torch.Tensor: + # Target all active environments if no specific IDs are provided + if env_ids is None: + return torch.arange(env.num_envs, device=env.device) + return env_ids + + +def _sample_body_scale( + env: EmbodiedEnv, + env_ids: torch.Tensor, + scale_factor_range: tuple[list[float], list[float]], + same_scale_all_axes: bool, +) -> torch.Tensor: + """Sample per-env body scale factors. + + Returns: + torch.Tensor: Shape (num_envs_selected, 3) scale factors for x/y/z. + """ + num_instance = len(env_ids) + if same_scale_all_axes: + low = torch.tensor(scale_factor_range[0][0], device=env.device) + high = torch.tensor(scale_factor_range[1][0], device=env.device) + s = sample_uniform(lower=low, upper=high, size=(num_instance,)) + return torch.stack([s, s, s], dim=1) + low = torch.tensor(scale_factor_range[0], device=env.device) + high = torch.tensor(scale_factor_range[1], device=env.device) + return sample_uniform(lower=low, upper=high, size=(num_instance, 3)) + + +def randomize_rigid_object_scale( + env: EmbodiedEnv, + env_ids: Union[torch.Tensor, None], + entity_cfg: SceneEntityCfg, + scale_factor_range: tuple[list[float], list[float]] | None = None, + same_scale_all_axes: bool = True, +) -> None: + """Randomize a rigid object's *body scale factors* (multiplicative, not absolute size). + + Args: + env: Environment instance. + env_ids: Target env ids. If None, applies to all envs. + entity_cfg: Scene entity config of the rigid object. + scale_factor_range: If same_scale_all_axes is True, should be [[s_min], [s_max]]. + Otherwise [[sx_min, sy_min, sz_min], [sx_max, sy_max, sz_max]]. + same_scale_all_axes: Whether to use same factor on x/y/z. + """ + if scale_factor_range is None: + return + if entity_cfg.uid not in env.sim.get_rigid_object_uid_list(): + return + + env_ids = _normalize_env_ids(env, env_ids) + rigid_object: RigidObject = env.sim.get_rigid_object(entity_cfg.uid) + scale = _sample_body_scale(env, env_ids, scale_factor_range, same_scale_all_axes) + rigid_object.set_body_scale(scale, env_ids=env_ids) + + +def randomize_rigid_objects_scale( + env: EmbodiedEnv, + env_ids: Union[torch.Tensor, None], + entity_cfgs: List[SceneEntityCfg], + scale_factor_range: tuple[list[float], list[float]] | None = None, + same_scale_all_axes: bool = True, + shared_sample: bool = False, +) -> None: + """Randomize body scale factors for multiple rigid objects. + + Args: + env: Environment instance. + env_ids: Target env ids. If None, applies to all envs. + entity_cfgs: List of scene entity configs (rigid objects). + scale_factor_range: Scale factor sampling range. + same_scale_all_axes: Whether to use same factor on x/y/z. + shared_sample: If True, sample one scale per-env and apply to *all* objects (sync). + If False, each object samples its own scales independently. + """ + if scale_factor_range is None: + return + + if not isinstance(entity_cfgs, list) or len(entity_cfgs) == 0: + return + + env_ids = _normalize_env_ids(env, env_ids) + + if shared_sample: + scale = _sample_body_scale( + env, env_ids, scale_factor_range, same_scale_all_axes + ) + for entity_cfg in entity_cfgs: + if entity_cfg.uid not in env.sim.get_rigid_object_uid_list(): + continue + rigid_object: RigidObject = env.sim.get_rigid_object(entity_cfg.uid) + rigid_object.set_body_scale(scale, env_ids=env_ids) + return + + for entity_cfg in entity_cfgs: + randomize_rigid_object_scale( + env=env, + env_ids=env_ids, + entity_cfg=entity_cfg, + scale_factor_range=scale_factor_range, + same_scale_all_axes=same_scale_all_axes, + ) + + +def randomize_rigid_object_body_scale( + env: EmbodiedEnv, + env_ids: Union[torch.Tensor, None], + entity_cfg: SceneEntityCfg, + scale_range: tuple[list[float], list[float]] | None = None, + same_scale_all_axes: bool = True, +) -> None: + """Deprecated. Use `randomize_rigid_object_scale` + `scale_factor_range`.""" + if scale_range is not None: + logger.log_warning( + "`randomize_rigid_object_body_scale` is deprecated. " + "Please migrate to `randomize_rigid_object_scale` with `scale_factor_range`." + ) + return randomize_rigid_object_scale( + env=env, + env_ids=env_ids, + entity_cfg=entity_cfg, + scale_factor_range=scale_range, + same_scale_all_axes=same_scale_all_axes, + ) + + +def randomize_rigid_objects_body_scale_sync( + env: EmbodiedEnv, + env_ids: Union[torch.Tensor, None], + entity_cfgs: List[SceneEntityCfg], + scale_range: tuple[list[float], list[float]] | None = None, + same_scale_all_axes: bool = True, +) -> None: + """Deprecated. Use `randomize_rigid_objects_scale(shared_sample=True)` + `scale_factor_range`.""" + if scale_range is not None: + logger.log_warning( + "`randomize_rigid_objects_body_scale_sync` is deprecated. " + "Please migrate to `randomize_rigid_objects_scale` with `shared_sample=true` " + "and `scale_factor_range`." + ) + return randomize_rigid_objects_scale( + env=env, + env_ids=env_ids, + entity_cfgs=entity_cfgs, + scale_factor_range=scale_range, + same_scale_all_axes=same_scale_all_axes, + shared_sample=True, + ) diff --git a/embodichain/lab/gym/envs/managers/randomization/visual.py b/embodichain/lab/gym/envs/managers/randomization/visual.py index d0c25d76..ce537f67 100644 --- a/embodichain/lab/gym/envs/managers/randomization/visual.py +++ b/embodichain/lab/gym/envs/managers/randomization/visual.py @@ -19,6 +19,7 @@ import torch import os import random +import copy from typing import TYPE_CHECKING, Literal, Union, Dict from embodichain.lab.sim.objects import Light, RigidObject, Articulation @@ -43,6 +44,54 @@ from embodichain.lab.gym.envs import EmbodiedEnv +__all__ = [ + "randomize_camera_extrinsics", + "randomize_light", + "randomize_camera_intrinsics", + "set_rigid_object_visual_material", + "randomize_visual_material", +] + + +def set_rigid_object_visual_material( + env: EmbodiedEnv, + env_ids: Union[torch.Tensor, None], + entity_cfg: SceneEntityCfg, + mat_cfg: Union[VisualMaterialCfg, Dict], +) -> None: + """Set a rigid object's visual material (deterministic, non-random). + + This helper exists to support configs that want fixed colors/materials during reset. + + Args: + env: Environment instance. + env_ids: Target env ids. If None, applies to all envs. + entity_cfg: Scene entity config (must point to a rigid object). + mat_cfg: Visual material configuration. Can be a VisualMaterialCfg object or a dict. + If a dict is provided, it will be converted to VisualMaterialCfg using from_dict(). + If uid is not specified in mat_cfg, it will default to "{entity_uid}_mat". + """ + if entity_cfg.uid not in env.sim.get_rigid_object_uid_list(): + return + + if env_ids is None: + env_ids = torch.arange(env.num_envs, device="cpu") + else: + env_ids = env_ids.cpu() + + if isinstance(mat_cfg, dict): + mat_cfg = VisualMaterialCfg.from_dict(mat_cfg) + + mat_cfg = copy.deepcopy(mat_cfg) + + if not mat_cfg.uid or mat_cfg.uid == "default_mat": + mat_cfg.uid = f"{entity_cfg.uid}_mat" + + mat = env.sim.create_visual_material(mat_cfg) + obj: RigidObject = env.sim.get_rigid_object(entity_cfg.uid) + obj.set_visual_material(mat, env_ids=env_ids) + + def randomize_camera_extrinsics( env: EmbodiedEnv, env_ids: Union[torch.Tensor, None], diff --git a/embodichain/lab/gym/envs/tasks/tableware/blocks_ranking_rgb.py b/embodichain/lab/gym/envs/tasks/tableware/blocks_ranking_rgb.py new file mode 100644 index 00000000..fc7eb442 --- /dev/null +++ b/embodichain/lab/gym/envs/tasks/tableware/blocks_ranking_rgb.py @@ -0,0 +1,88 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2025 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +import torch +import numpy as np + +from embodichain.lab.gym.envs import EmbodiedEnv, EmbodiedEnvCfg +from embodichain.lab.gym.utils.registration import register_env +from embodichain.utils import logger + +__all__ = ["BlocksRankingRGBEnv"] + + +@register_env("BlocksRankingRGB-v1", max_episode_steps=600) +class BlocksRankingRGBEnv(EmbodiedEnv): + def __init__(self, cfg: EmbodiedEnvCfg = None, **kwargs): + super().__init__(cfg, **kwargs) + + action_config = kwargs.get("action_config", None) + if action_config is not None: + self.action_config = action_config + + def is_task_success(self, **kwargs) -> torch.Tensor: + """Determine if the task is successfully completed. + + The task is successful if: + 1. Three blocks are arranged in RGB order from left to right: + - Red block (block_1) x < Green block (block_2) x < Blue block (block_3) x + 2. All blocks are close together (within tolerance) + + Args: + **kwargs: Additional arguments for task-specific success criteria. + + Returns: + torch.Tensor: A boolean tensor indicating success for each environment in the batch. + """ + try: + block1 = self.sim.get_rigid_object("block_1") # Red + block2 = self.sim.get_rigid_object("block_2") # Green + block3 = self.sim.get_rigid_object("block_3") # Blue + except Exception as e: + logger.log_warning(f"Blocks not found: {e}, returning False.") + return torch.zeros(self.num_envs, dtype=torch.bool, device=self.device) + + # Get block poses + block1_pose = block1.get_local_pose(to_matrix=True) + block2_pose = block2.get_local_pose(to_matrix=True) + block3_pose = block3.get_local_pose(to_matrix=True) + + # Extract positions (x, y, z) + block1_pos = block1_pose[:, :3, 3] # (num_envs, 3) + block2_pos = block2_pose[:, :3, 3] + block3_pos = block3_pose[:, :3, 3] + + # Tolerance for checking if blocks are close together + eps = torch.tensor([0.13, 0.03], dtype=torch.float32, device=self.device) + + # Check if blocks are close together in x-y plane + # block1 and block2 should be close + block1_block2_diff = torch.abs(block1_pos[:, :2] - block2_pos[:, :2]) + blocks_close_12 = torch.all(block1_block2_diff < eps.unsqueeze(0), dim=1) + + # block2 and block3 should be close + block2_block3_diff = torch.abs(block2_pos[:, :2] - block3_pos[:, :2]) + blocks_close_23 = torch.all(block2_block3_diff < eps.unsqueeze(0), dim=1) + + # Check RGB order: block1 (red) x < block2 (green) x < block3 (blue) x + rgb_order = (block1_pos[:, 0] < block2_pos[:, 0]) & ( + block2_pos[:, 0] < block3_pos[:, 0] + ) + + # Task succeeds if blocks are close together and in RGB order + success = blocks_close_12 & blocks_close_23 & rgb_order + + return success diff --git a/embodichain/lab/gym/envs/tasks/tableware/blocks_ranking_size.py b/embodichain/lab/gym/envs/tasks/tableware/blocks_ranking_size.py new file mode 100644 index 00000000..45b66997 --- /dev/null +++ b/embodichain/lab/gym/envs/tasks/tableware/blocks_ranking_size.py @@ -0,0 +1,89 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2025 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +import torch +import numpy as np + +from embodichain.lab.gym.envs import EmbodiedEnv, EmbodiedEnvCfg +from embodichain.lab.gym.utils.registration import register_env +from embodichain.utils import logger + +__all__ = ["BlocksRankingSizeEnv"] + + +@register_env("BlocksRankingSize-v1", max_episode_steps=600) +class BlocksRankingSizeEnv(EmbodiedEnv): + def __init__(self, cfg: EmbodiedEnvCfg = None, **kwargs): + super().__init__(cfg, **kwargs) + + action_config = kwargs.get("action_config", None) + if action_config is not None: + self.action_config = action_config + + def is_task_success(self, **kwargs) -> torch.Tensor: + """Determine if the task is successfully completed. + + The task is successful if: + 1. Three blocks are arranged in size order from left to right: + - Large block (block_1) x < Medium block (block_2) x < Small block (block_3) x + 2. All blocks are close together (within tolerance) + + Args: + **kwargs: Additional arguments for task-specific success criteria. + + Returns: + torch.Tensor: A boolean tensor indicating success for each environment in the batch. + """ + try: + block1 = self.sim.get_rigid_object("block_1") # Large + block2 = self.sim.get_rigid_object("block_2") # Medium + block3 = self.sim.get_rigid_object("block_3") # Small + except Exception as e: + logger.log_warning(f"Blocks not found: {e}, returning False.") + return torch.zeros(self.num_envs, dtype=torch.bool, device=self.device) + + # Get block poses + block1_pose = block1.get_local_pose(to_matrix=True) + block2_pose = block2.get_local_pose(to_matrix=True) + block3_pose = block3.get_local_pose(to_matrix=True) + + # Extract positions (x, y, z) + block1_pos = block1_pose[:, :3, 3] # (num_envs, 3) + block2_pos = block2_pose[:, :3, 3] + block3_pos = block3_pose[:, :3, 3] + + # Tolerance for checking if blocks are close together + # Same as RoboTwin: eps = [0.13, 0.03] + eps = torch.tensor([0.13, 0.03], dtype=torch.float32, device=self.device) + + # Check if blocks are close together in x-y plane + # block1 and block2 should be close + block1_block2_diff = torch.abs(block1_pos[:, :2] - block2_pos[:, :2]) + blocks_close_12 = torch.all(block1_block2_diff < eps.unsqueeze(0), dim=1) + + # block2 and block3 should be close + block2_block3_diff = torch.abs(block2_pos[:, :2] - block3_pos[:, :2]) + blocks_close_23 = torch.all(block2_block3_diff < eps.unsqueeze(0), dim=1) + + # Check size order: block1 (large) x < block2 (medium) x < block3 (small) x + size_order = (block1_pos[:, 0] < block2_pos[:, 0]) & ( + block2_pos[:, 0] < block3_pos[:, 0] + ) + + # All conditions must be satisfied + success = blocks_close_12 & blocks_close_23 & size_order + + return success diff --git a/embodichain/lab/gym/envs/tasks/tableware/match_object_container.py b/embodichain/lab/gym/envs/tasks/tableware/match_object_container.py new file mode 100644 index 00000000..2af1d374 --- /dev/null +++ b/embodichain/lab/gym/envs/tasks/tableware/match_object_container.py @@ -0,0 +1,176 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2025 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +import torch +import numpy as np + +from embodichain.lab.gym.envs import EmbodiedEnv, EmbodiedEnvCfg +from embodichain.lab.gym.utils.registration import register_env +from embodichain.utils import logger + +__all__ = ["MatchObjectContainerEnv"] + + +@register_env("MatchObjectContainer-v1", max_episode_steps=600) +class MatchObjectContainerEnv(EmbodiedEnv): + def __init__(self, cfg: EmbodiedEnvCfg = None, **kwargs): + super().__init__(cfg, **kwargs) + + action_config = kwargs.get("action_config", None) + if action_config is not None: + self.action_config = action_config + + def is_task_success(self, **kwargs) -> torch.Tensor: + """Determine if the task is successfully completed. + + This is a classification task: place blocks into matching shaped containers. + The task is successful if: + 1. Both cube blocks are inside container_cube (block_cube_1 and block_cube_2 -> container_cube) + 2. Both sphere blocks are inside container_sphere (block_sphere_1 and block_sphere_2 -> container_sphere) + 3. Both containers are up + + Args: + **kwargs: Additional arguments for task-specific success criteria. + + Returns: + torch.Tensor: A boolean tensor indicating success for each environment in the batch. + """ + try: + block_cube_1 = self.sim.get_rigid_object("block_cube_1") + block_sphere_1 = self.sim.get_rigid_object("block_sphere_1") + block_cube_2 = self.sim.get_rigid_object("block_cube_2") + block_sphere_2 = self.sim.get_rigid_object("block_sphere_2") + container_cube = self.sim.get_rigid_object("container_cube") + container_sphere = self.sim.get_rigid_object("container_sphere") + except Exception as e: + logger.log_warning(f"Blocks or containers not found: {e}, returning False.") + return torch.zeros(self.num_envs, dtype=torch.bool, device=self.device) + + # Get poses + block_cube_1_pose = block_cube_1.get_local_pose(to_matrix=True) + block_sphere_1_pose = block_sphere_1.get_local_pose(to_matrix=True) + block_cube_2_pose = block_cube_2.get_local_pose(to_matrix=True) + block_sphere_2_pose = block_sphere_2.get_local_pose(to_matrix=True) + container_cube_pose = container_cube.get_local_pose(to_matrix=True) + container_sphere_pose = container_sphere.get_local_pose(to_matrix=True) + + # Extract positions + block_cube_1_pos = block_cube_1_pose[:, :3, 3] # (num_envs, 3) + block_sphere_1_pos = block_sphere_1_pose[:, :3, 3] + block_cube_2_pos = block_cube_2_pose[:, :3, 3] + block_sphere_2_pos = block_sphere_2_pose[:, :3, 3] + container_cube_pos = container_cube_pose[:, :3, 3] + container_sphere_pos = container_sphere_pose[:, :3, 3] + + container_cube_fallen = self._is_fall(container_cube_pose) + container_sphere_fallen = self._is_fall(container_sphere_pose) + + # Check if blocks are inside their matching containers + # I got radius and height from the container_metal.obj file + container_bottom_radius = 0.1067 # Inner radius + z_tolerance = 0.05 # Vertical tolerance + container_height = 0.068 # Container height + container_half_height = container_height / 2 + + # Check if blocks are in their matching containers + cube_1_in_container = self._is_block_in_container( + block_cube_1_pos, + container_cube_pos, + container_bottom_radius, + container_half_height, + z_tolerance, + ) + cube_2_in_container = self._is_block_in_container( + block_cube_2_pos, + container_cube_pos, + container_bottom_radius, + container_half_height, + z_tolerance, + ) + sphere_1_in_container = self._is_block_in_container( + block_sphere_1_pos, + container_sphere_pos, + container_bottom_radius, + container_half_height, + z_tolerance, + ) + sphere_2_in_container = self._is_block_in_container( + block_sphere_2_pos, + container_sphere_pos, + container_bottom_radius, + container_half_height, + z_tolerance, + ) + + # Task success if cubes and spheres are in containers and containers are up + success = ( + cube_1_in_container + & cube_2_in_container + & sphere_1_in_container + & sphere_2_in_container + & ~container_cube_fallen + & ~container_sphere_fallen + ) + + return success + + def _is_block_in_container( + self, + block_pos: torch.Tensor, + container_pos: torch.Tensor, + container_bottom_radius: float, + container_half_height: float, + z_tolerance: float, + ) -> torch.Tensor: + """Check if a block is inside a container. + + Args: + block_pos: Block position (num_envs, 3) + container_pos: Container center position (num_envs, 3) + container_bottom_radius: Inner radius of container bottom in meters + container_half_height: Half height of container in meters + z_tolerance: Vertical tolerance for bottom check in meters + + Returns: + Boolean tensor indicating if block is in container (num_envs,) + """ + # XY plane distance check + xy_diff = torch.norm(block_pos[:, :2] - container_pos[:, :2], dim=1) + + # Z coordinate check: container center is at container_pos[:, 2] + # Container bottom = container_pos[:, 2] - container_half_height + # Container top = container_pos[:, 2] + container_half_height + z_above_bottom = ( + block_pos[:, 2] > container_pos[:, 2] - container_half_height - z_tolerance + ) + z_within_height = block_pos[:, 2] < container_pos[:, 2] + container_half_height + + return (xy_diff < container_bottom_radius) & z_above_bottom & z_within_height + + def _is_fall(self, pose: torch.Tensor) -> torch.Tensor: + # Extract z-axis from rotation matrix (last column, first 3 elements) + pose_rz = pose[:, :3, 2] + world_z_axis = torch.tensor([0, 0, 1], dtype=pose.dtype, device=pose.device) + + # Compute dot product for each batch element + dot_product = torch.sum(pose_rz * world_z_axis, dim=-1) # Shape: (batch_size,) + + # Clamp to avoid numerical issues with arccos + dot_product = torch.clamp(dot_product, -1.0, 1.0) + + # Compute angle and check if fallen + angle = torch.arccos(dot_product) + return angle >= torch.pi / 4 diff --git a/embodichain/lab/gym/envs/tasks/tableware/place_object_drawer.py b/embodichain/lab/gym/envs/tasks/tableware/place_object_drawer.py new file mode 100644 index 00000000..f83a4b70 --- /dev/null +++ b/embodichain/lab/gym/envs/tasks/tableware/place_object_drawer.py @@ -0,0 +1,83 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2025 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +import torch +import numpy as np + + +from embodichain.lab.gym.envs import EmbodiedEnv, EmbodiedEnvCfg +from embodichain.lab.gym.utils.registration import register_env +from embodichain.utils import logger + +__all__ = ["PlaceObjectDrawerEnv"] + + +@register_env("PlaceObjectDrawer-v1", max_episode_steps=600) +class PlaceObjectDrawerEnv(EmbodiedEnv): + def __init__(self, cfg: EmbodiedEnvCfg = None, **kwargs): + super().__init__(cfg, **kwargs) + + action_config = kwargs.get("action_config", None) + if action_config is not None: + self.action_config = action_config + + def is_task_success(self, **kwargs) -> torch.Tensor: + """Determine if the task is successfully completed. + + The task is successful if: + 1. Object is within drawer inner_box area + 2. Drawer has been closed + + Args: + **kwargs: Additional arguments for task-specific success criteria. + + Returns: + torch.Tensor: A boolean tensor indicating success for each environment in the batch. + """ + try: + object_obj = self.sim.get_rigid_object("object") + drawer = self.sim.get_articulation("drawer") + except Exception as e: + logger.log_warning(f"Object or drawer not found: {e}, returning False.") + return torch.zeros(self.num_envs, dtype=torch.bool, device=self.device) + + # Get poses + object_pose = object_obj.get_local_pose(to_matrix=True) + # Get drawer inner_box (drawer interior) pose, not outer_box + inner_box_pose = drawer.get_link_pose("inner_box", to_matrix=True) + + # Extract positions + object_pos = object_pose[:, :3, 3] # (num_envs, 3) + inner_box_pos = inner_box_pose[:, :3, 3] # (num_envs, 3) + + # Get drawer joint position + drawer_qpos = drawer.get_qpos() # (num_envs, num_joints) + drawer_joint_pos = drawer_qpos[:, 0] + + # Check if drawer has been closed + drawer_closed = drawer_joint_pos < 0.05 + + # Check if object is within drawer inner_box area + xy_diff = torch.abs(object_pos[:, :2] - inner_box_pos[:, :2]) # (num_envs, 2) + xy_tolerance = torch.tensor( + [0.03, 0.03], dtype=torch.float32, device=self.device + ) + object_in_drawer_xy = torch.all(xy_diff < xy_tolerance.unsqueeze(0), dim=1) + + # Object must be in drawer and drawer must be closed + success = drawer_closed & object_in_drawer_xy + + return success diff --git a/embodichain/lab/gym/envs/tasks/tableware/stack_blocks_two.py b/embodichain/lab/gym/envs/tasks/tableware/stack_blocks_two.py new file mode 100644 index 00000000..5637e105 --- /dev/null +++ b/embodichain/lab/gym/envs/tasks/tableware/stack_blocks_two.py @@ -0,0 +1,289 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2025 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +import torch +import numpy as np + +from embodichain.lab.gym.envs import EmbodiedEnv, EmbodiedEnvCfg +from embodichain.lab.gym.utils.registration import register_env +from embodichain.lab.sim.planners.motion_generator import MotionGenerator +from embodichain.lab.sim.planners.utils import TrajectorySampleMethod +from embodichain.utils import logger + +__all__ = ["StackBlocksTwoEnv"] + + +@register_env("StackBlocksTwo-v1", max_episode_steps=600) +class StackBlocksTwoEnv(EmbodiedEnv): + def __init__(self, cfg: EmbodiedEnvCfg = None, **kwargs): + super().__init__(cfg, **kwargs) + + action_config = kwargs.get("action_config", None) + if action_config is not None: + self.action_config = action_config + + def create_demo_action_list(self, *args, **kwargs): + """ + Create a demonstration action list for stacking block_2 on top of block_1. + + This expert trajectory: + 1. Moves to block_2 position + 2. Grasps block_2 + 3. Lifts block_2 + 4. Moves to block_1 position + 5. Places block_2 on top of block_1 + 6. Releases and retracts + + Returns: + list: A list of demo actions generated by the task. + """ + try: + block1 = self.sim.get_rigid_object("block_1") + block2 = self.sim.get_rigid_object("block_2") + except Exception as e: + logger.log_warning(f"Blocks not found: {e}, returning empty action list.") + return [] + + # Get block positions + block1_pose = block1.get_local_pose(to_matrix=True) + block2_pose = block2.get_local_pose(to_matrix=True) + + block1_pos = block1_pose[:, :3, 3] # (num_envs, 3) + block2_pos = block2_pose[:, :3, 3] # (num_envs, 3) + + right_arm_ids = self.robot.get_joint_ids(name="right_arm") + right_eef_ids = self.robot.get_joint_ids(name="right_eef") + + init_qpos = self.robot.get_qpos() + init_right_arm_qpos = init_qpos[:, right_arm_ids] + + init_right_arm_xpos = self.robot.compute_fk( + qpos=init_right_arm_qpos, name="right_arm", to_matrix=True + ) + + # 1. Approach block_2 + approach_block2_xpos = init_right_arm_xpos.clone() + approach_block2_xpos[:, :3, 3] = block2_pos + torch.tensor( + [0.02, 0.0, 0.1], dtype=torch.float32, device=self.device + ).unsqueeze(0) + + # 2. Pick block_2 + pick_block2_xpos = init_right_arm_xpos.clone() + pick_block2_xpos[:, :3, 3] = block2_pos + torch.tensor( + [0.02, 0.0, -0.025], dtype=torch.float32, device=self.device + ).unsqueeze(0) + + # 3. Lift block_2 + lift_block2_xpos = pick_block2_xpos.clone() + lift_block2_xpos[:, 2, 3] += 0.15 + + # 4. Approach block_1 + approach_block1_xpos = init_right_arm_xpos.clone() + target_place_pos = block1_pos.clone() + target_place_pos[:, 2] += 0.05 # block_1 height + block_2 height + approach_block1_xpos[:, :3, 3] = target_place_pos + torch.tensor( + [0.025, 0.0, 0.05], dtype=torch.float32, device=self.device + ).unsqueeze(0) + + # 5. Place block_2 on block_1 + place_block2_xpos = approach_block1_xpos.clone() + place_block2_xpos[:, :3, 3] = target_place_pos + torch.tensor( + [0.025, 0.0, 0.02], dtype=torch.float32, device=self.device + ).unsqueeze(0) + + # 6. Retract + retract_xpos = approach_block1_xpos.clone() + + retract_xpos[:, 2, 3] += 0.1 + + # Compute IK for each waypoint + waypoints = [ + ("approach_block2", approach_block2_xpos), + ("pick_block2", pick_block2_xpos), + ("lift_block2", lift_block2_xpos), + ("approach_block1", approach_block1_xpos), + ("place_block2", place_block2_xpos), + ("retract", retract_xpos), + ] + + qpos_waypoints = [] + current_qpos = init_right_arm_qpos + + for waypoint_name, waypoint_xpos in waypoints: + is_success, qpos = self.robot.compute_ik( + pose=waypoint_xpos, joint_seed=current_qpos, name="right_arm" + ) + + # Check IK success + success_flag = is_success.all() if isinstance(is_success, torch.Tensor) else is_success + if not success_flag: + logger.log_warning(f"IK failed for {waypoint_name}") + qpos = current_qpos + + qpos_waypoints.append(qpos) + current_qpos = qpos + + # Create motion generator for smooth trajectory + motion_gen = MotionGenerator( + robot=self.robot, + uid="right_arm", + planner_type="toppra", + default_velocity=0.2, + default_acceleration=0.5, + ) + + action_list = [] + + gripper_open = torch.tensor([0.05, 0.05], dtype=torch.float32, device=self.device) + gripper_close = torch.tensor([0.0, 0.0], dtype=torch.float32, device=self.device) + + # Convert waypoints to numpy for motion generator + qpos_waypoints_np = [q[0].cpu().numpy() for q in qpos_waypoints] + + # Generate trajectory segments + segments = [ + (0, 1, 30, gripper_open), + (1, 1, 10, gripper_close), + (1, 2, 30, gripper_close), + (2, 3, 40, gripper_close), + (3, 4, 30, gripper_close), + (4, 4, 10, gripper_open), + (4, 5, 30, gripper_open), + ] + + for start_idx, end_idx, num_steps, gripper_state in segments: + if start_idx == end_idx: + # Hold position (for gripper action) + qpos = qpos_waypoints[start_idx] + for _ in range(num_steps): + action = init_qpos.clone() + action[:, right_arm_ids] = qpos + action[:, right_eef_ids] = gripper_state.unsqueeze(0).expand( + self.num_envs, -1 + ) + action_list.append(action) + else: + # Generate smooth trajectory between waypoints + qpos_list = [qpos_waypoints_np[start_idx], qpos_waypoints_np[end_idx]] + + out_qpos_list, _ = motion_gen.create_discrete_trajectory( + qpos_list=qpos_list, + is_linear=False, + sample_method=TrajectorySampleMethod.QUANTITY, + sample_num=num_steps, + is_use_current_qpos=False, + ) + + # Convert to torch and add to action list + for qpos_item in out_qpos_list: + qpos = torch.as_tensor(qpos_item, dtype=torch.float32, device=self.device) + + qpos = qpos.flatten() + + if qpos.shape[0] != len(right_arm_ids): + logger.log_warning( + f"Qpos shape mismatch: got {qpos.shape[0]}, expected {len(right_arm_ids)}" + ) + continue + + qpos = qpos.unsqueeze(0).expand(self.num_envs, -1) + + action = init_qpos.clone() + action[:, right_arm_ids] = qpos + action[:, right_eef_ids] = gripper_state.unsqueeze(0).expand( + self.num_envs, -1 + ) + action_list.append(action) + + logger.log_info( + f"Generated {len(action_list)} demo actions for stacking blocks" + ) + self.action_length = len(action_list) + return action_list + + def is_task_success(self, **kwargs) -> torch.Tensor: + """Determine if the task is successfully completed. This is mainly used in the data generation process + of the imitation learning. + + The task is successful if: + 1. Block2 is stacked on top of Block1 + 2. Both blocks haven't fallen over + + Args: + **kwargs: Additional arguments for task-specific success criteria. + + Returns: + torch.Tensor: A boolean tensor indicating success for each environment in the batch. + """ + try: + block1 = self.sim.get_rigid_object("block_1") + block2 = self.sim.get_rigid_object("block_2") + except Exception as e: + logger.log_warning(f"Block1 or Block2 not found: {e}, returning False.") + return torch.zeros(self.num_envs, dtype=torch.bool, device=self.device) + + # Get poses + block1_pose = block1.get_local_pose(to_matrix=True) + block2_pose = block2.get_local_pose(to_matrix=True) + + # Extract positions + block1_pos = block1_pose[:, :3, 3] # (num_envs, 3) + block2_pos = block2_pose[:, :3, 3] + + # Check if blocks haven't fallen + block1_fallen = self._is_fall(block1_pose) + block2_fallen = self._is_fall(block2_pose) + + # Block2 should be on top of block1 + expected_block2_pos = torch.stack( + [ + block1_pos[:, 0], + block1_pos[:, 1], + block1_pos[:, 2] + 0.05, # block1 z + block height + ], + dim=1, + ) + + # Tolerance + eps = torch.tensor( + [0.025, 0.025, 0.012], dtype=torch.float32, device=self.device + ) + + # Check if block2 is within tolerance of expected position + position_diff = torch.abs(block2_pos - expected_block2_pos) # (num_envs, 3) + within_tolerance = torch.all( + position_diff < eps.unsqueeze(0), dim=1 + ) # (num_envs,) + + # Task succeeds if blocks are stacked correctly and haven't fallen + success = within_tolerance & ~block1_fallen & ~block2_fallen + + return success + + def _is_fall(self, pose: torch.Tensor) -> torch.Tensor: + # Extract z-axis from rotation matrix (last column, first 3 elements) + pose_rz = pose[:, :3, 2] + world_z_axis = torch.tensor([0, 0, 1], dtype=pose.dtype, device=pose.device) + + # Compute dot product for each batch element + dot_product = torch.sum(pose_rz * world_z_axis, dim=-1) # Shape: (batch_size,) + + # Clamp to avoid numerical issues with arccos + dot_product = torch.clamp(dot_product, -1.0, 1.0) + + # Compute angle and check if fallen + angle = torch.arccos(dot_product) + return angle >= torch.pi / 4 diff --git a/embodichain/lab/gym/envs/tasks/tableware/stack_cups.py b/embodichain/lab/gym/envs/tasks/tableware/stack_cups.py new file mode 100644 index 00000000..99c3937d --- /dev/null +++ b/embodichain/lab/gym/envs/tasks/tableware/stack_cups.py @@ -0,0 +1,118 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2025 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +import torch +import numpy as np + +from embodichain.lab.gym.envs import EmbodiedEnv, EmbodiedEnvCfg +from embodichain.lab.gym.utils.registration import register_env +from embodichain.utils import logger + +__all__ = ["StackCupsEnv"] + + +@register_env("StackCups-v1", max_episode_steps=600) +class StackCupsEnv(EmbodiedEnv): + def __init__(self, cfg: EmbodiedEnvCfg = None, **kwargs): + super().__init__(cfg, **kwargs) + + action_config = kwargs.get("action_config", None) + if action_config is not None: + self.action_config = action_config + + def is_task_success(self, **kwargs) -> torch.Tensor: + """Determine if the task is successfully completed. + + The task is successful if: + 1. Both cups haven't fallen over + 2. Cups are aligned in xy plane + 3. Top cup is only slightly higher than base cup + + Args: + **kwargs: Additional arguments for task-specific success criteria. + + Returns: + torch.Tensor: A boolean tensor indicating success for each environment in the batch. + """ + try: + cup1 = self.sim.get_rigid_object("cup_1") + cup2 = self.sim.get_rigid_object("cup_2") + except Exception as e: + logger.log_warning(f"Cups not found: {e}, returning False.") + return torch.zeros(self.num_envs, dtype=torch.bool, device=self.device) + + # Get cup poses + cup1_pose = cup1.get_local_pose(to_matrix=True) + cup2_pose = cup2.get_local_pose(to_matrix=True) + + # Extract positions + cup1_pos = cup1_pose[:, :3, 3] # (num_envs, 3) + cup2_pos = cup2_pose[:, :3, 3] # (num_envs, 3) + + # Check if cups haven't fallen + cup1_fallen = self._is_fall(cup1_pose) + cup2_fallen = self._is_fall(cup2_pose) + + # Determine which cup is lower (should be cup1) + cup1_is_lower = cup1_pos[:, 2] < cup2_pos[:, 2] + + # Use the lower cup as base, calculate expected position of top cup + base_cup_pos = torch.where( + cup1_is_lower.unsqueeze(1).expand_as(cup1_pos), cup1_pos, cup2_pos + ) + top_cup_pos = torch.where( + cup1_is_lower.unsqueeze(1).expand_as(cup2_pos), cup2_pos, cup1_pos + ) + + # Top cup should be slightly higher than base cup + cup_rim_offset = 0.015 + expected_top_pos = torch.stack( + [ + base_cup_pos[:, 0], + base_cup_pos[:, 1], + base_cup_pos[:, 2] + cup_rim_offset, + ], + dim=1, + ) + + # Tolerance + eps = torch.tensor([0.04, 0.04, 0.02], dtype=torch.float32, device=self.device) + + # Check if top cup is within tolerance of expected position + position_diff = torch.abs(top_cup_pos - expected_top_pos) # (num_envs, 3) + stacked_correctly = torch.all( + position_diff < eps.unsqueeze(0), dim=1 + ) # (num_envs,) + + # Task succeeds if cups are stacked correctly and haven't fallen + success = stacked_correctly & ~cup1_fallen & ~cup2_fallen + + return success + + def _is_fall(self, pose: torch.Tensor) -> torch.Tensor: + # Extract z-axis from rotation matrix (last column, first 3 elements) + pose_rz = pose[:, :3, 2] + world_z_axis = torch.tensor([0, 0, 1], dtype=pose.dtype, device=pose.device) + + # Compute dot product for each batch element + dot_product = torch.sum(pose_rz * world_z_axis, dim=-1) # Shape: (batch_size,) + + # Clamp to avoid numerical issues with arccos + dot_product = torch.clamp(dot_product, -1.0, 1.0) + + # Compute angle and check if fallen + angle = torch.arccos(dot_product) + return angle >= torch.pi / 4 diff --git a/embodichain/lab/gym/utils/gym_utils.py b/embodichain/lab/gym/utils/gym_utils.py index 40ced87f..0aa13340 100644 --- a/embodichain/lab/gym/utils/gym_utils.py +++ b/embodichain/lab/gym/utils/gym_utils.py @@ -449,20 +449,27 @@ class ComponentCfg: for dataset_name, dataset_params in config["env"]["dataset"].items(): dataset_params_modified = deepcopy(dataset_params) - # Find the function from multiple modules using the utility function - dataset_func = find_function_from_modules( - dataset_params["func"], - dataset_modules, - raise_if_not_found=True, - ) + # Check if this is a functor configuration (has "func" field) or a plain config + if "func" in dataset_params: + # Find the function from multiple modules using the utility function + dataset_func = find_function_from_modules( + dataset_params["func"], + dataset_modules, + raise_if_not_found=True, + ) - dataset = DatasetFunctorCfg( - func=dataset_func, - mode=dataset_params_modified["mode"], - params=dataset_params_modified["params"], - ) + from embodichain.lab.gym.envs.managers import DatasetFunctorCfg + + dataset = DatasetFunctorCfg( + func=dataset_func, + mode=dataset_params_modified["mode"], + params=dataset_params_modified["params"], + ) - setattr(env_cfg.dataset, dataset_name, dataset) + setattr(env_cfg.dataset, dataset_name, dataset) + else: + # Plain configuration (e.g., robot_meta), set directly + setattr(env_cfg.dataset, dataset_name, dataset_params_modified) env_cfg.events = ComponentCfg() if "events" in config["env"]: @@ -481,6 +488,12 @@ class ComponentCfg: **event_params_modified["params"]["entity_cfg"] ) event_params_modified["params"]["entity_cfg"] = entity_cfg + if "entity_cfgs" in event_params["params"]: + entity_cfgs = [ + SceneEntityCfg(**cfg) + for cfg in event_params_modified["params"]["entity_cfgs"] + ] + event_params_modified["params"]["entity_cfgs"] = entity_cfgs # Find the function from multiple modules using the utility function event_func = find_function_from_modules( diff --git a/embodichain/lab/sim/sim_manager.py b/embodichain/lab/sim/sim_manager.py index c1aec881..8f79a234 100644 --- a/embodichain/lab/sim/sim_manager.py +++ b/embodichain/lab/sim/sim_manager.py @@ -158,6 +158,8 @@ class SimulationManager: - physics simulation management, eg. time step, manual update, etc. - interactive control via gizmo and window callbacks events. + This class implements the singleton pattern to ensure only one instance exists at a time. + Args: sim_config (SimulationManagerCfg, optional): simulation configuration. Defaults to SimulationManagerCfg(). """ diff --git a/examples/gym/run_stack_blocks_two.sh b/examples/gym/run_stack_blocks_two.sh new file mode 100644 index 00000000..a24fc561 --- /dev/null +++ b/examples/gym/run_stack_blocks_two.sh @@ -0,0 +1,5 @@ +python -m embodichain.lab.scripts.run_env \ + --gym_config configs/gym/stack_blocks_two/cobot_magic_3cam.json \ + --num_envs 1 \ + --device cpu \ + --debug_mode