swangnice/rotor_py_control
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Allow user to change quad color.

Browse Source
This commit is contained in:
spencerfolk
2024-01-04 12:18:53 -05:00
parent 6856a4f146
commit 1302741252
1 changed files with 6 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
rotorpy/learning/quadrotor_environments.py
View File

@@ -38,6 +38,7 @@ class QuadrotorEnv(gym.Env):
render_mode: render the quadrotor. render_mode: render the quadrotor.
render_fps: rendering frames per second, lower this for faster visualization. render_fps: rendering frames per second, lower this for faster visualization.
ax: for plotting purposes, you can supply an axis object that the quadrotor will visualize on. ax: for plotting purposes, you can supply an axis object that the quadrotor will visualize on.
color: choose the color of the quadrotor. If none, it will randomly select a color.
""" """
metadata = {"render_modes": ["None", "3D", "console"], metadata = {"render_modes": ["None", "3D", "console"],
@@ -63,6 +64,7 @@ class QuadrotorEnv(gym.Env):
render_fps = 30, # The rendering frames per second. Lower this for faster visualization. render_fps = 30, # The rendering frames per second. Lower this for faster visualization.
fig = None, # Figure for rendering. Optional. fig = None, # Figure for rendering. Optional.
ax = None, # Axis for rendering. Optional. ax = None, # Axis for rendering. Optional.
color = None, # The color of the quadrotor.
): ):
super(QuadrotorEnv, self).__init__() super(QuadrotorEnv, self).__init__()
@@ -159,7 +161,10 @@ class QuadrotorEnv(gym.Env):
else: else:
self.fig = fig self.fig = fig
self.ax = ax self.ax = ax
if color is None:
colors = list(mcolors.CSS4_COLORS) colors = list(mcolors.CSS4_COLORS)
else:
colors = [color]
self.quad_obj = Quadrotor(self.ax, wind=True, color=np.random.choice(colors)) self.quad_obj = Quadrotor(self.ax, wind=True, color=np.random.choice(colors))
self.world_artists = None self.world_artists = None
self.title_artist = self.ax.set_title('t = {}'.format(self.t)) self.title_artist = self.ax.set_title('t = {}'.format(self.t))