예제 #1
0
파일: drake_dot_sim.py 프로젝트: gizatt/dot
def setup_dot_diagram(builder, args):
    ''' Using an existing DiagramBuilder, adds a sim of the Dot
    robot. Args comes from argparse.

    The returned controller will need its first port connected to
    a setpoint source.'''

    print(
        "TODO: load in servo calibration dict to a servo calibration object that gets shared"
    )

    with open(args.yaml_path, "r") as f:
        config_dict = yaml.load(f, Loader=Loader)
    sdf_path = os.path.join(os.path.dirname(args.yaml_path),
                            config_dict["sdf_path"])
    plant, scene_graph = AddMultibodyPlantSceneGraph(builder, 0.0005)
    parser = Parser(plant)
    model = parser.AddModelFromFile(sdf_path)

    # Set initial pose floating above the ground.
    plant.SetDefaultFreeBodyPose(plant.GetBodyByName("body"),
                                 RigidTransform(p=[0., 0., 0.25]))
    if args.welded:
        plant.WeldFrames(plant.world_frame(),
                         plant.GetBodyByName("body").body_frame())
    else:
        add_ground(plant)

    plant.Finalize()

    controller = builder.AddSystem(ServoController(plant, config_dict))
    # Fixed control-rate controller with a low pass filter on its torque output.
    zoh = builder.AddSystem(
        ZeroOrderHold(period_sec=0.001, vector_size=controller.n_servos))
    filter = builder.AddSystem(
        FirstOrderLowPassFilter(time_constant=0.02, size=controller.n_servos))
    builder.Connect(plant.get_state_output_port(),
                    controller.get_input_port(1))
    builder.Connect(controller.get_output_port(0), zoh.get_input_port(0))
    builder.Connect(zoh.get_output_port(0), filter.get_input_port(0))
    builder.Connect(filter.get_output_port(0),
                    plant.get_actuation_input_port())
    return plant, scene_graph, controller
예제 #2
0
def main():
    builder = DiagramBuilder()
    plant, scene_graph = AddMultibodyPlantSceneGraph(builder, time_step=0.0)
    path = subprocess.check_output([
        'venv/share/drake/common/resource_tool',
        '-print_resource_path',
        'drake/examples/multibody/cart_pole/cart_pole.sdf',
    ]).strip()
    Parser(plant).AddModelFromFile(path)
    plant.Finalize()

    # Add to visualizer.
    DrakeVisualizer.AddToBuilder(builder, scene_graph)

    # Add controller.
    controller = builder.AddSystem(BalancingLQR(plant))
    builder.Connect(plant.get_state_output_port(),
                    controller.get_input_port(0))
    builder.Connect(controller.get_output_port(0),
                    plant.get_actuation_input_port())

    diagram = builder.Build()

    # Set up a simulator to run this diagram.
    simulator = Simulator(diagram)
    simulator.set_target_realtime_rate(1.0)
    context = simulator.get_mutable_context()

    # Simulate.
    duration = 8.0
    for i in range(5):
        initial_state = UPRIGHT_STATE + 0.1 * np.random.randn(4)
        print(f"Iteration: {i}. Initial: {initial_state}")
        context.SetContinuousState(initial_state)
        context.SetTime(0.0)
        simulator.Initialize()
        simulator.AdvanceTo(duration)
예제 #3
0
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
예제 #4
0
                    type=float,
                    help="Duration to run sim.",
                    default=10.0)
args = parser.parse_args()

builder = DiagramBuilder()
plant, scene_graph = AddMultibodyPlantSceneGraph(builder)

# Load the double pendulum from Universal Robot Description Format
parser = Parser(plant, scene_graph)
parser.AddModelFromFile(FindResource("double_pendulum/double_pendulum.urdf"))
plant.AddForceElement(UniformGravityFieldElement())
plant.Finalize()

controller = builder.AddSystem(Controller(plant, args.gravity))
builder.Connect(plant.get_state_output_port(), controller.get_input_port(0))
builder.Connect(controller.get_output_port(0),
                plant.get_actuation_input_port())

visualizer = builder.AddSystem(
    PlanarSceneGraphVisualizer(scene_graph, xlim=[-2.8, 2.8], ylim=[-2.8,
                                                                    2.8]))
builder.Connect(scene_graph.get_pose_bundle_output_port(),
                visualizer.get_input_port(0))

diagram = builder.Build()

# Set up a simulator to run this diagram
simulator = Simulator(diagram)
simulator.set_target_realtime_rate(1.0)
예제 #5
0
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
예제 #6
0
Hence, by running the following cell multiple times, you would actually try to wire the blocks in the diagram more than once.
This is not acceptable, and Drake will raise the error `RuntimeError: Input port is already wired.`
If you wish to modify the next cell and rerun the program to see the effects of your modification, you must rerun the cell where the `builder` is initialized first (i.e. the cell with the line `builder = DiagramBuilder()`).
"""

# instantiate controllers
inner_controller = builder.AddSystem(InnerController(vibrating_pendulum))
outer_controller = builder.AddSystem(OuterController(vibrating_pendulum, pendulum_torque))

# instantiate visualizer
visualizer = builder.AddSystem(
    PlanarSceneGraphVisualizer(scene_graph, xlim=[-2.5, 2.5], ylim=[-1.5, 2.5], show=False)
)

# logger that records the state trajectory during simulation
logger = LogOutput(vibrating_pendulum.get_state_output_port(), builder)

# mux block to squeeze the (base + pendulum) state and
# the outer control signal in a single cable
mux = builder.AddSystem(Multiplexer([4, 1]))

# selector that extracts the pendulum state from the state of the base and the pendulum
selector = builder.AddSystem(MatrixGain(np.array([
    [0, 1, 0, 0], # selects the angle of the pendulum
    [0, 0, 0, 1] # selects the angular velocity of the pendulum
])))

# (base + pendulum) state, outer control -> mux
builder.Connect(vibrating_pendulum.get_state_output_port(), mux.get_input_port(0))
builder.Connect(outer_controller.get_output_port(0), mux.get_input_port(1))
예제 #7
0
# set the operating point (vertical unstable equilibrium)
context = cartpole.CreateDefaultContext()
context.get_mutable_continuous_state_vector().SetFromVector(x_star)

# fix the input port to zero and get its index for the lqr function
cartpole.get_actuation_input_port().FixValue(context, [0])
input_i = cartpole.get_actuation_input_port().get_index()

# synthesize lqr controller directly from
# the nonlinear system and the operating point
lqr = LinearQuadraticRegulator(cartpole, context, Q, R, input_port_index=input_i)
lqr = builder.AddSystem(lqr)

# the following two lines are not needed here...
output_i = cartpole.get_state_output_port().get_index()
cartpole_lin = Linearize(cartpole, context, input_port_index=input_i, output_port_index=output_i)

# wire cart-pole and lqr
builder.Connect(cartpole.get_state_output_port(), lqr.get_input_port(0))
builder.Connect(lqr.get_output_port(0), cartpole.get_actuation_input_port())

# add a visualizer and wire it
visualizer = builder.AddSystem(
    PlanarSceneGraphVisualizer(scene_graph, xlim=[-3., 3.], ylim=[-1.2, 1.2], show=False)
)
builder.Connect(scene_graph.get_pose_bundle_output_port(), visualizer.get_input_port(0))

# finish building the block diagram
diagram = builder.Build()
예제 #8
0
parser.add_argument("-T", "--duration",
                    type=float,
                    help="Duration to run sim.",
                    default=10.0)
args = parser.parse_args()

builder = DiagramBuilder()
plant, scene_graph = AddMultibodyPlantSceneGraph(builder)

# Load the double pendulum from Universal Robot Description Format
parser = Parser(plant, scene_graph)
parser.AddModelFromFile(FindResource("double_pendulum/double_pendulum.urdf"))
plant.Finalize()

controller = builder.AddSystem(Controller(plant, args.gravity))
builder.Connect(plant.get_state_output_port(),
                controller.get_input_port(0))
builder.Connect(controller.get_output_port(0),
                plant.get_actuation_input_port())

visualizer = builder.AddSystem(PlanarSceneGraphVisualizer(scene_graph,
                                                          xlim=[-2.8, 2.8],
                                                          ylim=[-2.8, 2.8]))
builder.Connect(scene_graph.get_pose_bundle_output_port(),
                visualizer.get_input_port(0))

diagram = builder.Build()

# Set up a simulator to run this diagram
simulator = Simulator(diagram)
simulator.set_target_realtime_rate(1.0)
예제 #9
0
class iiwa_sys():
    def __init__(self,
                 builder,
                 dt=5e-4,
                 N=150,
                 params=None,
                 trj_decay=0.7,
                 x_w_cov=1e-5,
                 door_angle_ref=1.0,
                 visualize=False):
        self.plant_derivs = MultibodyPlant(time_step=dt)
        parser = Parser(self.plant_derivs)
        self.derivs_iiwa, _, _ = self.add_models(self.plant_derivs,
                                                 parser,
                                                 params=params)
        self.plant_derivs.Finalize()
        self.plant_derivs_context = self.plant_derivs.CreateDefaultContext()
        self.plant_derivs.get_actuation_input_port().FixValue(
            self.plant_derivs_context, [0., 0., 0., 0., 0., 0., 0.])
        null_force = BasicVector([0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0])
        self.plant_derivs.GetInputPort("applied_generalized_force").FixValue(
            self.plant_derivs_context, null_force)

        self.plant, scene_graph = AddMultibodyPlantSceneGraph(builder,
                                                              time_step=dt)
        parser = Parser(self.plant, scene_graph)
        self.iiwa, self.hinge, self.bushing = self.add_models(self.plant,
                                                              parser,
                                                              params=params)
        self.plant.Finalize(
        )  # Finalize will assign ports for compatibility w/ the scene_graph; could be cause of the issue w/ first order taylor.

        self.meshcat = ConnectMeshcatVisualizer(builder,
                                                scene_graph,
                                                zmq_url=zmq_url)

        self.sym_derivs = False  # If system should use symbolic derivatives; if false, autodiff
        self.custom_sim = False  # If rollouts should be gathered with sys.dyn() calls

        nq = self.plant.num_positions()
        nv = self.plant.num_velocities()
        self.n_x = nq + nv
        self.n_u = self.plant.num_actuators()
        self.n_y = self.plant.get_state_output_port(self.iiwa).size()

        self.N = N
        self.dt = dt
        self.decay = trj_decay
        self.V = 1e-2 * np.ones(self.n_y)
        self.W = np.concatenate((1e-7 * np.ones(nq), 1e-4 * np.ones(nv)))
        self.W0 = np.concatenate((1e-9 * np.ones(nq), 1e-6 * np.ones(nv)))
        self.x_w_cov = x_w_cov
        self.door_angle_ref = door_angle_ref

        self.q0 = np.array(
            [-3.12, -0.17, 0.52, -3.11, 1.22, -0.75, -1.56, 0.55])
        #self.q0 = np.array([-3.12, -0.27, 0.52, -3.11, 1.22, -0.75, -1.56, 0.55])
        self.x0 = np.concatenate((self.q0, np.zeros(nv)))
        self.door_index = None

        self.phi = {}

    def ports_init(self, context):
        self.plant_context = self.plant.GetMyMutableContextFromRoot(context)
        self.plant.SetPositionsAndVelocities(self.plant_context, self.x0)

        #door_angle = self.plant.GetPositionsFromArray(self.hinge, self.x0[:8])
        self.door_index = self.plant.GetJointByName(
            'right_door_hinge').position_start()

        null_force = BasicVector([0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0])
        self.plant.GetInputPort("applied_generalized_force").FixValue(
            self.plant_context, null_force)
        self.plant.get_actuation_input_port().FixValue(
            self.plant_context, [0., 0., 0., 0., 0., 0., 0.])
        self.W0[self.door_index] = self.x_w_cov

    def add_models(self, plant, parser, params=None):
        iiwa = parser.AddModelFromFile(
            "/home/hanikevi/drake/manipulation/models/iiwa_description/sdf/iiwa14_no_collision_no_grav.sdf"
        )
        plant.WeldFrames(plant.world_frame(),
                         plant.GetFrameByName("iiwa_link_0"))

        #box = Box(10., 10., 10.)
        #X_WBox = RigidTransform([0, 0, -5])
        #mu = 0.6
        #plant.RegisterCollisionGeometry(plant.world_body(), X_WBox, box, "ground", CoulombFriction(mu, mu))
        #plant.RegisterVisualGeometry(plant.world_body(), X_WBox, box, "ground", [.9, .9, .9, 1.0])
        #planar_joint_frame = plant.AddFrame(FixedOffsetFrame("planar_joint_frame", plant.world_frame(), RigidTransform(RotationMatrix.MakeXRotation(np.pi/2))))

        X_WCylinder = RigidTransform([-0.75, 0, 0.5])
        hinge = parser.AddModelFromFile(
            "/home/hanikevi/drake/examples/manipulation_station/models/simple_hinge.sdf"
        )
        plant.WeldFrames(plant.world_frame(), plant.GetFrameByName("base"),
                         X_WCylinder)
        #cupboard_door_spring = plant.AddForceElement(RevoluteSpring_[float](plant.GetJointByName("right_door_hinge"), nominal_angle = -0.4, stiffness = 10))
        if params is None:
            bushing = LinearBushingRollPitchYaw_[float](
                plant.GetFrameByName("iiwa_link_7"),
                plant.GetFrameByName("handle"),
                [50, 50, 50],  # Torque stiffness
                [2., 2., 2.],  # Torque damping
                [5e4, 5e4, 5e4],  # Linear stiffness
                [80, 80, 80],  # Linear damping
            )
        else:
            print('setting custom stiffnesses')
            bushing = LinearBushingRollPitchYaw_[float](
                plant.GetFrameByName("iiwa_link_7"),
                plant.GetFrameByName("handle"),
                [params['k4'], params['k5'], params['k6']],  # Torque stiffness
                [2, 2, 2],  # Torque damping
                [params['k1'], params['k2'], params['k3']],  # Linear stiffness
                [100, 100, 100],  # Linear damping
            )
        bushing_element = plant.AddForceElement(bushing)

        return iiwa, hinge, bushing

    def cost_stage(self, x, u):
        ctrl = 1e-5 * np.sum(u**2)
        pos = 15.0 * (x[self.door_index] - self.door_angle_ref)**2
        vel = 1e-5 * np.sum(x[8:]**2)
        return pos + ctrl + vel

    def cost_final(self, x):
        return 50 * (1.0 * (x[self.door_index] - self.door_angle_ref)**2 +
                     np.sum(2.5e-4 * x[8:]**2))

    def get_deriv(self, x, u):
        self.plant_derivs.SetPositionsAndVelocities(self.plant_derivs_context,
                                                    x)
        lin = FirstOrderTaylorApproximation(
            self.plant_derivs, self.plant_derivs_context,
            self.plant.get_actuation_input_port().get_index(),
            self.plant.get_state_output_port(self.iiwa).get_index())
        return lin.A(), lin.B(), lin.C()

    def get_param_deriv(self, x, u):
        # Using a closed-form solution; as currently DRAKE doesn't do support autodiff on parameters for LinearBusing.
        # Note dC is 0 - stiffness does not affect measurements of q, v
        W = self.plant.world_frame()
        I = self.plant.GetFrameByName("iiwa_link_7")
        H = self.plant.GetFrameByName("handle")
        self.plant.SetPositionsAndVelocities(self.plant_context, x)
        M = self.plant.CalcMassMatrixViaInverseDynamics(self.plant_context)
        Jac_I = self.plant.CalcJacobianSpatialVelocity(
            self.plant_context, JacobianWrtVariable.kQDot, I, [0, 0, 0], W, W)
        Jac_H = self.plant.CalcJacobianSpatialVelocity(
            self.plant_context, JacobianWrtVariable.kQDot, H, [0, 0, 0], W, W)
        dA = np.zeros((self.n_x**2, 3))
        dA_sq = np.zeros((self.n_x, self.n_x))
        for param_ind in range(3):
            JH = np.outer(Jac_H[param_ind, :], Jac_H[param_ind, :])
            JI = np.outer(Jac_I[param_ind, :], Jac_I[param_ind, :])
            dA_sq[8:, :8] = self.dt * np.linalg.inv(M).dot(JH + JI)
            dA[:, param_ind] = deepcopy(dA_sq.ravel())

    #print(np.sum(np.abs(dA), axis=0))
        return dA

    def reset(self):
        x_traj_new = np.zeros((self.N + 1, self.n_x))
        x_traj_new[0, :] = self.x0 + np.multiply(np.sqrt(self.W0),
                                                 np.random.randn(self.n_x))
        u_traj_new = np.zeros((self.N, self.n_u))
        #self.plant_context.SetDiscreteState(x_traj_new[0,:])
        self.plant.SetPositionsAndVelocities(self.plant_context,
                                             x_traj_new[0, :])
        return x_traj_new, u_traj_new

    def rollout(self):
        self.u_trj = np.random.randn(self.N, self.n_u) * 0.001
        self.x_trj, _ = self.reset()
        for i in range(self.N):
            self.x_trj[i + 1, :] = self.dyn(self.x_trj[i, :],
                                            self.u_trj[i],
                                            noise=True)

    def dyn(self, x, u, phi=None, noise=False):
        x_next = self.plant.AllocateDiscreteVariables()
        self.plant.SetPositionsAndVelocities(self.plant_context, x)
        self.plant.get_actuation_input_port(self.iiwa).FixValue(
            self.plant_context, u)
        self.plant.CalcDiscreteVariableUpdates(self.plant_context, x_next)
        #print(x_next.get_mutable_vector().get_value())
        x_new = x_next.get_mutable_vector().get_value()
        if noise:
            noise = np.multiply(np.sqrt(self.W), np.random.randn(self.n_x))
            x_new += noise
        return x_new

    def obs(self, x, mode=None, noise=False, phi=None):
        y = self.plant.get_state_output_port(self.iiwa).Eval(
            self.plant_context)
        #print(self.bushing.CalcBushingSpatialForceOnFrameA(self.plant_context).translational())
        if noise:
            y += np.multiply(np.sqrt(self.V), np.random.randn(self.n_y))
        return y

    def cost(self, x_trj=None, u_trj=None):
        cost_trj = 0.0
        if x_trj is None:
            for i in range(self.N):
                cost_trj += self.cost_stage(self.x_trj[i, :], self.u_trj[i, :])
            cost_trj += self.cost_final(self.sys.plant, self.x_trj[-1, :])
        else:
            for i in range(self.N):
                cost_trj += self.cost_stage(x_trj[i, :], u_trj[i, :])
            cost_trj += self.cost_final(x_trj[-1, :])
        return cost_trj
parser.add_argument("-T", "--duration",
                    type=float,
                    help="Duration to run sim.",
                    default=10.0)
args = parser.parse_args()

builder = DiagramBuilder()
plant, scene_graph = AddMultibodyPlantSceneGraph(builder)

# Load the double pendulum from Universal Robot Description Format
parser = Parser(plant, scene_graph)
parser.AddModelFromFile(FindResource("double_pendulum/double_pendulum.urdf"))
plant.Finalize()

controller = builder.AddSystem(Controller(plant, args.gravity))
builder.Connect(plant.get_state_output_port(),
                controller.get_input_port(0))
builder.Connect(controller.get_output_port(0),
                plant.get_actuation_input_port())

visualizer = builder.AddSystem(PlanarSceneGraphVisualizer(scene_graph,
                                                          xlim=[-2.8, 2.8],
                                                          ylim=[-2.8, 2.8]))
builder.Connect(scene_graph.get_pose_bundle_output_port(),
                visualizer.get_input_port(0))

diagram = builder.Build()

# Set up a simulator to run this diagram
simulator = Simulator(diagram)
simulator.set_target_realtime_rate(1.0)