From fa4b0f07ad9912825ed9bdc337b9d01283c35be4 Mon Sep 17 00:00:00 2001 From: Michel Breyer Date: Tue, 7 Sep 2021 22:29:42 +0200 Subject: [PATCH] Refactor simulated scenes --- CMakeLists.txt | 3 + active_grasp/bbox.py | 4 +- active_grasp/simulation.py | 147 +++++++++++++++++++++++-------------- cfg/active_grasp.yaml | 1 + package.xml | 2 + scripts/bt_sim_node.py | 5 +- 6 files changed, 104 insertions(+), 58 deletions(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index 576ceb9..962fb78 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -2,7 +2,9 @@ cmake_minimum_required(VERSION 3.1) project(active_grasp) find_package(catkin REQUIRED COMPONENTS + rospy geometry_msgs + std_msgs message_generation ) @@ -22,4 +24,5 @@ add_service_files( generate_messages( DEPENDENCIES geometry_msgs + std_msgs ) diff --git a/active_grasp/bbox.py b/active_grasp/bbox.py index 2747f4e..a64ab68 100644 --- a/active_grasp/bbox.py +++ b/active_grasp/bbox.py @@ -7,8 +7,8 @@ from robot_helpers.ros.conversions import to_point_msg, from_point_msg class AABBox: def __init__(self, bbox_min, bbox_max): - self.min = bbox_min - self.max = bbox_max + self.min = np.asarray(bbox_min) + self.max = np.asarray(bbox_max) @property def corners(self): diff --git a/active_grasp/simulation.py b/active_grasp/simulation.py index 2133a31..7332cc6 100644 --- a/active_grasp/simulation.py +++ b/active_grasp/simulation.py @@ -6,17 +6,20 @@ import rospkg from active_grasp.bbox import AABBox from robot_helpers.bullet import * from robot_helpers.model import Model -from robot_helpers.spatial import Rotation, Transform +from robot_helpers.spatial import Rotation + + +rospack = rospkg.RosPack() class Simulation: - def __init__(self, gui): + """Robot is placed s.t. world and base frames are the same""" + + def __init__(self, gui, scene_id): self.configure_physics_engine(gui, 60, 4) self.configure_visualizer() - self.find_urdfs() - self.load_table() self.load_robot() - self.object_uids = [] + self.scene = get_scene(scene_id) def configure_physics_engine(self, gui, rate, sub_step_count): self.rate = rate @@ -28,41 +31,25 @@ class Simulation: def configure_visualizer(self): # p.configureDebugVisualizer(p.COV_ENABLE_GUI, 0) - p.resetDebugVisualizerCamera(1.4, 50, -35, [0.0, 0.0, 0.6]) - - def find_urdfs(self): - rospack = rospkg.RosPack() - assets_path = Path(rospack.get_path("active_grasp")) / "assets" - self.panda_urdf = assets_path / "urdfs/franka/panda_arm_hand.urdf" - root = Path(rospack.get_path("vgn")) / "assets/urdfs/packed/test" - self.urdfs = [str(f) for f in root.iterdir() if f.suffix == ".urdf"] - - def load_table(self): - p.loadURDF("plane.urdf") - ori = Rotation.from_rotvec(np.array([0, 0, np.pi / 2])).as_quat() - p.loadURDF("table/table.urdf", baseOrientation=ori, useFixedBase=True) - self.length = 0.3 - self.origin = [-0.3, -0.5 * self.length, 0.5] + p.resetDebugVisualizerCamera(1.2, 30, -30, [0.4, 0.0, 0.2]) def load_robot(self): - self.T_world_base = Transform.translation(np.r_[-0.6, 0.0, 0.4]) - self.arm = BtPandaArm(self.panda_urdf, self.T_world_base) + path = Path(rospack.get_path("active_grasp")) + panda_urdf = path / "assets/urdfs/franka/panda_arm_hand.urdf" + self.arm = BtPandaArm(panda_urdf) self.gripper = BtPandaGripper(self.arm) - self.model = Model(self.panda_urdf, self.arm.base_frame, self.arm.ee_frame) + self.model = Model(panda_urdf, self.arm.base_frame, self.arm.ee_frame) self.camera = BtCamera(320, 240, 1.047, 0.1, 1.0, self.arm.uid, 11) def seed(self, seed): self.rng = np.random.default_rng(seed) def reset(self): - self.remove_all_objects() + self.scene.reset(rng=self.rng) self.set_initial_arm_configuration() - self.load_random_object_arrangement() uid = self.select_target() - return self.get_target_bbox(uid) - - def step(self): - p.stepSimulation() + bbox = self.get_target_bbox(uid) + return bbox def set_initial_arm_configuration(self): q = [ @@ -79,6 +66,40 @@ class Simulation: p.resetJointState(self.arm.uid, 9, 0.04, 0) p.resetJointState(self.arm.uid, 10, 0.04, 0) + def select_target(self): + _, _, mask = self.camera.get_image() + uids, counts = np.unique(mask, return_counts=True) + mask = np.isin(uids, self.scene.object_uids) # remove ids of the floor, etc + uids, counts = uids[mask], counts[mask] + target_uid = uids[np.argmin(counts)] + p.changeVisualShape(target_uid, -1, rgbaColor=[1, 0, 0, 1]) + return target_uid + + def get_target_bbox(self, uid): + aabb_min, aabb_max = p.getAABB(uid) + return AABBox(aabb_min, aabb_max) + + def step(self): + p.stepSimulation() + + +class Scene: + def __init__(self): + self.urdfs_path = Path(rospack.get_path("vgn")) / "assets/urdfs" + self.origin = None + self.support_uid = -1 + self.object_uids = [] + + def load_support(self, urdf, ori, pos, scale): + uid = p.loadURDF(str(urdf), pos, ori.as_quat(), globalScaling=scale) + self.support_uid = uid + return uid + + def remove_support(self): + if self.support_uid != -1: + p.removeBody(self.support_uid) + self.support_uid = -1 + def load_object(self, urdf, ori, pos, scale=1.0): uid = p.loadURDF(str(urdf), pos, ori.as_quat(), globalScaling=scale) self.object_uids.append(uid) @@ -92,22 +113,44 @@ class Simulation: for uid in list(self.object_uids): self.remove_object(uid) - def load_random_object_arrangement(self, attempts=10): - origin = np.r_[self.origin[:2], 0.625] - scale = 0.8 - urdfs = self.rng.choice(self.urdfs, 4) + def reset(self, rng): + self.remove_support() + self.remove_all_objects() + self.load(rng) + + def load(self, rng): + raise NotImplementedError + + +def find_urdfs(root): + # Scans a dir for URDF assets + return [str(f) for f in root.iterdir() if f.suffix == ".urdf"] + + +class RandomScene(Scene): + def __init__(self): + super().__init__() + self.center = np.r_[[0.5, 0.0, 0.2]] + self.length = 0.3 + self.origin = self.center - np.r_[0.5 * self.length, 0.5 * self.length, 0.0] + self.object_urdfs = find_urdfs(self.urdfs_path / "packed" / "test") + self.support_urdf = self.urdfs_path / "setup/plane.urdf" + + def load(self, rng, attempts=10): + self.load_support(self.support_urdf, Rotation.identity(), self.center, 0.3) + + urdfs, scale = rng.choice(self.object_urdfs, 4), 0.8 for urdf in urdfs: - # Load the object - uid = self.load_object(urdf, Rotation.identity(), [0.0, 0.0, 0.0], scale) + uid = self.load_object(urdf, Rotation.identity(), np.zeros(3), scale) lower, upper = p.getAABB(uid) z_offset = 0.5 * (upper[2] - lower[2]) + 0.002 state_id = p.saveState() for _ in range(attempts): - # Try to place the object without collision - ori = Rotation.from_rotvec([0.0, 0.0, self.rng.uniform(0, 2 * np.pi)]) - offset = np.r_[self.rng.uniform(0.2, 0.8, 2) * self.length, z_offset] - p.resetBasePositionAndOrientation(uid, origin + offset, ori.as_quat()) - self.step() + # Try to place and check for collisions + ori = Rotation.from_rotvec([0.0, 0.0, rng.uniform(0, 2 * np.pi)]) + pos = np.r_[rng.uniform(0.2, 0.8, 2) * self.length, z_offset] + p.resetBasePositionAndOrientation(uid, self.origin + pos, ori.as_quat()) + p.stepSimulation() if not p.getContactPoints(uid): break else: @@ -116,18 +159,14 @@ class Simulation: # No placement found, remove the object self.remove_object(uid) - def select_target(self): - _, _, mask = self.camera.get_image() - uids, counts = np.unique(mask, return_counts=True) - mask = np.isin(uids, self.object_uids) # remove ids of the floor, table, etc - uids, counts = uids[mask], counts[mask] - target_uid = uids[np.argmin(counts)] - p.changeVisualShape(target_uid, -1, rgbaColor=[1, 0, 0, 1]) - return target_uid - def get_target_bbox(self, uid): - aabb_min, aabb_max = p.getAABB(uid) - # Transform the coordinates to base_frame - aabb_min = np.array(aabb_min) - self.T_world_base.translation - aabb_max = np.array(aabb_max) - self.T_world_base.translation - return AABBox(aabb_min, aabb_max) +scenes = { + "random": RandomScene, +} + + +def get_scene(scene_id): + if scene_id in scenes: + return scenes[scene_id]() + else: + raise ValueError("Unknown scene {}.".format(scene_id)) diff --git a/cfg/active_grasp.yaml b/cfg/active_grasp.yaml index 195aa23..4e4993f 100644 --- a/cfg/active_grasp.yaml +++ b/cfg/active_grasp.yaml @@ -2,6 +2,7 @@ bt_sim: gui: False cam_noise: True calib_error: True + scene: random grasp_controller: base_frame_id: panda_link0 diff --git a/package.xml b/package.xml index 6880fa4..7986b46 100644 --- a/package.xml +++ b/package.xml @@ -13,5 +13,7 @@ message_runtime + geometry_msgs rospy + std_msgs diff --git a/scripts/bt_sim_node.py b/scripts/bt_sim_node.py index ef95aeb..736c7c2 100755 --- a/scripts/bt_sim_node.py +++ b/scripts/bt_sim_node.py @@ -21,8 +21,9 @@ from robot_helpers.ros.conversions import * class BtSimNode: def __init__(self): - self.gui = rospy.get_param("~gui", True) - self.sim = Simulation(gui=self.gui) + gui = rospy.get_param("~gui") + scene_id = rospy.get_param("~scene") + self.sim = Simulation(gui=gui, scene_id=scene_id) self.init_plugins() self.advertise_services()