import os import numpy as np import pybullet as p import subprocess as sp import time from pyquaternion import Quaternion from collections import namedtuple from itertools import groupby from math_utils import Transform import utils # handle to a simulated rigid body Body = namedtuple('Body', ['id', 'bid', 'vid', 'cid', 'static']) # a body-body contact record Contact = namedtuple( 'Contact', ['flags', 'idA', 'idB', 'linkIndexA', 'linkIndexB', 'positionOnAInWS', 'positionOnBInWS', 'contactNormalOnBInWS', 'distance', 'normalForce'] ) # ray intersection record Intersection = namedtuple('Intersection', ['id', 'linkIndex', 'ray_fraction', 'position', 'normal']) class Simulator: def __init__(self, mode='direct', bullet_server_binary=None, data_dir_base=None, verbose=False): self._mode = mode self._verbose = verbose module_dir = os.path.dirname(os.path.abspath(__file__)) data_root_dir = utils.get_data_root_dir() if data_dir_base: self._data_dir_base = data_dir_base else: self._data_dir_base = os.path.join(data_root_dir, 'suncg_data') if bullet_server_binary: self._bullet_server_binary = bullet_server_binary else: self._bullet_server_binary = os.path.join(module_dir, '..', 'bullet_shared_memory_server') self._obj_id_to_body = {} self._bid_to_body = {} self._pid = None self._bullet_server = None self.connect() def connect(self): # disconnect and kill existing servers if self._pid: p.disconnect(physicsClientId=self._pid) self._pid = None if self._bullet_server: print(f'Restarting by killing bullet server pid={self._bullet_server.pid}...') self._bullet_server.kill() time.sleep(1) # seems necessary to prevent deadlock on re-connection attempt self._bullet_server = None # reconnect to appropriate server type if self._mode == 'gui': self._pid = p.connect(p.GUI) elif self._mode == 'direct': self._pid = p.connect(p.DIRECT) elif self._mode == 'shared_memory': print(f'Restarting bullet server process...') self._bullet_server = sp.Popen([self._bullet_server_binary]) time.sleep(1) # seems necessary to prevent deadlock on connection attempt self._pid = p.connect(p.SHARED_MEMORY) else: raise RuntimeError(f'Unknown simulator server mode={self._mode}') # reset and initialize gui if needed p.resetSimulation(physicsClientId=self._pid) if self._mode == 'gui': p.configureDebugVisualizer(p.COV_ENABLE_Y_AXIS_UP, 1, physicsClientId=self._pid) p.configureDebugVisualizer(p.COV_ENABLE_RGB_BUFFER_PREVIEW, 0, physicsClientId=self._pid) p.configureDebugVisualizer(p.COV_ENABLE_DEPTH_BUFFER_PREVIEW, 0, physicsClientId=self._pid) p.configureDebugVisualizer(p.COV_ENABLE_SEGMENTATION_MARK_PREVIEW, 0, physicsClientId=self._pid) # disable rendering during loading -> much faster p.configureDebugVisualizer(p.COV_ENABLE_RENDERING, 0, physicsClientId=self._pid) def __del__(self): if self._bullet_server: print(f'Process terminating. Killing bullet server pid={self._bullet_server.pid}...') self._bullet_server.kill() def run(self): if self._mode == 'gui': # infinite ground plane and gravity # plane_cid = p.createCollisionShape(p.GEOM_PLANE, planeNormal=[0, 1, 0]) # plane_bid = p.createMultiBody(baseMass=0, baseCollisionShapeIndex=plane_cid) p.setGravity(0, -10, 0, physicsClientId=self._pid) p.setRealTimeSimulation(1, physicsClientId=self._pid) self.set_gui_rendering(enabled=True) while True: contacts = self.get_contacts(include_collision_with_static=True) if self._verbose: print(f'#contacts={len(contacts)}, contact_pairs={contacts.keys()}') else: self.step() contacts = self.get_contacts(include_collision_with_static=True) if self._verbose: for (k, c) in contacts.items(): cp = self.get_closest_point(obj_id_a=c.idA, obj_id_b=c.idB) print(f'contact pair={k} record={cp}') print(f'#contacts={len(contacts)}, contact_pairs={contacts.keys()}') def set_gui_rendering(self, enabled): if not self._mode == 'gui': return False center = np.array([0.0, 0.0, 0.0]) num_obj = 0 for obj_id in self._obj_id_to_body.keys(): pos, _ = self.get_state(obj_id) if not np.allclose(pos, [0, 0, 0]): # try to ignore room object "identity" transform num_obj += 1 center += pos center /= num_obj p.resetDebugVisualizerCamera(cameraDistance=10.0, cameraYaw=45.0, cameraPitch=-30.0, cameraTargetPosition=center, physicsClientId=self._pid) p.configureDebugVisualizer(p.COV_ENABLE_RENDERING, 1 if enabled else 0, physicsClientId=self._pid) return enabled def add_mesh(self, obj_id, obj_file, transform, vis_mesh_file=None, static=False): if static: cid = p.createCollisionShape(p.GEOM_MESH, fileName=obj_file, meshScale=transform.scale, flags=p.GEOM_FORCE_CONCAVE_TRIMESH, physicsClientId=self._pid) else: cid = p.createCollisionShape(p.GEOM_MESH, fileName=obj_file, meshScale=transform.scale, physicsClientId=self._pid) vid = -1 if vis_mesh_file: vid = p.createVisualShape(p.GEOM_MESH, fileName=vis_mesh_file, meshScale=transform.scale, physicsClientId=self._pid) rot_q = np.roll(transform.rotation.elements, -1) # w,x,y,z -> x,y,z,w (which pybullet expects) mass = 0 if static else 1 bid = p.createMultiBody(baseMass=mass, baseCollisionShapeIndex=cid, baseVisualShapeIndex=vid, basePosition=transform.translation, baseOrientation=rot_q, physicsClientId=self._pid) body = Body(id=obj_id, bid=bid, vid=vid, cid=cid, static=static) self._obj_id_to_body[obj_id] = body self._bid_to_body[bid] = body return body def add_box(self, obj_id, half_extents, transform, static=False): cid = p.createCollisionShape(p.GEOM_BOX, halfExtents=half_extents, physicsClientId=self._pid) rot_q = np.roll(transform.rotation.elements, -1) # w,x,y,z -> x,y,z,w (which pybullet expects) mass = 0 if static else 1 bid = p.createMultiBody(baseMass=mass, baseCollisionShapeIndex=cid, basePosition=transform.translation, baseOrientation=rot_q, physicsClientId=self._pid) body = Body(id=obj_id, bid=bid, vid=-1, cid=cid, static=static) self._obj_id_to_body[obj_id] = body self._bid_to_body[bid] = body return body # House-specific functions def add_object(self, node, create_vis_mesh=False, static=False): model_id = node.modelId.replace('_mirror', '') # TODO: need to otherwise account for mirror? object_dir = os.path.join(self._data_dir_base, 'object') basename = f'{object_dir}/{model_id}/{model_id}' vis_obj_filename = f'{basename}.obj' if create_vis_mesh else None col_obj_filename = f'{basename}.vhacd.obj' if not os.path.exists(col_obj_filename): print('WARNING: collision mesh {col_obj_filename} unavailable, using visual mesh instead.') col_obj_filename = f'{basename}.obj' return self.add_mesh(obj_id=node.id, obj_file=col_obj_filename, transform=Transform.from_node(node), vis_mesh_file=vis_obj_filename, static=static) def add_wall(self, node): h = node['height'] p0 = np.transpose(np.matrix(node['points'][0])) p1 = np.transpose(np.matrix(node['points'][1])) c = (p0 + p1) * 0.5 c[1] = h * 0.5 dp = p1 - p0 dp_l = np.linalg.norm(dp) dp = dp / dp_l angle = np.arccos(dp[0]) rot_q = Quaternion(axis=[0, 1, 0], radians=angle) half_extents = np.array([dp_l, h, node['depth']]) * 0.5 return self.add_box(obj_id=node['id'], half_extents=half_extents, transform=Transform(translation=c, rotation=rot_q), static=True) def add_room(self, node, wall=True, floor=True, ceiling=False): def add_architecture(n, obj_file, suffix): return self.add_mesh(obj_id=n.id + suffix, obj_file=obj_file, transform=Transform(), vis_mesh_file=None, static=True) room_id = node.modelId room_dir = os.path.join(self._data_dir_base, 'room') basename = f'{room_dir}/{node.house_id}/{room_id}' body_ids = [] if wall: body_wall = add_architecture(node, f'{basename}w.obj', '') # treat walls as room (=room.id, no suffix) body_ids.append(body_wall) if floor: body_floor = add_architecture(node, f'{basename}f.obj', 'f') body_ids.append(body_floor) if ceiling: body_ceiling = add_architecture(node, f'{basename}c.obj', 'c') body_ids.append(body_ceiling) return body_ids def add_house(self, house, no_walls=False, no_ceil=False, no_floor=False, use_separate_walls=False, static=False): for node in house.nodes: if not node.valid: continue if not hasattr(node, 'body'): node.body = None if node.type == 'Object': node.body = self.add_object(node, static=static) if node.type == 'Room': ceil = False if no_ceil else not (hasattr(node, 'hideCeiling') and node.hideCeiling == 1) wall = False if (no_walls or use_separate_walls) else not (hasattr(node, 'hideWalls') and node.hideWalls == 1) floor = False if no_floor else not (hasattr(node, 'hideFloor') and node.hideFloor == 1) node.body = self.add_room(node, wall=wall, floor=floor, ceiling=ceil) if node.type == 'Box': half_widths = list(map(lambda x: 0.5 * x, node.dimensions)) node.body = self.add_box(obj_id=node.id, half_extents=half_widths, transform=Transform.from_node(node), static=static) if use_separate_walls and not no_walls: for wall in house.walls: wall['body'] = self.add_wall(wall) def add_house_room_only(self, house, room, no_walls=False, no_ceil=True, no_floor=False, use_separate_walls=False, only_architecture=False, static=False): #walls, ceil, floor logic not fleshed out due to current limited use case room_node = [node for node in house.nodes if node.id == room.id] if len(room_node) < 1: raise Exception("Missing Room") if only_architecture: house.nodes = room_node else: house.nodes = [node for node in room.nodes] house.nodes.append(room_node[0]) for node in house.nodes: if not node.valid: continue if not hasattr(node, 'body'): node.body = None if node.type == 'Object': node.body = self.add_object(node, static=static) if node.type == 'Room': ceil = False if no_ceil else not (hasattr(node, 'hideCeiling') and node.hideCeiling == 1) wall = False if (no_walls or use_separate_walls) else not (hasattr(node, 'hideWalls') and node.hideWalls == 1) floor = False if no_floor else not (hasattr(node, 'hideFloor') and node.hideFloor == 1) node.body = self.add_room(node, wall=wall, floor=floor, ceiling=ceil) if node.type == 'Box': half_widths = list(map(lambda x: 0.5 * x, node.dimensions)) node.body = self.add_box(obj_id=node.id, half_extents=half_widths, transform=Transform.from_node(node), static=static) if use_separate_walls and not no_walls: for wall in house.walls: wall['body'] = self.add_wall(wall) def remove(self, obj_id): body = self._obj_id_to_body[obj_id] p.removeBody(bodyUniqueId=body.bid, physicsClientId=self._pid) del self._obj_id_to_body[obj_id] del self._bid_to_body[body.bid] def set_state(self, obj_id, position, rotation_q): body = self._obj_id_to_body[obj_id] rot_q = np.roll(rotation_q.elements, -1) # w,x,y,z -> x,y,z,w (which pybullet expects) p.resetBasePositionAndOrientation(bodyUniqueId=body.bid, posObj=position, ornObj=rot_q, physicsClientId=self._pid) def get_state(self, obj_id): body = self._obj_id_to_body[obj_id] pos, q = p.getBasePositionAndOrientation(bodyUniqueId=body.bid, physicsClientId=self._pid) rotation = Quaternion(x=q[0], y=q[1], z=q[2], w=q[3]) return pos, rotation def step(self): p.stepSimulation(physicsClientId=self._pid) def reset(self): p.resetSimulation(physicsClientId=self._pid) self._obj_id_to_body = {} self._bid_to_body = {} def get_closest_point(self, obj_id_a, obj_id_b, max_distance=np.inf): """ Return record with distance between closest points between pair of nodes if within max_distance or None. """ bid_a = self._obj_id_to_body[obj_id_a].bid bid_b = self._obj_id_to_body[obj_id_b].bid cps = p.getClosestPoints(bodyA=bid_a, bodyB=bid_b, distance=max_distance, physicsClientId=self._pid) cp = None if len(cps) > 0: closest_points = self._convert_contacts(cps) cp = min(closest_points, key=lambda x: x.distance) del cps # NOTE force garbage collection of pybullet objects return cp def get_contacts(self, obj_id_a=None, obj_id_b=None, only_closest_contact_per_pair=True, include_collision_with_static=True): """ Return all current contacts. When include_collision_with_statics is true, include contacts with static bodies """ bid_a = self._obj_id_to_body[obj_id_a].bid if obj_id_a else -1 bid_b = self._obj_id_to_body[obj_id_b].bid if obj_id_b else -1 cs = p.getContactPoints(bodyA=bid_a, bodyB=bid_b, physicsClientId=self._pid) contacts = self._convert_contacts(cs) del cs # NOTE force garbage collection of pybullet objects if not include_collision_with_static: def not_contact_with_static(c): static_a = self._obj_id_to_body[c.idA].static static_b = self._obj_id_to_body[c.idB].static return not static_a and not static_b contacts = filter(not_contact_with_static, contacts) # print(f'#all_contacts={len(all_contacts)} to #non_static_contacts={len(non_static_contacts)}') if only_closest_contact_per_pair: def bid_pair_key(x): return str(x.idA) + '_' + str(x.idB) contacts = sorted(contacts, key=bid_pair_key) min_dist_contact_by_pair = {} for k, g in groupby(contacts, key=bid_pair_key): min_dist_contact = min(g, key=lambda x: x.distance) min_dist_contact_by_pair[k] = min_dist_contact contacts = min_dist_contact_by_pair.values() # convert into dictionary of form (id_a, id_b) -> Contact contacts_dict = {} for c in contacts: key = (c.idA, c.idB) contacts_dict[key] = c return contacts_dict def _convert_contacts(self, contacts): out = [] for c in contacts: bid_a = c[1] bid_b = c[2] if bid_a not in self._bid_to_body or bid_b not in self._bid_to_body: continue id_a = self._bid_to_body[bid_a].id id_b = self._bid_to_body[bid_b].id o = Contact(flags=c[0], idA=id_a, idB=id_b, linkIndexA=c[3], linkIndexB=c[4], positionOnAInWS=c[5], positionOnBInWS=c[6], contactNormalOnBInWS=c[7], distance=c[8], normalForce=c[9]) out.append(o) return out def ray_test(self, from_pos, to_pos): hit = p.rayTest(rayFromPosition=from_pos, rayToPosition=to_pos, physicsClientId=self._pid) intersection = Intersection._make(*hit) del hit # NOTE force garbage collection of pybullet objects if intersection.id >= 0: # if intersection, replace bid with id intersection = intersection._replace(id=self._bid_to_body[intersection.id].id) return intersection if __name__ == '__main__': from data import House sim = Simulator(mode='gui', verbose=True) h = House(id_='d119e6e0bd567d923aea774c2a984bf0', include_arch_information=False) # h = House(index=5) sim.add_house(h, no_walls=False, no_ceil=True, use_separate_walls=False, static=False) sim.run()