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 | class DroneEnv(BaseEnv):
"""Drone Environment class
This provides an environment where the drone is contained to lie within a safe box,
and must avoid a movable obstacle.
Args:
xyz_min (ArrayLike): Minimum bounds of the safe region, shape (3,)
xyz_max (ArrayLike): Maximum bounds of the safe region, shape (3,)
"""
# Constants
RADIUS = 0.1 # TODO tune
def __init__(
self,
xyz_min: ArrayLike = (-0.5, -0.5, 0.5),
xyz_max: ArrayLike = (0.5, 0.5, 1.5),
):
# Suppress pybullet output + a Gym warning about float32 precision
with stdout_redirected(restore=False):
with warnings.catch_warnings():
warnings.simplefilter("ignore")
self._env = VelocityAviary(
# drone_model="cf2x",
num_drones=1,
initial_xyzs=np.array([[0, 0, 1]]),
initial_rpys=np.array([[0, 0, 0]]),
# physics="pyb",
neighbourhood_radius=np.inf,
pyb_freq=240,
ctrl_freq=48,
gui=True,
record=False,
obstacles=False,
user_debug_gui=False,
)
# Hack: Create same client interface as other envs
self.client: pybullet = object.__new__(BulletClient)
self.client._client = self._env.CLIENT
self.xyz_min = np.array(xyz_min)
self.xyz_max = np.array(xyz_max)
assert len(self.xyz_min) == len(self.xyz_max) == 3
self.client.setAdditionalSearchPath(find_assets_dir())
self.obstacle = self.client.loadURDF("point_robot.urdf", basePosition=(1, 1, 1))
self.client.configureDebugVisualizer(self.client.COV_ENABLE_GUI, 0)
self.client.resetDebugVisualizerCamera(1.80, 37.60, -25.00, (0.05, 0.03, 0.75))
self.robot = self._env.DRONE_IDS[0]
self.client.changeVisualShape(self.obstacle, -1, rgbaColor=[1, 0, 0, 1])
self.client.changeDynamics(self.obstacle, -1, angularDamping=10)
self.box = visualize_3D_box(
[self.xyz_min - self.RADIUS, self.xyz_max + self.RADIUS],
rgba=(0, 1, 0, 0.5),
)
# For color determination
self.is_in_box = True
self.tol = 1e-3
self.action = np.array([[0.0, 0.0, 0.0, 0.0]])
self.obs, self.reward, self.terminated, self.truncated, self.info = (
self._env.step(self.action)
)
def _update_color(self, robot_pos: ArrayLike) -> None:
"""Update the color of the box depending on if the robot is inside or not (Green inside, red outside)"""
robot_pos = np.array(robot_pos)
if np.any(robot_pos < self.xyz_min - self.tol) or np.any(
robot_pos > self.xyz_max + self.tol
):
if self.is_in_box:
self.client.changeVisualShape(self.box, -1, rgbaColor=[1, 0, 0, 0.5])
self.is_in_box = False
else:
if not self.is_in_box:
self.client.changeVisualShape(self.box, -1, rgbaColor=[0, 1, 0, 0.5])
self.is_in_box = True
def get_state(self) -> Array:
robot_pos = self.client.getBasePositionAndOrientation(self.robot)[0]
robot_vel = self.client.getBaseVelocity(self.robot)[0]
self._update_color(robot_pos)
obstacle_pos = self.client.getBasePositionAndOrientation(self.obstacle)[0]
obstacle_vel = self.client.getBaseVelocity(self.obstacle)[0]
return np.array([*robot_pos, *robot_vel]), np.array(
[*obstacle_pos, *obstacle_vel]
)
def get_desired_state(self) -> Array:
return np.array([0, 0, 1, 0, 0, 0])
def apply_control(self, u: Array) -> None:
# Note: the gym-pybullet-drones API has a weird format for the "velocity action" to take
self.action = np.array([[*u, np.linalg.norm(u)]])
def step(self):
# Step the gym-pybullet-drones environment using the last stored action
self.obs, self.reward, self.terminated, self.truncated, self.info = (
self._env.step(self.action)
)
|