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
