From 130274125213b31786cd31c8ab052dd8249e2771 Mon Sep 17 00:00:00 2001 From: spencerfolk Date: Thu, 4 Jan 2024 12:18:53 -0500 Subject: [PATCH] Allow user to change quad color. --- rotorpy/learning/quadrotor_environments.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/rotorpy/learning/quadrotor_environments.py b/rotorpy/learning/quadrotor_environments.py index 7c76b14..c8c4d28 100644 --- a/rotorpy/learning/quadrotor_environments.py +++ b/rotorpy/learning/quadrotor_environments.py @@ -38,6 +38,7 @@ class QuadrotorEnv(gym.Env): render_mode: render the quadrotor. 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. + color: choose the color of the quadrotor. If none, it will randomly select a color. """ 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. fig = None, # Figure for rendering. Optional. ax = None, # Axis for rendering. Optional. + color = None, # The color of the quadrotor. ): super(QuadrotorEnv, self).__init__() @@ -159,7 +161,10 @@ class QuadrotorEnv(gym.Env): else: self.fig = fig self.ax = ax - colors = list(mcolors.CSS4_COLORS) + if color is None: + colors = list(mcolors.CSS4_COLORS) + else: + colors = [color] self.quad_obj = Quadrotor(self.ax, wind=True, color=np.random.choice(colors)) self.world_artists = None self.title_artist = self.ax.set_title('t = {}'.format(self.t))