-
Notifications
You must be signed in to change notification settings - Fork 55
/
user_code.py
82 lines (65 loc) · 2.38 KB
/
user_code.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
#!/usr/bin/python3
from pickle import NONE
from utils import AgileCommandMode, AgileCommand
from rl_example import rl_example
def compute_command_vision_based(state, img):
################################################
# !!! Begin of user code !!!
# TODO: populate the command message
################################################
print("Computing command vision-based!")
# print(state)
# print("Image shape: ", img.shape)
# Example of SRT command
command_mode = 0
command = AgileCommand(command_mode)
command.t = state.t
command.rotor_thrusts = [1.0, 1.0, 1.0, 1.0]
# Example of CTBR command
command_mode = 1
command = AgileCommand(command_mode)
command.t = state.t
command.collective_thrust = 15.0
command.bodyrates = [0.0, 0.0, 0.0]
# Example of LINVEL command (velocity is expressed in world frame)
command_mode = 2
command = AgileCommand(command_mode)
command.t = state.t
command.velocity = [1.0, 0.0, 0.0]
command.yawrate = 0.0
################################################
# !!! End of user code !!!
################################################
return command
def compute_command_state_based(state, obstacles, rl_policy=None):
################################################
# !!! Begin of user code !!!
# TODO: populate the command message
################################################
print("Computing command based on obstacle information!")
# print(state)
# print("Obstacles: ", obstacles)
# Example of SRT command
command_mode = 0
command = AgileCommand(command_mode)
command.t = state.t
command.rotor_thrusts = [1.0, 1.0, 1.0, 1.0]
# Example of CTBR command
command_mode = 1
command = AgileCommand(command_mode)
command.t = state.t
command.collective_thrust = 10.0
command.bodyrates = [0.0, 0.0, 0.0]
# Example of LINVEL command (velocity is expressed in world frame)
command_mode = 2
command = AgileCommand(command_mode)
command.t = state.t
command.velocity = [1.0, 0.0, 0.0]
command.yawrate = 0.0
# If you want to test your RL policy
if rl_policy is not None:
command = rl_example(state, obstacles, rl_policy)
################################################
# !!! End of user code !!!
################################################
return command