drake-tutorials/Lesson-6/example3.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
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
import numpy as np
import matplotlib.pyplot as plt

from pydrake.all import (
    plot_system_graphviz,
    DiagramBuilder,
    Simulator,
    LeafSystem,
    SimulatorConfig,
    ApplySimulatorConfig,
    ZeroOrderHold,
    LogVectorOutput,
)

dt = 1/50

class ExampleContinuousPlant(LeafSystem):

    def __init__(self):
        LeafSystem.__init__(self)
        self.DeclareVectorInputPort("u", 1)
        self.state_index = self.DeclareContinuousState(1,1,0)
        self.DeclareVectorOutputPort("y", 2, self.Output)

    def DoCalcTimeDerivatives(self, context, derivatives):
        t = context.get_time()
        u = self.get_input_port().Eval(context)[0]
        theta = context.get_continuous_state_vector().GetAtIndex(0)
        theta_dot = context.get_continuous_state_vector().GetAtIndex(1)

        # pendulum dynamics
        theta_ddot = -9.81 * np.sin(theta) + u
        derivatives.SetFromVector([theta_dot, theta_ddot])
        print(f"Plant.DoCalcTimeDerivatives(): \tt={t:0.6f},                                     \tu={u:0.6f}")


    def Output(self, context, output):
        t = context.get_time()
        theta = context.get_continuous_state_vector().GetAtIndex(0)
        theta_dot = context.get_continuous_state_vector().GetAtIndex(1)

        output.SetFromVector([theta, theta_dot])
        print(f"Plant.Output(): \t\tt={t:0.6f}, theta={theta:0.6f}, theta_dot={theta_dot:0.6f}")

class ExampleDiscreteController(LeafSystem):

    def __init__(self):
        LeafSystem.__init__(self)
        self.DeclareVectorInputPort("y_sampled", 2)
        self.DeclareVectorOutputPort("u", 1, self.Output)

    def Output(self, context, output):
        t = context.get_time()
        theta = self.get_input_port().Eval(context)[0]
        theta_dot = self.get_input_port().Eval(context)[1]

        # simple proportional controller
        u = -2.0 * theta - 0.5 * theta_dot
        output.SetFromVector([u])
        print(f"Controller.Output(): \t\tt={t:0.6f}, theta={theta:0.6f}, theta_dot={theta_dot:0.6f}  \tu={u:0.6f}")

class MyZeroOrderHold(LeafSystem):

    def __init__(self, dt, size):
        LeafSystem.__init__(self)
        self.DeclareVectorInputPort("x", size)
        self.state_index = self.DeclareDiscreteState(size)
        self.DeclarePeriodicDiscreteUpdateEvent(dt, 0.0, self.LatchInputVectorToState)
        self.DeclareStateOutputPort("y", self.state_index)

    def LatchInputVectorToState(self, context, discrete_state):
        t = context.get_time()
        x = self.get_input_port().Eval(context)
        discrete_state.set_value(x)
        print(f"ZOH.LatchInputVectorToState(): \tt={t:0.6f}, theta={x[0]:0.6f}, theta_dot={x[1]:0.6f}")


if __name__ == "__main__":

    # create the diagram
    builder = DiagramBuilder()
    plant = builder.AddSystem(ExampleContinuousPlant())
    zoh = builder.AddSystem(MyZeroOrderHold(dt, 2))
    controller = builder.AddSystem(ExampleDiscreteController())
    builder.Connect(plant.get_output_port(), zoh.get_input_port())
    builder.Connect(zoh.get_output_port(), controller.get_input_port())
    builder.Connect(controller.get_output_port(), plant.get_input_port())

    diagram = builder.Build()

    # plot_system_graphviz(diagram)
    # plt.show()

    # set initial conditions
    # Initial state: 30 degrees, 0 rad/s
    context = diagram.CreateDefaultContext()
    plant_context = diagram.GetMutableSubsystemContext(plant, context)
    plant_context.SetContinuousState(np.array([0.1, 0]))


    # create the simulator
    simulator = Simulator(diagram, context)
    simulator_config = SimulatorConfig(max_step_size=dt/2)
    ApplySimulatorConfig(simulator_config, simulator)

    # run the simulation
    simulator.Initialize()
    print('Initialize complete!\nStarting simulation...')
    simulator.AdvanceTo(3*dt)  # Run for 3 time steps
    print('Simulation complete.')