SWIXMujocoSim/main.py at main · ncsurobotics/SWIXMujocoSim · 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
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
import mujoco
import mujoco.viewer
import time
import numpy as np
import os
from sub import Sub, SubParams
from telemetry import Telemetry
import matplotlib.pyplot as plt
import time
from mppi import *
from trajectory import Trajectory
from dynamics import *

BASE_DIR = os.path.dirname(os.path.abspath(__file__))
xml_path = os.path.join(BASE_DIR, "project.xml")


model = mujoco.MjModel.from_xml_path(xml_path)
data = mujoco.MjData(model)
mujoco.mj_resetData(model, data)


# Initialize submarine
m = 22.0
Ix = 1/12 * m * (0.4**2 + 0.4**2)
Iy = 1/12 * m * (0.4**2 + 0.8**2)
Iz = 1/12 * m * (0.8**2 + 0.4**2)

grav = 9.81
vol = 0.178
rho = 1000

W = m*grav
B = rho * grav * vol

eta = np.array([0, 0, 5, 0, 0, 0])
v = np.zeros(6)
Mrb = np.diag([m, m, m, Ix, Iy, Iz])

def iface(thrusts):
    for i, thrust in enumerate(thrusts):
        data.ctrl[i] = thrust

telemetry = Telemetry(data, model)
params = SubParams(Mrb, Ma)

sub = Sub(eta, v, iface, telemetry, params)

waypoints = np.array([
    [0.0,0.0,5.0,0.0,0.0,0.0, 0.0],
    [0.0,0.0,8.0,0.0,0.0,0.0, 4.0],
    [0.0,0.0,8.0,0.0,0.0,0.0, 8.0],
    [0.0,0.0,5.0,0.0,0.0,0.0, 12.0]
])

traj = Trajectory(waypoints)

def warmup():
    print("Warming up...")
    mppi(np.concat([eta, np.zeros(6)]), traj, 0.15, 1000, 12, 1, 0.1, m, Ix, Iy, Iz, W, B, params)

warmup()


predicted_state = sub.eta
t = []
goal_states = []
true_states = []
plotting_started = True
sim_time = 0

ctrl = np.zeros(6)
count = 0
start = time.perf_counter()
with mujoco.viewer.launch_passive(model, data) as viewer:
    dt = model.opt.timestep
    while viewer.is_running():

        now = time.perf_counter()
        print(now-start)
        if count % 100 == 0:
            ctrl = mppi(sub.telemetry.x(), traj, sim_time, 2500, 20, 1, 0.01, m, Ix, Iy, Iz, W, B, params)
        # Apply controller input
        sub.control(ctrl)
        mujoco.mj_step(model, data)

        goal_state = traj.sample_trajectory(sim_time)
        # Save predicted and true measurements
        goal_states.append(goal_state)
        true_state = np.concatenate([sub.telemetry.pos(), sub.telemetry.rot()]).tolist()
        true_states.append(true_state)
        t.append(sim_time)

        viewer.sync()

        if now - start > 60:
            break
        count +=1
        sim_time += dt
        time.sleep(dt)


predicted_states = np.array(goal_states)
true_states = np.array(true_states)
print(sim_time)

plt.subplot(231)
plt.plot(t, predicted_states[:, 0], label="Predicted")
plt.plot(t, true_states[:, 0], label="Measured")
plt.legend()

plt.subplot(232)
plt.plot(t, predicted_states[:, 1])
plt.plot(t, true_states[:, 1])

plt.subplot(233)
plt.plot(t, predicted_states[:, 2])
plt.plot(t, true_states[:, 2])

plt.subplot(234)
plt.plot(t, predicted_states[:, 3])
plt.plot(t, true_states[:, 5])

plt.subplot(235)
plt.plot(t, predicted_states[:, 4])
plt.plot(t, true_states[:, 4])

plt.subplot(236)
plt.plot(t, predicted_states[:, 5])
plt.plot(t, true_states[:, 3])

plt.show()