def PendulumExample(): builder = DiagramBuilder() plant = builder.AddSystem(PendulumPlant()) scene_graph = builder.AddSystem(SceneGraph()) PendulumGeometry.AddToBuilder(builder, plant.get_state_output_port(), scene_graph) MeshcatVisualizerCpp.AddToBuilder( builder, scene_graph, meshcat, MeshcatVisualizerParams(publish_period=np.inf)) builder.ExportInput(plant.get_input_port(), "torque") # Note: The Pendulum-v0 in gym outputs [cos(theta), sin(theta), thetadot] builder.ExportOutput(plant.get_state_output_port(), "state") # Add a camera (I have sugar for this in the manip repo) X_WC = RigidTransform(RotationMatrix.MakeXRotation(-np.pi / 2), [0, -1.5, 0]) rgbd = AddRgbdSensor(builder, scene_graph, X_WC) builder.ExportOutput(rgbd.color_image_output_port(), "camera") diagram = builder.Build() simulator = Simulator(diagram) def reward(system, context): plant_context = plant.GetMyContextFromRoot(context) state = plant_context.get_continuous_state_vector() u = plant.get_input_port().Eval(plant_context)[0] theta = state[0] % 2 * np.pi # Wrap to 2*pi theta_dot = state[1] return (theta - np.pi)**2 + 0.1 * theta_dot**2 + 0.001 * u max_torque = 3 env = DrakeGymEnv(simulator, time_step=0.05, action_space=gym.spaces.Box(low=np.array([-max_torque]), high=np.array([max_torque])), observation_space=gym.spaces.Box( low=np.array([-np.inf, -np.inf]), high=np.array([np.inf, np.inf])), reward=reward, render_rgb_port_id="camera") check_env(env) if show: env.reset() image = env.render(mode='rgb_array') fig, ax = plt.subplots() ax.imshow(image) plt.show()
def plant_system(self): ''' Implements the plant_system method in DrakeEnv by constructing a RigidBodyPlant ''' # Add Systems builder = DiagramBuilder() self.scene_graph = builder.AddSystem(SceneGraph()) self.mbp = builder.AddSystem(MultibodyPlant()) # Load the model from the file AddModelFromSdfFile(file_name=self.model_path, plant=self.mbp, scene_graph=self.scene_graph) self.mbp.AddForceElement(UniformGravityFieldElement([0, 0, -9.81])) self.mbp.Finalize(self.scene_graph) assert self.mbp.geometry_source_is_registered() # Visualizer must be initialized after Finalize() and before CreateDefaultContext() self.init_visualizer() builder.Connect( self.mbp.get_geometry_poses_output_port(), self.scene_graph.get_source_pose_port(self.mbp.get_source_id())) # Expose the inputs and outputs and build the diagram self._input_port_index_action = builder.ExportInput( self.mbp.get_actuation_input_port()) self._output_port_index_state = builder.ExportOutput( self.mbp.get_continuous_state_output_port()) self.diagram = builder.Build() self._output = self.mbp.AllocateOutput() return self.diagram
def MakeManipulationStation(time_step=0.002, plant_setup_callback=None, camera_prefix="camera"): """ Sets up a manipulation station with the iiwa + wsg + controllers [+ cameras]. Cameras will be added to each body with a name starting with "camera". Args: time_step: the standard MultibodyPlant time step. plant_setup_callback: should be a python callable that takes one argument: `plant_setup_callback(plant)`. It will be called after the iiwa and WSG are added to the plant, but before the plant is finalized. Use this to add additional geometry. camera_prefix: Any bodies in the plant (created during the plant_setup_callback) starting with this prefix will get a camera attached. """ builder = DiagramBuilder() # Add (only) the iiwa, WSG, and cameras to the scene. plant, scene_graph = AddMultibodyPlantSceneGraph(builder, time_step=time_step) iiwa = AddIiwa(plant) wsg = AddWsg(plant, iiwa) if plant_setup_callback: plant_setup_callback(plant) plant.Finalize() num_iiwa_positions = plant.num_positions(iiwa) # I need a PassThrough system so that I can export the input port. iiwa_position = builder.AddSystem(PassThrough(num_iiwa_positions)) builder.ExportInput(iiwa_position.get_input_port(), "iiwa_position") builder.ExportOutput(iiwa_position.get_output_port(), "iiwa_position_command") # Export the iiwa "state" outputs. demux = builder.AddSystem( Demultiplexer(2 * num_iiwa_positions, num_iiwa_positions)) builder.Connect(plant.get_state_output_port(iiwa), demux.get_input_port()) builder.ExportOutput(demux.get_output_port(0), "iiwa_position_measured") builder.ExportOutput(demux.get_output_port(1), "iiwa_velocity_estimated") builder.ExportOutput(plant.get_state_output_port(iiwa), "iiwa_state_estimated") # Make the plant for the iiwa controller to use. controller_plant = MultibodyPlant(time_step=time_step) controller_iiwa = AddIiwa(controller_plant) AddWsg(controller_plant, controller_iiwa, welded=True) controller_plant.Finalize() # Add the iiwa controller iiwa_controller = builder.AddSystem( InverseDynamicsController(controller_plant, kp=[100] * num_iiwa_positions, ki=[1] * num_iiwa_positions, kd=[20] * num_iiwa_positions, has_reference_acceleration=False)) iiwa_controller.set_name("iiwa_controller") builder.Connect(plant.get_state_output_port(iiwa), iiwa_controller.get_input_port_estimated_state()) # Add in the feed-forward torque adder = builder.AddSystem(Adder(2, num_iiwa_positions)) builder.Connect(iiwa_controller.get_output_port_control(), adder.get_input_port(0)) # Use a PassThrough to make the port optional (it will provide zero values # if not connected). torque_passthrough = builder.AddSystem(PassThrough([0] * num_iiwa_positions)) builder.Connect(torque_passthrough.get_output_port(), adder.get_input_port(1)) builder.ExportInput(torque_passthrough.get_input_port(), "iiwa_feedforward_torque") builder.Connect(adder.get_output_port(), plant.get_actuation_input_port(iiwa)) # Add discrete derivative to command velocities. desired_state_from_position = builder.AddSystem( StateInterpolatorWithDiscreteDerivative( num_iiwa_positions, time_step, suppress_initial_transient=True)) desired_state_from_position.set_name("desired_state_from_position") builder.Connect(desired_state_from_position.get_output_port(), iiwa_controller.get_input_port_desired_state()) builder.Connect(iiwa_position.get_output_port(), desired_state_from_position.get_input_port()) # Export commanded torques. builder.ExportOutput(adder.get_output_port(), "iiwa_torque_commanded") builder.ExportOutput(adder.get_output_port(), "iiwa_torque_measured") builder.ExportOutput(plant.get_generalized_contact_forces_output_port(iiwa), "iiwa_torque_external") # Wsg controller. wsg_controller = builder.AddSystem(SchunkWsgPositionController()) wsg_controller.set_name("wsg_controller") builder.Connect(wsg_controller.get_generalized_force_output_port(), plant.get_actuation_input_port(wsg)) builder.Connect(plant.get_state_output_port(wsg), wsg_controller.get_state_input_port()) builder.ExportInput(wsg_controller.get_desired_position_input_port(), "wsg_position") builder.ExportInput(wsg_controller.get_force_limit_input_port(), "wsg_force_limit") wsg_mbp_state_to_wsg_state = builder.AddSystem( MakeMultibodyStateToWsgStateSystem()) builder.Connect(plant.get_state_output_port(wsg), wsg_mbp_state_to_wsg_state.get_input_port()) builder.ExportOutput(wsg_mbp_state_to_wsg_state.get_output_port(), "wsg_state_measured") builder.ExportOutput(wsg_controller.get_grip_force_output_port(), "wsg_force_measured") # Cameras. AddRgbdSensors(builder, plant, scene_graph, model_instance_prefix=camera_prefix) # Export "cheat" ports. builder.ExportOutput(scene_graph.get_query_output_port(), "geometry_query") builder.ExportOutput(plant.get_contact_results_output_port(), "contact_results") builder.ExportOutput(plant.get_state_output_port(), "plant_continuous_state") builder.ExportOutput(plant.get_body_poses_output_port(), "body_poses") diagram = builder.Build() diagram.set_name("ManipulationStation") return diagram
def make_box_flipup(generator, observations="state", meshcat=None, time_limit=10): builder = DiagramBuilder() plant, scene_graph = AddMultibodyPlantSceneGraph(builder, time_step=0.001) # TODO(russt): randomize parameters. box = AddPlanarBinAndSimpleBox(plant) finger = AddPointFinger(plant) plant.Finalize() plant.set_name("plant") SetTransparency(scene_graph, alpha=0.5, source_id=plant.get_source_id()) controller_plant = MultibodyPlant(time_step=0.005) AddPointFinger(controller_plant) if meshcat: MeshcatVisualizerCpp.AddToBuilder(builder, scene_graph, meshcat) meshcat.Set2dRenderMode(xmin=-.35, xmax=.35, ymin=-0.1, ymax=0.3) ContactVisualizer.AddToBuilder( builder, plant, meshcat, ContactVisualizerParams(radius=0.005, newtons_per_meter=40.0)) # Use the controller plant to visualize the set point geometry. controller_scene_graph = builder.AddSystem(SceneGraph()) controller_plant.RegisterAsSourceForSceneGraph(controller_scene_graph) SetColor(controller_scene_graph, color=[1.0, 165.0 / 255, 0.0, 1.0], source_id=controller_plant.get_source_id()) controller_vis = MeshcatVisualizerCpp.AddToBuilder( builder, controller_scene_graph, meshcat, MeshcatVisualizerParams(prefix="controller")) controller_vis.set_name("controller meshcat") controller_plant.Finalize() # Stiffness control. (For a point finger with unit mass, the # InverseDynamicsController is identical) N = controller_plant.num_positions() kp = [100] * N ki = [1] * N kd = [2 * np.sqrt(kp[0])] * N controller = builder.AddSystem( InverseDynamicsController(controller_plant, kp, ki, kd, False)) builder.Connect(plant.get_state_output_port(finger), controller.get_input_port_estimated_state()) actions = builder.AddSystem(PassThrough(N)) positions_to_state = builder.AddSystem(Multiplexer([N, N])) builder.Connect(actions.get_output_port(), positions_to_state.get_input_port(0)) zeros = builder.AddSystem(ConstantVectorSource([0] * N)) builder.Connect(zeros.get_output_port(), positions_to_state.get_input_port(1)) builder.Connect(positions_to_state.get_output_port(), controller.get_input_port_desired_state()) builder.Connect(controller.get_output_port(), plant.get_actuation_input_port()) if meshcat: positions_to_poses = builder.AddSystem( MultibodyPositionToGeometryPose(controller_plant)) builder.Connect( positions_to_poses.get_output_port(), controller_scene_graph.get_source_pose_port( controller_plant.get_source_id())) builder.ExportInput(actions.get_input_port(), "actions") if observations == "state": builder.ExportOutput(plant.get_state_output_port(), "observations") # TODO(russt): Add 'time', and 'keypoints' else: raise ValueError("observations must be one of ['state']") class RewardSystem(LeafSystem): def __init__(self): LeafSystem.__init__(self) self.DeclareVectorInputPort("box_state", 6) self.DeclareVectorInputPort("finger_state", 4) self.DeclareVectorInputPort("actions", 2) self.DeclareVectorOutputPort("reward", 1, self.CalcReward) def CalcReward(self, context, output): box_state = self.get_input_port(0).Eval(context) finger_state = self.get_input_port(1).Eval(context) actions = self.get_input_port(2).Eval(context) angle_from_vertical = (box_state[2] % np.pi) - np.pi / 2 cost = 2 * angle_from_vertical**2 # box angle cost += 0.1 * box_state[5]**2 # box velocity effort = actions - finger_state[:2] cost += 0.1 * effort.dot(effort) # effort # finger velocity cost += 0.1 * finger_state[2:].dot(finger_state[2:]) # Add 10 to make rewards positive (to avoid rewarding simulator # crashes). output[0] = 10 - cost reward = builder.AddSystem(RewardSystem()) builder.Connect(plant.get_state_output_port(box), reward.get_input_port(0)) builder.Connect(plant.get_state_output_port(finger), reward.get_input_port(1)) builder.Connect(actions.get_output_port(), reward.get_input_port(2)) builder.ExportOutput(reward.get_output_port(), "reward") # Set random state distributions. uniform_random = Variable(name="uniform_random", type=Variable.Type.RANDOM_UNIFORM) box_joint = plant.GetJointByName("box") x, y = box_joint.get_default_translation() box_joint.set_random_pose_distribution([.2 * uniform_random - .1 + x, y], 0) diagram = builder.Build() simulator = Simulator(diagram) # Termination conditions: def monitor(context): if context.get_time() > time_limit: return EventStatus.ReachedTermination(diagram, "time limit") return EventStatus.Succeeded() simulator.set_monitor(monitor) return simulator