vincekurtz / drake_ddp

Differential Dynamic Programming and Iterative LQR in Drake
MIT License
64 stars 12 forks source link

Difference equation assertion error #6

Closed kamiradi closed 2 days ago

kamiradi commented 4 days ago

Hi I am trying to use your iLQR implementation for my own system. The discrete system assertion is being called whenever I try and run the code on my system. To the best of my knowledge, this is if you have given a time_step of 0.0 to your multibody plant when intialising it. However, I give it a discrete time step of 1e-2. Any ideas of where I could be going wrong?

The following is the function I use instead of your create_system_model function

def MakeInsertionAssemblyStation(plant_config, config, cfg):

    # sim related settings

    # create a diagram
    directives = "./assembly.dmd.yaml"
    builder = DiagramBuilder()

    # plant, scene_graph = AddMultibodyPlant(
    #     config=plant_config,
    #     builder=builder
    # )
    plant, scene_graph = AddMultibodyPlant(
        config=plant_config,
        builder=builder
    )
    parser = Parser(plant)

    # add the robot
    iiwa_dmd = './assembly.dmd.yaml'

    model_indices = parser.AddModels(iiwa_dmd)

    # add some transforms to denote the hole
    hole_index = plant.GetModelInstanceByName("hole")
    hole_frame = plant.AddFrame(
        FixedOffsetFrame(
            "hole_frame",
            plant.GetFrameByName(
                "base_link",
                hole_index),
            RigidTransform(RotationMatrix.MakeXRotation(-np.pi/2))
        )
    )

    # adding peg to the scene
    G_PPcm_P = UnitInertia.SolidCylinder(
        config["peg_width"]*0.5,
        config["peg_length"],
        [0, 0, 1]
    )
    M_PPcm_P = SpatialInertia(
        config["peg_mass"],
        [0, 0, 0],
        G_PPcm_P
    )
    peg_geom = Cylinder(
        config["peg_width"]*0.5,
        config["peg_length"]
    )

    peg = plant.AddRigidBody(
        'peg',
        plant.GetModelInstanceByName(
            'iiwa'),
        M_PPcm_P
    )
    plant.RegisterVisualGeometry(
        peg,
        RigidTransform(),
        peg_geom,
        'peg',
        [0.5, 0.5, 0.5, 1.0]
    )
    plant.RegisterCollisionGeometry(
        peg,
        RigidTransform(),
        peg_geom,
        'peg',
        config['peg_props']
    )
    plant.WeldFrames(
        plant.GetFrameByName("iiwa_link_7"),
        peg.body_frame(),
        RigidTransform([0, 0, 0.06])
    )
    stiffness_frame = plant.AddFrame(
        FixedOffsetFrame(
            "stiffness_frame",
            peg.body_frame(),
            RigidTransform(
                RotationMatrix.MakeXRotation(-np.pi),
                [0, 0, config["peg_length"]/2]))
    )
    plant.RegisterVisualGeometry(
        peg,
        RigidTransform(
            [0, 0, config['peg_length']/2]),
        Sphere(0.002),
        "peg_tip",
        [1.0, 0.0, 0.0, 1.0]
    )

    # load up uncertain hole just for visualisation
    # load up the noisy hole to plan to
    hole_uncertain = "/root/workspace/model/peg_uncertain.urdf"
    hole_plan = "/root/workspace/model/peg_plan.urdf"
    sigma_parser = Parser(plant, "sigma")
    plan_parser = Parser(plant, "plan_hole")
    sigma_parser.SetAutoRenaming(True)
    plan_parser.SetAutoRenaming(True)
    for i in range(10):
        x = np.random.normal(0.5, 0.005)
        y = np.random.normal(0.0, 0.005)
        z = np.random.normal(0.08, 0.0025)
        rot = np.random.normal(-np.pi/2, 0.005)
        hole_index = sigma_parser.AddModels(hole_uncertain)[0]
        plant.WeldFrames(
            plant.world_frame(),
            plant.GetFrameByName("base_link", hole_index),
            RigidTransform(
            RotationMatrix.MakeXRotation(np.pi/2),
            [x, y, z])
        )
        hole_frame = plant.AddFrame(
            FixedOffsetFrame(
                "hole_frame_{}".format(i),
                plant.GetFrameByName(
                    "base_link",
                    hole_index),
                RigidTransform(RotationMatrix.MakeXRotation(-np.pi/2))
            )
        )

    x = np.random.normal(0.5, 0.0025)
    y = np.random.normal(0.0, 0.0025)
    # z = np.random.normal(0.08, 0.0025)
    z = 0.08
    x_rot = np.random.normal(0.0, 2.5)*0
    y_rot = np.random.normal(0.0, 2.5)*0
    z_rot = np.random.normal(0.0, 2.5)*0
    drake_rotation = RotationMatrix.MakeXRotation(np.pi/2 + np.deg2rad(x_rot)).MakeYRotation(np.deg2rad(y_rot)).MakeZRotation(np.deg2rad(z_rot))
    plan_hole_index = plan_parser.AddModels(hole_plan)[0]
    plant.WeldFrames(
        plant.world_frame(),
        plant.GetFrameByName("base_link", plan_hole_index),
        RigidTransform(
        RotationMatrix.MakeXRotation(np.pi/2 + np.deg2rad(x_rot)),
        [x, y, z])
    )
    hole_frame = plant.AddFrame(
        FixedOffsetFrame(
            "hole_frame",
            plant.GetFrameByName(
                "base_link",
                plan_hole_index),
            RigidTransform(RotationMatrix.MakeXRotation(-np.pi/2))
        )
    )
    plant.Finalize()
    iiwa_instance = plant.GetModelInstanceByName('iiwa')

    # add controller
    controller_plant = MultibodyPlant(time_step=plant_config.time_step)
    control_parser = Parser(controller_plant)

    # add the robot
    iiwa_controller_dmd = './iiwa.dmd.yaml'

    iiwa_control_indices = control_parser.AddModels(iiwa_controller_dmd)

    # add the peg
    peg_control = controller_plant.AddRigidBody(
        'peg_control',
        controller_plant.GetModelInstanceByName('iiwa_controller'),
        M_PPcm_P
    )
    controller_plant.WeldFrames(
        controller_plant.GetFrameByName("iiwa_link_7"),
        peg_control.body_frame(),
        RigidTransform([0, 0, 0.06])
    )
    controller_stiffness_frame = controller_plant.AddFrame(
        FixedOffsetFrame(
            "stiffness_frame",
            peg.body_frame(),
            RigidTransform(
                RotationMatrix.MakeXRotation(-np.pi),
                [0, 0, config['peg_length']/2])
        )
    )
    controller_plant.Finalize()
    iiwa_controller_instance = controller_plant.GetModelInstanceByName('iiwa_controller')
    num_iiwa_positions = controller_plant.num_positions(iiwa_controller_instance)

    # add the controller
    iiwa_controller = builder.AddSystem(
        # InverseDynamicsController(
        #     controller_plant,
        #     kp=[100] * num_iiwa_positions,
        #     kd=[20] * num_iiwa_positions,
        #     ki=[1] * num_iiwa_positions,
        #     has_reference_acceleration=False
        # )
        JointStiffnessController(
            controller_plant,
            # kp=[100] * num_iiwa_positions,
            # kd=[20] * num_iiwa_positions
            kp=[cfg.plant.Kp] * num_iiwa_positions,
            kd=[cfg.plant.Kd] * num_iiwa_positions
        )
    )
    iiwa_controller.set_name("inverse dynamics controller")

    demux = builder.AddNamedSystem(
        'StateDemuxer',
        Demultiplexer(2 * num_iiwa_positions, num_iiwa_positions)
        )

    builder.Connect(
            plant.get_state_output_port(iiwa_instance),
            demux.get_input_port())

    # NOTE: Connect controller to adder
    adder = builder.AddNamedSystem(
            "ControlAdder",
            Adder(2, num_iiwa_positions)
            )
    # NOTE: If running InverseDynamicsController
    # builder.Connect(
    #         iiwa_controller.get_output_port_control(),
    #         adder.get_input_port(0)
    #         )
    # NOTE: If running JointStiffnessController
    builder.Connect(
            iiwa_controller.get_output_port_generalized_force(),
            adder.get_input_port(0)
            )

    # NOTE: Connect the passed torque to adder
    torque_passthrough = builder.AddNamedSystem(
            "FeedForwardTorquePass",
            PassThrough([0] * num_iiwa_positions)
            )
    # TODO: Need to add a state interpolator, desired position passthrough
    position_pass = builder.AddNamedSystem(
            'DesiredPostionPass',
            PassThrough(num_iiwa_positions)
            )
    desired_state_from_position = builder.AddNamedSystem(
            'StateInterpolator',
            StateInterpolatorWithDiscreteDerivative(
                num_iiwa_positions,
                plant_config.time_step,
                suppress_initial_transient=True)
            )
    builder.Connect(
            torque_passthrough.get_output_port(),
            adder.get_input_port(1)
            )
    builder.ExportInput(
            torque_passthrough.get_input_port(),
            "iiwa_feedforward_torque")
    # NOTE: Connect the adder to the plant's actuation
    builder.Connect(
            adder.get_output_port(),
            plant.get_actuation_input_port(iiwa_instance)
            )
    # NOTE: Connect plant to controller
    builder.Connect(
            plant.get_state_output_port(iiwa_instance),
            iiwa_controller.get_input_port_estimated_state())
    builder.ExportInput(
            position_pass.get_input_port(),
            "iiwa_position_desired"
            )
    builder.Connect(
            desired_state_from_position.get_output_port(),
            iiwa_controller.get_input_port_desired_state()
            )
    builder.Connect(
            position_pass.get_output_port(),
            desired_state_from_position.get_input_port()
            )

    # export ports
    builder.ExportOutput(
        scene_graph.get_query_output_port(),
        "query_object"
    )
    builder.ExportOutput(
        plant.get_body_poses_output_port(),
        "body_poses"
    )
    builder.ExportOutput(
        plant.get_contact_results_output_port(),
        "contact_results"
    )
    builder.ExportOutput(
            demux.get_output_port(0),
            'iiwa_position_measured'
            )
    builder.ExportOutput(
            demux.get_output_port(1),
            'iiwa_velocity_measured'
            )
    builder.ExportOutput(
            plant.get_state_output_port(iiwa_instance),
            'robot_state'
            )

    # build the diagram
    diagram = builder.Build()
    diagram.set_name("Assembly Station")

    return model_indices, diagram

I have left comments on the code for it to be more readable. Let me know if I can create a smaller example for this to be grokable. I then pass the returned diagram to the iLQR class. Do let me know if you want further clarification!

vincekurtz commented 4 days ago

Hi @kamiradi, thanks for your interest in our work! I'm not able to tell what the issue is from this code, and I'm also not sure what "the discrete system assertion" is referring to. A smaller and self contained example would be helpful, preferably with a simpler system as well.

kamiradi commented 3 days ago

Hi @vincekurtz, thanks for your response.

In your ilqr.py file, this line makes sure that the diagram you have passed into this system is a discrete-time system. The method that I wrote in the main description of the issue replaces your create_system_model function with my own model. I now see that it is a bit too large so I have created smaller function that loads my system, without changing how I am instantiating the MultibodyPlant. This is as follows

def CreateSystemModel(plant, scene_graph, config):

    # create system model for iLQR implementation
    iiwa_dmd = './assembly.dmd.yaml'
    parser = Parser(plant)

    # add iiwa robot, tables, and hole
    model_indices = parser.AddModels(iiwa_dmd)

    # add some transforms to denote the hole
    hole_index = plant.GetModelInstanceByName("hole")
    hole_frame = plant.AddFrame(
        FixedOffsetFrame(
            "hole_frame",
            plant.GetFrameByName(
                "base_link",
                hole_index),
            RigidTransform(RotationMatrix.MakeXRotation(-np.pi/2))
        )
    )

    # adding peg to the scene
    G_PPcm_P = UnitInertia.SolidCylinder(
        config["peg_width"]*0.5,
        config["peg_length"],
        [0, 0, 1]
    )
    M_PPcm_P = SpatialInertia(
        config["peg_mass"],
        [0, 0, 0],
        G_PPcm_P
    )
    peg_geom = Cylinder(
        config["peg_width"]*0.5,
        config["peg_length"]
    )

    peg = plant.AddRigidBody(
        'peg',
        plant.GetModelInstanceByName(
            'iiwa'),
        M_PPcm_P
    )
    plant.RegisterVisualGeometry(
        peg,
        RigidTransform(),
        peg_geom,
        'peg',
        [0.5, 0.5, 0.5, 1.0]
    )
    plant.RegisterCollisionGeometry(
        peg,
        RigidTransform(),
        peg_geom,
        'peg',
        config['peg_props']
    )
    plant.WeldFrames(
        plant.GetFrameByName("iiwa_link_7"),
        peg.body_frame(),
        RigidTransform([0, 0, 0.06])
    )
    stiffness_frame = plant.AddFrame(
        FixedOffsetFrame(
            "stiffness_frame",
            peg.body_frame(),
            RigidTransform(
                RotationMatrix.MakeXRotation(-np.pi),
                [0, 0, config["peg_length"]/2]))
    )
    plant.RegisterVisualGeometry(
        peg,
        RigidTransform(
            [0, 0, config['peg_length']/2]),
        Sphere(0.002),
        "peg_tip",
        [1.0, 0.0, 0.0, 1.0]
    )

    plant.Finalize()
    logger.info(f" time step: {plant.time_step()}")
    assert plant.IsDifferenceEquationSystem() "Has to be discrete time"
    return plant, scene_graph

plant_config = MultibodyPlantConfig() 
plant_config.time_step = 1e-2 
plant_config.contact_model = "hydroelastic_with_fallback" 
plant_config.contact_surface_representation = "polygon" 
plant_config.penetration_allowance = 1e-5
plant, scene_graph = AddMultibodyPlant(
        config=plant_config,
        builder=builder
    )
plant, scene_graph = CreateSystemModel(plant, scene_graph, config)

I also put the same assert just before returning plant and scene_graph, and it still fires false. I posted this issue in-case it was a drake bug. Let me know if you need further clarification.

kamiradi commented 2 days ago

Hi, this has now been solved. Turns out it was an issue with how I was instantiating the MultibodyPlant. More here