-
Notifications
You must be signed in to change notification settings - Fork 0
/
mdn_main.py
101 lines (85 loc) · 4.35 KB
/
mdn_main.py
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
import argparse
import numpy as np
import gc
from tqdm import tqdm
import Modules.MonteCarloTools as MCTools
from Modules.run_monte_carlo_sweep import MonteCarlo
from Modules.Structures import CoordinateSystemType, FrameType
from Modules.Sensors import TVCAction
from Modules.main import Simulation
def has_reached_burnout(sim: Simulation):
return (np.abs(sim.propulsion.getForce(sim.state)[0][0]) == 0) and (sim.state.time > sim.state.dt)
def run_mdn_sim(mc: MonteCarlo, inputs: tuple) -> None:
name, config, i = inputs
sim = mc._initialise_simulation(name, config, i)
# Initial states
sBI__G = sim.state.getPosition(CoordinateSystemType.GEOCENTRIC)
T_gG = sim.state.getTransformationMatrix(CoordinateSystemType.GEOGRAPHIC, CoordinateSystemType.GEOCENTRIC)
sBI__G_initial = sim.state.getPosition(CoordinateSystemType.GEOCENTRIC)
initial_alt = sim.state.getAltitudeGeometric()
step_count = 0
# We terminate the sim at burnout since TVC loses control authority beyond this point (and hence RL receives poor gradients)
while sim.should_continue_loop(sim.state) and not has_reached_burnout(sim):
# Only giving the model access every few timesteps (based on the hardware's clockspeed) and prior to burnout
if step_count % int(1/(sim.tvc_clockspeed*sim.state.dt)) == 0:
# Returns current state info
#TODO IN THE EARTH FRAME - IRL MAY BE IN THE ROCKET'S FRAME (e.g. acceleration measured on board)
altitude = sim.state.getAltitudeGeometric() - initial_alt # Zeroed for RL
sBI__G = sim.state.getPosition(CoordinateSystemType.GEOCENTRIC)
displacement = T_gG @ (sBI__G - sBI__G_initial)
vB_E_B = sim.state.getVelocity(coord=CoordinateSystemType.BODY, frame=FrameType.EARTH)
aB_E_B = sim.state.getAcceleration(coord=CoordinateSystemType.BODY, frame=FrameType.EARTH)
euler_angles = np.rad2deg(sim.state.euler_angles)
wB_E_B = sim.state.getAngularVelocity(coord=CoordinateSystemType.BODY, frame=FrameType.EARTH)
# Agent chooses action
action = None #TVCAction.TILT_UP # @MDN THIS IS WHERE YOU PASS IN YOUR ACTION
if action == TVCAction.TILT_UP:
sim.tilt += sim.tilt_stepsize
elif action == TVCAction.TILT_DOWN:
sim.tilt -= sim.tilt_stepsize
elif action == TVCAction.PAN_UP:
sim.pan += sim.pan_stepsize
elif action == TVCAction.PAN_DOWN:
sim.pan -= sim.pan_stepsize
else:
pass
#raise ValueError(f'Action {action} not recognised. Please use TU, TD, PU, PD')
# Enforcing a hard limit on tilt and pan
if sim.tilt > sim.max_tilt:
sim.tilt -= sim.tilt_stepsize
elif sim.tilt < -sim.max_tilt:
sim.tilt += sim.tilt_stepsize
if sim.pan > sim.max_pan:
sim.pan -= sim.pan_stepsize
elif sim.pan < -sim.max_pan:
sim.pan += sim.pan_stepsize
# Updating the state given the action
sim.iteration()
step_count += 1
# Keeping memory usage in check
del sim.aerodynamics.force_coefficients
del sim.aerodynamics.moment_coefficients
try:
del sim.atmosphere.pyatm
except AttributeError:
pass
gc.collect() # Runs the garbage collection.
state_history, _ = sim.finalise_run()
mc.results[sim] = state_history.df
if __name__=='__main__':
parser = argparse.ArgumentParser()
parser.add_argument('--vehicle_config', type=str, help='.yaml file which specifies the initial parameters', default='mdn_vehicleConfig.yaml')
parser.add_argument('--mc_config', type=str, help='.yaml file which specifies the monte carlo parameters', default='mdn_mcConfig.yaml')
parser.add_argument('--save_loc', type=str, help='Path to Monte Carlo save location. Not saved by default', default=None)
args = parser.parse_args()
mc = MonteCarlo(args.mc_config, args.vehicle_config)
cases = mc._generate_cases()
for inputs in tqdm(cases):
run_mdn_sim(mc, inputs)
# Saving sim data for postprocessing
if args.save_loc is not None:
res = MCTools.MonteCarloResults(mc.results)
res.save(f'{args.save_loc}.pkl')
import utils
utils.print_all_open_files()
print(utils.get_open_file_count())