def load_box(pybullet_client: bullet_client.BulletClient, half_extents: Sequence[float] = (1, 1, 1), position: Sequence[float] = (0, 0, 0), orientation: Sequence[float] = (0, 0, 0, 1), rgba_color: Sequence[float] = (0.3, 0.3, 0.3, 1), mass: float = 0) -> int: """Loads a visible and tangible box. Args: half_extents: Half the dimension of the box in meters in each direction. position: Global coordinates of the center of the box. orientation: As a quaternion. rgba_color: The color and transparency of the box, each in the range [0,1]. Defaults to opaque gray. mass: Mass in kg. Mass 0 fixes the box in place. Returns: Unique integer referring to the loaded box. """ col_box_id = pybullet_client.createCollisionShape(pybullet.GEOM_BOX, halfExtents=half_extents) visual_box_id = pybullet_client.createVisualShape(pybullet.GEOM_BOX, halfExtents=half_extents, rgbaColor=rgba_color) return pybullet_client.createMultiBody(baseMass=mass, baseCollisionShapeIndex=col_box_id, baseVisualShapeIndex=visual_box_id, basePosition=position, baseOrientation=orientation)
class ModularEnv(gym.Env): """Abstract modular environment""" metadata = {'render.modes': ['human']} def __init__(self, terrain=None): # Create pybullet interfaces self.client = BulletClient(connection_mode=pyb.DIRECT) self._last_render = time.time() # Setup information needed for simulation self.dt = 1. / 240. self.client.setAdditionalSearchPath(ASSET_PATH) self.terrain_type = terrain # Stored for user interactions self.morphology = None self.multi_id = None self._joint_ids = [] self._joints = [] # Used for user interaction: self._real_time = False # Run setup self.setup() def load_terrain(self): """Helper method to load terrain for ground""" if not self.terrain_type: return None elif type(self.terrain_type) is PNGTerrain: shape = self.client.createCollisionShape( shapeType=pyb.GEOM_HEIGHTFIELD, meshScale=self.terrain_type.scale, fileName=self.terrain_type.height_field) assert shape >= 0, "Could not load PNGTerrain shape" terrain = self.client.createMultiBody(0, shape) if self.terrain_type.texture: texture = self.client.loadTexture(self.terrain_type.texture) assert texture >= 0, "Could not load PNGTerrain texture" self.client.changeVisualShape(terrain, -1, textureUniqueId=texture) if self.terrain_type.position: self.client.resetBasePositionAndOrientation( terrain, self.terrain_type.position, [0, 0, 0, 1]) elif type(self.terrain_type) is ArrayTerrain: shape = self.client.createCollisionShape( shapeType=pyb.GEOM_HEIGHTFIELD, meshScale=self.terrain_type.scale, heightfieldTextureScaling=(self.terrain_type.rows - 1) / 2., heightfieldData=self.terrain_type.field.flatten(), numHeightfieldRows=self.terrain_type.rows, numHeightfieldColumns=self.terrain_type.cols) assert shape >= 0, "Could not create ArrayTerrain shape" terrain = self.client.createMultiBody(0, shape) self.client.resetBasePositionAndOrientation( terrain, [0, 0, 0], [0, 0, 0, 1]) assert terrain >= 0, "Could not load terrain" self.client.changeVisualShape(terrain, -1, rgbaColor=[1, 1, 1, 1]) return terrain def setup(self): """Helper method to initialize default environment""" # This is abstracted out from '__init__' because we need to do this # first time 'render' is called self.client.resetSimulation() self.client.setGravity(0, 0, -9.81) # Load ground plane for robots to walk on self.plane_id = self.client.loadURDF('plane/plane.urdf', useMaximalCoordinates=1) assert self.plane_id >= 0, "Could not load 'plane.urdf'" # Change dynamics parameters from: # https://github.com/bulletphysics/bullet3/blob/master/examples/pybullet/gym/pybullet_envs/deep_mimic/env/pybullet_deep_mimic_env.py#L45 self.client.changeDynamics(self.plane_id, -1, lateralFriction=0.9) # self.client.setPhysicsEngineParameter(numSolverIterations=10) # self.client.setPhysicsEngineParameter(minimumSolverIslandSize=100) self.client.setPhysicsEngineParameter(contactERP=0) # Extract time step for sleep during rendering self.dt = self.client.getPhysicsEngineParameters()['fixedTimeStep'] def close(self): self.client.disconnect() def reset(self, morphology=None, max_size=None): # Disable rendering during spawning self.client.configureDebugVisualizer(pyb.COV_ENABLE_RENDERING, 0) # Reset the environment by running all setup code again self.setup() # Reset internal state self.morphology = None self.multi_id = None self._joint_ids = [] self._joints = [] # Check method arguments for simple logic errors if morphology is None: raise TypeError("Morphology cannot be 'None'!") if max_size is not None and max_size < 1: raise ValueError("'max_size' must be larger than 0") # NOTE: Utilize deep copy here to avoid interfering with morphology, # this means that if we want to do structural changes we can do so on # our own private instance. Which is nice for maximum size etc. self.morphology = copy.deepcopy(morphology.root) # Mapping from module to PyBullet ID spawned_ids = {} # NOTE: We are using explicit queue handling here so that we can # ignore children of overlapping modules queue = deque([self.morphology.root]) while len(queue) > 0: module = queue.popleft() assert isinstance(module, Module3D), "{} does not inherit\ from 3D Module".format(type(module)) # Spawn module in world m_id = module.spawn(self.client) # Check if the module overlaps aabb_min, aabb_max = self.client.getAABB(m_id) # Check overlapping modules overlap = self.client.getOverlappingObjects(aabb_min, aabb_max) # NOTE: An object always collides with it self overlapping_modules = False if overlap: overlapping_modules = any( [u_id != m_id for u_id, _ in overlap]) # Check against plane aabb_min, aabb_max = self.client.getAABB(self.plane_id) plane_overlap = self.client.getOverlappingObjects( aabb_min, aabb_max) overlapping_plane = False if plane_overlap: overlapping_plane = any( [u_id != self.plane_id for u_id, _ in plane_overlap]) # If overlap is detected de-spawn module and continue if overlapping_modules or overlapping_plane: # Remove from simulation self.client.removeBody(m_id) # Remove from our private copy parent = module.parent if parent: del parent[module] else: raise RuntimeError( "Trying to remove root module due to collision!") continue # Add children to queue for processing queue.extend(module.children) # Add ID to spawned IDs so that we can remove them later spawned_ids[module] = m_id # Add joint if present if module.joint: self._joints.append(module.joint) # Check size constraints if max_size is not None and len(spawned_ids) >= max_size: # If we are above max desired spawn size drain queue and remove for module in queue: parent = module.parent if parent: del parent[module] else: raise RuntimeError("Trying to prune root link!") break # Create multi body builder and instantiate with morphological tree and # module IDs builder = MultiBodyBuilder(self.morphology, spawned_ids, self.client, flags=pyb.URDF_USE_SELF_COLLISION) # Before spawning multi body we remove all modules so far # NOTE: We must do that here after all modules have been spawned to not # interfere with collision detection for m_id in spawned_ids.values(): self.client.removeBody(m_id) # Instantiate the builder creating a multi body from the morphology multi_id = builder.build() assert multi_id >= 0, "Could not spawn multibody!" self.multi_id = multi_id for jid in range(self.client.getNumJoints(self.multi_id)): j_info = self.client.getJointInfo(self.multi_id, jid) if j_info[2] != pyb.JOINT_FIXED: self._joint_ids.append(jid) self.observation_space = gym.spaces.Box(-100., 100., shape=(3 * len(self._joints), ), dtype=np.float64) lows = np.repeat(-1.57, len(self._joints)) self.action_space = gym.spaces.Box(lows, -1. * lows, dtype=np.float64) # Load additional terrain if requested # NOTE: We load terrain after collision detection to avoid collisions # with terrain itself if self.load_terrain() is not None: # If a different terrain was loaded we remove the plane self.client.removeBody(self.plane_id) self.plane_id = None # Enable rendering here self.client.configureDebugVisualizer(pyb.COV_ENABLE_RENDERING, 1) return self.observation() def act(self, action): """Helper function to apply actions to robot""" assert action.shape[0] == len(self._joint_ids), 'Action not equal to\ number of joints' for j_id, cfg, act in zip(self._joint_ids, self._joints, action): # Create a local copy so we can delete from it l_cfg = cfg.copy() # If joint should be limited we check that here if 'limit' in l_cfg: act = max(l_cfg['limit'][0], min(l_cfg['limit'][1], act)) del l_cfg['limit'] # Extract type of 'target' for joint l_cfg[l_cfg['target']] = act del l_cfg['target'] self.client.setJointMotorControl2(self.multi_id, j_id, **l_cfg) def step(self, action): self.act(np.asarray(action)) self.client.stepSimulation() return self.observation(), self.reward(), False, {} def observation(self): """Get observation from environment The default observation is `numpy.concatenate([positions, velocities, torques])` for all movable joints.""" # If the body does not have any movable joints: if not self._joint_ids: return np.array([]) positions = [] velocities = [] torques = [] states = self.client.getJointStates(self.multi_id, self._joint_ids) for state in states: positions.append(state[0]) velocities.append(state[1]) torques.append(state[3]) return np.concatenate([positions, velocities, torques]) def reward(self): """Estimate current reward""" # Get state of root link pos, _ = self.client.getBasePositionAndOrientation(self.multi_id) # Extract position position = np.array(pos) # We are only interested in distance in (X, Y) position[-1] = 0.0 # Return the root's distance from World center return np.linalg.norm(position) def render(self, mode="human"): info = self.client.getConnectionInfo() if info['connectionMethod'] != pyb.GUI and mode == 'human': # Close current simulation and start new with GUI self.close() self.client = BulletClient(connection_mode=pyb.GUI) # Disable rendering during spawning self.client.configureDebugVisualizer(pyb.COV_ENABLE_RENDERING, 0) # Reset if morphology is set if self.morphology is not None: self.reset(self.morphology) else: self.setup() # Configure viewport self.client.configureDebugVisualizer(pyb.COV_ENABLE_GUI, 0) self.client.configureDebugVisualizer( pyb.COV_ENABLE_PLANAR_REFLECTION, 1) self.client.resetDebugVisualizerCamera(0.5, 50.0, -35.0, (0, 0, 0)) # Setup timing for correct sleeping when rendering self._last_render = time.time() elif info['connectionMethod'] == pyb.GUI: # Handle interaction with simulation self.handle_interaction() # Calculate time to sleep now = time.time() diff = now - self._last_render to_sleep = self.dt - diff if to_sleep > 0: time.sleep(to_sleep) self._last_render = now else: raise RuntimeError( "Unknown pybullet mode ({}) or render mode ({})".format( info['connectionMethod'], mode)) def handle_interaction(self): """Handle user interaction with simulation This function is called in the 'render' call and is only called if the GUI is visible""" keys = self.client.getKeyboardEvents() # If 'n' is pressed then we want to step the simulation next_key = ord('n') if next_key in keys and keys[next_key] & pyb.KEY_WAS_TRIGGERED: self.client.stepSimulation() # If space is pressed we start/stop real-time simulation space = ord(' ') if space in keys and keys[space] & pyb.KEY_WAS_TRIGGERED: self._real_time = not self._real_time self.client.setRealTimeSimulation(self._real_time) # if 'r' is pressed we restart simulation r = ord('r') if r in keys and keys[r] & pyb.KEY_WAS_TRIGGERED: real_time = self._real_time self.client.setRealTimeSimulation(False) self.reset(self.morphology) self.client.setRealTimeSimulation(real_time)