drake-tutorials/Lesson-5/Lesson-5-Task-4-solution.py at main · logdog/drake-tutorials · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
import numpy as np

from pydrake.all import (
    plot_system_graphviz,
    DiagramBuilder,
    LeafSystem,
    LogVectorOutput,
    Meshcat,
    MeshcatParams,
    RigidTransform,
    ConstantVectorSource,
    Parser,
    SpatialForce,
    AddMultibodyPlantSceneGraph,
    Propeller,
    ExternallyAppliedSpatialForceMultiplexer,
    Simulator,
    DiagramBuilder,
    AddDefaultVisualization,
    PdControllerGains,
    AbstractValue,
    JointIndex,
    FrameIndex,
    BodyIndex,
    RpyFloatingJoint,
    GeometryInstance,
    RotationMatrix,
    Cylinder,
    IllustrationProperties,
    Rgba,
    MakePhongIllustrationProperties,
    Sphere,
)

import os
import matplotlib.pyplot as plt

if __name__ == "__main__":
    urdf_path = os.path.join(os.path.dirname(__file__), "TP.urdf")

    meshcat = Meshcat()
    meshcat.SetCameraPose(
        camera_in_world=np.array([0, -2, 0.5]),
        target_in_world=np.array([0, 0, 0]),
    )

    builder = DiagramBuilder()
    plant, scene_graph = AddMultibodyPlantSceneGraph(builder, time_step=0.0)
    parser = Parser(plant)
    parser.AddModels(urdf_path)
    plant.Finalize()

    AddDefaultVisualization(builder, meshcat)
    diagram = builder.Build()

    # manually add axes to the upper_arm/lower_arm frames
    # this is useful for debugging joints in your URDF file
    inspector = scene_graph.model_inspector()
    plant_source_id = plant.get_source_id()
    for frame_id in inspector.FramesForSource(plant_source_id):
        frame_name = inspector.GetName(frame_id)

        radius = 0.002
        length = 0.25
        x_axis = scene_graph.RegisterGeometry(plant_source_id, frame_id,
            GeometryInstance(RigidTransform(RotationMatrix.MakeYRotation(np.pi/2), np.array([length/2,0,0])), Cylinder(radius,length), frame_name+"_xaxis"))
        scene_graph.AssignRole(plant_source_id, x_axis, MakePhongIllustrationProperties([1, 0, 0, 1.0]))

        y_axis = scene_graph.RegisterGeometry(plant_source_id, frame_id,
            GeometryInstance(RigidTransform(RotationMatrix.MakeXRotation(np.pi/2), np.array([0,length/2,0])), Cylinder(radius,length), frame_name+"_yaxis"))
        scene_graph.AssignRole(plant_source_id, y_axis, MakePhongIllustrationProperties([0, 1, 0, 1.0]))

        z_axis = scene_graph.RegisterGeometry(plant_source_id, frame_id,
            GeometryInstance(RigidTransform(RotationMatrix(), np.array([0,0,length/2])), Cylinder(radius,length), frame_name+"_zaxis"))
        scene_graph.AssignRole(plant_source_id, z_axis, MakePhongIllustrationProperties([0, 0, 1, 1.0]))


    # Set initial conditions
    context = diagram.CreateDefaultContext()
    context.SetContinuousState(np.deg2rad([30, 30, 30, 0, 0, 0]))

    # simulate the multibody plant
    simulator = Simulator(diagram, context)
    simulator.Initialize()
    simulator.set_target_realtime_rate(1.) # optional
    meshcat.StartRecording()
    simulator.AdvanceTo(15.0)
    meshcat.PublishRecording()

    while True:
        pass