├── LICENSE ├── README.md ├── gym_rocketlander ├── __init__.py └── envs │ ├── __init__.py │ └── rocket_lander_env.py └── setup.py /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2021 Sven Niederberger 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # ![GIF](https://thumbs.gfycat.com/CoarseEmbellishedIsopod-max-14mb.gif) 2 | [Click here for higher quality video](https://gfycat.com/CoarseEmbellishedIsopod) 3 | 4 | The objective of this environment is to land a rocket on a ship. The environment is highly customizable and takes discrete or continuous inputs. 5 | 6 | ### Installation 7 | ``` 8 | cd gym-rocketlander 9 | pip install -e . 10 | ``` 11 | 12 | ### Usage 13 | Once the has been installed, it can be used like any other Gym environment: 14 | ```py 15 | env = gym.make("gym_rocketlander:rocketlander-v0") 16 | ``` 17 | 18 | ### STATE VARIABLES 19 | The state consists of the following variables: 20 | * x position 21 | * y position 22 | * angle 23 | * first leg ground contact indicator 24 | * second leg ground contact indicator 25 | * throttle 26 | * engine gimbal 27 | 28 | If VEL_STATE is set to true, the velocities are included: 29 | * x velocity 30 | * y velocity 31 | * angular velocity 32 | 33 | ### CONTROL INPUTS 34 | Discrete control inputs are: 35 | * gimbal left 36 | * gimbal right 37 | * throttle up 38 | * throttle down 39 | * use first control thruster 40 | * use second control thruster 41 | * no action 42 | 43 | Continuous control inputs are: 44 | * gimbal (left/right) 45 | * throttle (up/down) 46 | * control thruster (left/right) 47 | -------------------------------------------------------------------------------- /gym_rocketlander/__init__.py: -------------------------------------------------------------------------------- 1 | from gym.envs.registration import register 2 | 3 | register( 4 | id="rocketlander-v0", 5 | entry_point="gym_rocketlander.envs:RocketLander", 6 | ) -------------------------------------------------------------------------------- /gym_rocketlander/envs/__init__.py: -------------------------------------------------------------------------------- 1 | from gym_rocketlander.envs.rocket_lander_env import RocketLander 2 | -------------------------------------------------------------------------------- /gym_rocketlander/envs/rocket_lander_env.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import Box2D 3 | from Box2D.b2 import ( 4 | fixtureDef, 5 | polygonShape, 6 | revoluteJointDef, 7 | distanceJointDef, 8 | contactListener, 9 | ) 10 | import gym 11 | from gym import spaces 12 | from gym.utils import seeding 13 | from gym.envs.classic_control import rendering 14 | 15 | """ 16 | The objective of this environment is to land a rocket on a ship. 17 | 18 | STATE VARIABLES 19 | The state consists of the following variables: 20 | - x position 21 | - y position 22 | - angle 23 | - first leg ground contact indicator 24 | - second leg ground contact indicator 25 | - throttle 26 | - engine gimbal 27 | If VEL_STATE is set to true, the velocities are included: 28 | - x velocity 29 | - y velocity 30 | - angular velocity 31 | all state variables are roughly in the range [-1, 1] 32 | 33 | CONTROL INPUTS 34 | Discrete control inputs are: 35 | - gimbal left 36 | - gimbal right 37 | - throttle up 38 | - throttle down 39 | - use first control thruster 40 | - use second control thruster 41 | - no action 42 | 43 | Continuous control inputs are: 44 | - gimbal (left/right) 45 | - throttle (up/down) 46 | - control thruster (left/right) 47 | """ 48 | 49 | CONTINUOUS = True 50 | VEL_STATE = True # Add velocity info to state 51 | FPS = 60 52 | SCALE_S = 0.35 # Temporal Scaling, lower is faster - adjust forces appropriately 53 | INITIAL_RANDOM = 0.4 # Random scaling of initial velocity, higher is more difficult 54 | 55 | START_HEIGHT = 800.0 56 | START_SPEED = 80.0 57 | 58 | # ROCKET 59 | MIN_THROTTLE = 0.4 60 | GIMBAL_THRESHOLD = 0.4 61 | MAIN_ENGINE_POWER = 1600 * SCALE_S 62 | SIDE_ENGINE_POWER = 100 / FPS * SCALE_S 63 | 64 | ROCKET_WIDTH = 3.66 * SCALE_S 65 | ROCKET_HEIGHT = ROCKET_WIDTH / 3.7 * 47.9 66 | ENGINE_HEIGHT = ROCKET_WIDTH * 0.5 67 | ENGINE_WIDTH = ENGINE_HEIGHT * 0.7 68 | THRUSTER_HEIGHT = ROCKET_HEIGHT * 0.86 69 | 70 | # LEGS 71 | LEG_LENGTH = ROCKET_WIDTH * 2.2 72 | BASE_ANGLE = -0.27 73 | SPRING_ANGLE = 0.27 74 | LEG_AWAY = ROCKET_WIDTH / 2 75 | 76 | # SHIP 77 | SHIP_HEIGHT = ROCKET_WIDTH 78 | SHIP_WIDTH = SHIP_HEIGHT * 40 79 | 80 | # VIEWPORT 81 | VIEWPORT_H = 720 82 | VIEWPORT_W = 500 83 | H = 1.1 * START_HEIGHT * SCALE_S 84 | W = float(VIEWPORT_W) / VIEWPORT_H * H 85 | 86 | # SMOKE FOR VISUALS 87 | MAX_SMOKE_LIFETIME = 2 * FPS 88 | 89 | 90 | class ContactDetector(contactListener): 91 | def __init__(self, env): 92 | contactListener.__init__(self) 93 | self.env = env 94 | 95 | def BeginContact(self, contact): 96 | if ( 97 | self.env.water in [contact.fixtureA.body, contact.fixtureB.body] 98 | or self.env.lander in [contact.fixtureA.body, contact.fixtureB.body] 99 | or self.env.containers[0] in [contact.fixtureA.body, contact.fixtureB.body] 100 | or self.env.containers[1] in [contact.fixtureA.body, contact.fixtureB.body] 101 | ): 102 | self.env.game_over = True 103 | else: 104 | for i in range(2): 105 | if self.env.legs[i] in [contact.fixtureA.body, contact.fixtureB.body]: 106 | self.env.legs[i].ground_contact = True 107 | 108 | def EndContact(self, contact): 109 | for i in range(2): 110 | if self.env.legs[i] in [contact.fixtureA.body, contact.fixtureB.body]: 111 | self.env.legs[i].ground_contact = False 112 | 113 | 114 | class RocketLander(gym.Env): 115 | metadata = {"render.modes": ["human", "rgb_array"], "video.frames_per_second": FPS} 116 | 117 | def __init__(self): 118 | self._seed() 119 | self.viewer = None 120 | self.episode_number = 0 121 | 122 | self.world = Box2D.b2World() 123 | self.water = None 124 | self.lander = None 125 | self.engine = None 126 | self.ship = None 127 | self.legs = [] 128 | 129 | high = np.array([1, 1, 1, 1, 1, 1, 1, np.inf, np.inf, np.inf], dtype=np.float32) 130 | low = -high 131 | if not VEL_STATE: 132 | high = high[0:7] 133 | low = low[0:7] 134 | 135 | self.observation_space = spaces.Box(low, high, dtype=np.float32) 136 | 137 | if CONTINUOUS: 138 | self.action_space = spaces.Box(-1.0, +1.0, (3,), dtype=np.float32) 139 | else: 140 | self.action_space = spaces.Discrete(7) 141 | 142 | self.reset() 143 | 144 | def _seed(self, seed=None): 145 | self.np_random, seed = seeding.np_random(seed) 146 | return [seed] 147 | 148 | def _destroy(self): 149 | if not self.water: 150 | return 151 | self.world.contactListener = None 152 | self.world.DestroyBody(self.water) 153 | self.water = None 154 | self.world.DestroyBody(self.lander) 155 | self.lander = None 156 | self.world.DestroyBody(self.ship) 157 | self.ship = None 158 | self.world.DestroyBody(self.legs[0]) 159 | self.world.DestroyBody(self.legs[1]) 160 | self.legs = [] 161 | self.world.DestroyBody(self.containers[0]) 162 | self.world.DestroyBody(self.containers[1]) 163 | self.containers = [] 164 | 165 | def reset(self): 166 | self._destroy() 167 | self.world.contactListener_keepref = ContactDetector(self) 168 | self.world.contactListener = self.world.contactListener_keepref 169 | self.game_over = False 170 | self.prev_shaping = None 171 | self.throttle = 0 172 | self.gimbal = 0.0 173 | self.landed_ticks = 0 174 | self.stepnumber = 0 175 | self.smoke = [] 176 | 177 | # self.terrainheigth = self.np_random.uniform(H / 20, H / 10) 178 | self.terrainheigth = H / 20 179 | self.shipheight = self.terrainheigth + SHIP_HEIGHT 180 | # ship_pos = self.np_random.uniform(0, SHIP_WIDTH / SCALE) + SHIP_WIDTH / SCALE 181 | ship_pos = W / 2 182 | self.helipad_x1 = ship_pos - SHIP_WIDTH / 2 183 | self.helipad_x2 = self.helipad_x1 + SHIP_WIDTH 184 | self.helipad_y = self.terrainheigth + SHIP_HEIGHT 185 | 186 | self.water = self.world.CreateStaticBody( 187 | fixtures=fixtureDef( 188 | shape=polygonShape( 189 | vertices=( 190 | (0, 0), 191 | (W, 0), 192 | (W, self.terrainheigth), 193 | (0, self.terrainheigth), 194 | ) 195 | ), 196 | friction=0.1, 197 | restitution=0.0, 198 | ) 199 | ) 200 | self.water.color1 = rgb(70, 96, 176) 201 | 202 | self.ship = self.world.CreateStaticBody( 203 | fixtures=fixtureDef( 204 | shape=polygonShape( 205 | vertices=( 206 | (self.helipad_x1, self.terrainheigth), 207 | (self.helipad_x2, self.terrainheigth), 208 | (self.helipad_x2, self.terrainheigth + SHIP_HEIGHT), 209 | (self.helipad_x1, self.terrainheigth + SHIP_HEIGHT), 210 | ) 211 | ), 212 | friction=0.5, 213 | restitution=0.0, 214 | ) 215 | ) 216 | 217 | self.containers = [] 218 | for side in [-1, 1]: 219 | self.containers.append( 220 | self.world.CreateStaticBody( 221 | fixtures=fixtureDef( 222 | shape=polygonShape( 223 | vertices=( 224 | ( 225 | ship_pos + side * 0.95 * SHIP_WIDTH / 2, 226 | self.helipad_y, 227 | ), 228 | ( 229 | ship_pos + side * 0.95 * SHIP_WIDTH / 2, 230 | self.helipad_y + SHIP_HEIGHT, 231 | ), 232 | ( 233 | ship_pos 234 | + side * 0.95 * SHIP_WIDTH / 2 235 | - side * SHIP_HEIGHT, 236 | self.helipad_y + SHIP_HEIGHT, 237 | ), 238 | ( 239 | ship_pos 240 | + side * 0.95 * SHIP_WIDTH / 2 241 | - side * SHIP_HEIGHT, 242 | self.helipad_y, 243 | ), 244 | ) 245 | ), 246 | friction=0.2, 247 | restitution=0.0, 248 | ) 249 | ) 250 | ) 251 | self.containers[-1].color1 = rgb(206, 206, 2) 252 | 253 | self.ship.color1 = (0.2, 0.2, 0.2) 254 | 255 | initial_x = W / 2 + W * np.random.uniform(-0.3, 0.3) 256 | initial_y = H * 0.95 257 | self.lander = self.world.CreateDynamicBody( 258 | position=(initial_x, initial_y), 259 | angle=0.0, 260 | fixtures=fixtureDef( 261 | shape=polygonShape( 262 | vertices=( 263 | (-ROCKET_WIDTH / 2, 0), 264 | (+ROCKET_WIDTH / 2, 0), 265 | (ROCKET_WIDTH / 2, +ROCKET_HEIGHT), 266 | (-ROCKET_WIDTH / 2, +ROCKET_HEIGHT), 267 | ) 268 | ), 269 | density=1.0, 270 | friction=0.5, 271 | categoryBits=0x0010, 272 | maskBits=0x001, 273 | restitution=0.0, 274 | ), 275 | ) 276 | 277 | self.lander.color1 = rgb(230, 230, 230) 278 | 279 | for i in [-1, +1]: 280 | leg = self.world.CreateDynamicBody( 281 | position=(initial_x - i * LEG_AWAY, initial_y + ROCKET_WIDTH * 0.2), 282 | angle=(i * BASE_ANGLE), 283 | fixtures=fixtureDef( 284 | shape=polygonShape( 285 | vertices=( 286 | (0, 0), 287 | (0, LEG_LENGTH / 25), 288 | (i * LEG_LENGTH, 0), 289 | (i * LEG_LENGTH, -LEG_LENGTH / 20), 290 | (i * LEG_LENGTH / 3, -LEG_LENGTH / 7), 291 | ) 292 | ), 293 | density=1, 294 | restitution=0.0, 295 | friction=0.2, 296 | categoryBits=0x0020, 297 | maskBits=0x001, 298 | ), 299 | ) 300 | leg.ground_contact = False 301 | leg.color1 = (0.25, 0.25, 0.25) 302 | rjd = revoluteJointDef( 303 | bodyA=self.lander, 304 | bodyB=leg, 305 | localAnchorA=(i * LEG_AWAY, ROCKET_WIDTH * 0.2), 306 | localAnchorB=(0, 0), 307 | enableLimit=True, 308 | maxMotorTorque=2500.0, 309 | motorSpeed=-0.05 * i, 310 | enableMotor=True, 311 | ) 312 | djd = distanceJointDef( 313 | bodyA=self.lander, 314 | bodyB=leg, 315 | anchorA=(i * LEG_AWAY, ROCKET_HEIGHT / 8), 316 | anchorB=leg.fixtures[0].body.transform * (i * LEG_LENGTH, 0), 317 | collideConnected=False, 318 | frequencyHz=0.01, 319 | dampingRatio=0.9, 320 | ) 321 | if i == 1: 322 | rjd.lowerAngle = -SPRING_ANGLE 323 | rjd.upperAngle = 0 324 | else: 325 | rjd.lowerAngle = 0 326 | rjd.upperAngle = +SPRING_ANGLE 327 | leg.joint = self.world.CreateJoint(rjd) 328 | leg.joint2 = self.world.CreateJoint(djd) 329 | 330 | self.legs.append(leg) 331 | 332 | self.lander.linearVelocity = ( 333 | -self.np_random.uniform(0, INITIAL_RANDOM) 334 | * START_SPEED 335 | * (initial_x - W / 2) 336 | / W, 337 | -START_SPEED, 338 | ) 339 | 340 | self.lander.angularVelocity = (1 + INITIAL_RANDOM) * np.random.uniform(-1, 1) 341 | 342 | self.drawlist = ( 343 | self.legs + [self.water] + [self.ship] + self.containers + [self.lander] 344 | ) 345 | 346 | if CONTINUOUS: 347 | return self.step([0, 0, 0])[0] 348 | else: 349 | return self.step(6)[0] 350 | 351 | def step(self, action): 352 | 353 | self.force_dir = 0 354 | 355 | if CONTINUOUS: 356 | np.clip(action, -1, 1) 357 | self.gimbal += action[0] * 0.15 / FPS 358 | self.throttle += action[1] * 0.5 / FPS 359 | if action[2] > 0.5: 360 | self.force_dir = 1 361 | elif action[2] < -0.5: 362 | self.force_dir = -1 363 | else: 364 | if action == 0: 365 | self.gimbal += 0.01 366 | elif action == 1: 367 | self.gimbal -= 0.01 368 | elif action == 2: 369 | self.throttle += 0.01 370 | elif action == 3: 371 | self.throttle -= 0.01 372 | elif action == 4: # left 373 | self.force_dir = -1 374 | elif action == 5: # right 375 | self.force_dir = 1 376 | 377 | self.gimbal = np.clip(self.gimbal, -GIMBAL_THRESHOLD, GIMBAL_THRESHOLD) 378 | self.throttle = np.clip(self.throttle, 0.0, 1.0) 379 | self.power = ( 380 | 0 381 | if self.throttle == 0.0 382 | else MIN_THROTTLE + self.throttle * (1 - MIN_THROTTLE) 383 | ) 384 | 385 | # main engine force 386 | force_pos = (self.lander.position[0], self.lander.position[1]) 387 | force = ( 388 | -np.sin(self.lander.angle + self.gimbal) * MAIN_ENGINE_POWER * self.power, 389 | np.cos(self.lander.angle + self.gimbal) * MAIN_ENGINE_POWER * self.power, 390 | ) 391 | self.lander.ApplyForce(force=force, point=force_pos, wake=False) 392 | 393 | # control thruster force 394 | force_pos_c = self.lander.position + THRUSTER_HEIGHT * np.array( 395 | (np.sin(self.lander.angle), np.cos(self.lander.angle)) 396 | ) 397 | force_c = ( 398 | -self.force_dir * np.cos(self.lander.angle) * SIDE_ENGINE_POWER, 399 | self.force_dir * np.sin(self.lander.angle) * SIDE_ENGINE_POWER, 400 | ) 401 | self.lander.ApplyLinearImpulse(impulse=force_c, point=force_pos_c, wake=False) 402 | 403 | self.world.Step(1.0 / FPS, 60, 60) 404 | 405 | pos = self.lander.position 406 | vel_l = np.array(self.lander.linearVelocity) / START_SPEED 407 | vel_a = self.lander.angularVelocity 408 | x_distance = (pos.x - W / 2) / W 409 | y_distance = (pos.y - self.shipheight) / (H - self.shipheight) 410 | 411 | angle = (self.lander.angle / np.pi) % 2 412 | if angle > 1: 413 | angle -= 2 414 | 415 | state = [ 416 | 2 * x_distance, 417 | 2 * (y_distance - 0.5), 418 | angle, 419 | 1.0 if self.legs[0].ground_contact else 0.0, 420 | 1.0 if self.legs[1].ground_contact else 0.0, 421 | 2 * (self.throttle - 0.5), 422 | (self.gimbal / GIMBAL_THRESHOLD), 423 | ] 424 | if VEL_STATE: 425 | state.extend([vel_l[0], vel_l[1], vel_a]) 426 | 427 | # REWARD ------------------------------------------------------------------------------------------------------- 428 | 429 | # state variables for reward 430 | distance = np.linalg.norm( 431 | (3 * x_distance, y_distance) 432 | ) # weight x position more 433 | speed = np.linalg.norm(vel_l) 434 | groundcontact = self.legs[0].ground_contact or self.legs[1].ground_contact 435 | brokenleg = ( 436 | self.legs[0].joint.angle < 0 or self.legs[1].joint.angle > -0 437 | ) and groundcontact 438 | outside = abs(pos.x - W / 2) > W / 2 or pos.y > H 439 | fuelcost = 0.1 * (0.5 * self.power + abs(self.force_dir)) / FPS 440 | landed = ( 441 | self.legs[0].ground_contact and self.legs[1].ground_contact and speed < 0.1 442 | ) 443 | done = False 444 | 445 | reward = -fuelcost 446 | 447 | if outside or brokenleg: 448 | self.game_over = True 449 | 450 | if self.game_over: 451 | done = True 452 | else: 453 | # reward shaping 454 | shaping = -0.5 * (distance + speed + abs(angle) ** 2 + abs(vel_a) ** 2) 455 | shaping += 0.1 * (self.legs[0].ground_contact + self.legs[1].ground_contact) 456 | if self.prev_shaping is not None: 457 | reward += shaping - self.prev_shaping 458 | self.prev_shaping = shaping 459 | 460 | if landed: 461 | self.landed_ticks += 1 462 | else: 463 | self.landed_ticks = 0 464 | if self.landed_ticks == FPS: 465 | reward = 1.0 466 | done = True 467 | 468 | if done: 469 | reward += max(-1, 0 - 2 * (speed + distance + abs(angle) + abs(vel_a))) 470 | 471 | reward = np.clip(reward, -1, 1) 472 | 473 | # REWARD ------------------------------------------------------------------------------------------------------- 474 | 475 | self.stepnumber += 1 476 | 477 | return np.array(state), reward, done, {} 478 | 479 | def render(self, mode="human", close=False): 480 | if close: 481 | if self.viewer is not None: 482 | self.viewer.close() 483 | self.viewer = None 484 | return 485 | 486 | if self.viewer is None: 487 | 488 | self.viewer = rendering.Viewer(VIEWPORT_W, VIEWPORT_H) 489 | self.viewer.set_bounds(0, W, 0, H) 490 | 491 | sky = rendering.FilledPolygon(((0, 0), (0, H), (W, H), (W, 0))) 492 | self.sky_color = rgb(126, 150, 233) 493 | sky.set_color(*self.sky_color) 494 | self.sky_color_half_transparent = ( 495 | np.array((np.array(self.sky_color) + rgb(255, 255, 255))) / 2 496 | ) 497 | self.viewer.add_geom(sky) 498 | 499 | self.rockettrans = rendering.Transform() 500 | 501 | engine = rendering.FilledPolygon( 502 | ( 503 | (0, 0), 504 | (ENGINE_WIDTH / 2, -ENGINE_HEIGHT), 505 | (-ENGINE_WIDTH / 2, -ENGINE_HEIGHT), 506 | ) 507 | ) 508 | self.enginetrans = rendering.Transform() 509 | engine.add_attr(self.enginetrans) 510 | engine.add_attr(self.rockettrans) 511 | engine.set_color(0.4, 0.4, 0.4) 512 | self.viewer.add_geom(engine) 513 | 514 | self.fire = rendering.FilledPolygon( 515 | ( 516 | (ENGINE_WIDTH * 0.4, 0), 517 | (-ENGINE_WIDTH * 0.4, 0), 518 | (-ENGINE_WIDTH * 1.2, -ENGINE_HEIGHT * 5), 519 | (0, -ENGINE_HEIGHT * 8), 520 | (ENGINE_WIDTH * 1.2, -ENGINE_HEIGHT * 5), 521 | ) 522 | ) 523 | self.fire.set_color(*rgb(255, 230, 107)) 524 | self.firescale = rendering.Transform(scale=(1, 1)) 525 | self.firetrans = rendering.Transform(translation=(0, -ENGINE_HEIGHT)) 526 | self.fire.add_attr(self.firescale) 527 | self.fire.add_attr(self.firetrans) 528 | self.fire.add_attr(self.enginetrans) 529 | self.fire.add_attr(self.rockettrans) 530 | 531 | smoke = rendering.FilledPolygon( 532 | ( 533 | (ROCKET_WIDTH / 2, THRUSTER_HEIGHT * 1), 534 | (ROCKET_WIDTH * 3, THRUSTER_HEIGHT * 1.03), 535 | (ROCKET_WIDTH * 4, THRUSTER_HEIGHT * 1), 536 | (ROCKET_WIDTH * 3, THRUSTER_HEIGHT * 0.97), 537 | ) 538 | ) 539 | smoke.set_color(*self.sky_color_half_transparent) 540 | self.smokescale = rendering.Transform(scale=(1, 1)) 541 | smoke.add_attr(self.smokescale) 542 | smoke.add_attr(self.rockettrans) 543 | self.viewer.add_geom(smoke) 544 | 545 | self.gridfins = [] 546 | for i in (-1, 1): 547 | finpoly = ( 548 | (i * ROCKET_WIDTH * 1.1, THRUSTER_HEIGHT * 1.01), 549 | (i * ROCKET_WIDTH * 0.4, THRUSTER_HEIGHT * 1.01), 550 | (i * ROCKET_WIDTH * 0.4, THRUSTER_HEIGHT * 0.99), 551 | (i * ROCKET_WIDTH * 1.1, THRUSTER_HEIGHT * 0.99), 552 | ) 553 | gridfin = rendering.FilledPolygon(finpoly) 554 | gridfin.add_attr(self.rockettrans) 555 | gridfin.set_color(0.25, 0.25, 0.25) 556 | self.gridfins.append(gridfin) 557 | 558 | if self.stepnumber % round(FPS / 10) == 0 and self.power > 0: 559 | s = [ 560 | MAX_SMOKE_LIFETIME * self.power, # total lifetime 561 | 0, # current lifetime 562 | self.power * (1 + 0.2 * np.random.random()), # size 563 | np.array(self.lander.position) 564 | + self.power 565 | * ROCKET_WIDTH 566 | * 10 567 | * np.array( 568 | ( 569 | np.sin(self.lander.angle + self.gimbal), 570 | -np.cos(self.lander.angle + self.gimbal), 571 | ) 572 | ) 573 | + self.power * 5 * (np.random.random(2) - 0.5), 574 | ] # position 575 | self.smoke.append(s) 576 | 577 | for s in self.smoke: 578 | s[1] += 1 579 | if s[1] > s[0]: 580 | self.smoke.remove(s) 581 | continue 582 | t = rendering.Transform(translation=(s[3][0], s[3][1] + H * s[1] / 2000)) 583 | self.viewer.draw_circle( 584 | radius=0.05 * s[1] + s[2], 585 | color=self.sky_color 586 | + (1 - (2 * s[1] / s[0] - 1) ** 2) 587 | / 3 588 | * (self.sky_color_half_transparent - self.sky_color), 589 | ).add_attr(t) 590 | 591 | self.viewer.add_onetime(self.fire) 592 | for g in self.gridfins: 593 | self.viewer.add_onetime(g) 594 | 595 | for obj in self.drawlist: 596 | for f in obj.fixtures: 597 | trans = f.body.transform 598 | path = [trans * v for v in f.shape.vertices] 599 | self.viewer.draw_polygon(path, color=obj.color1) 600 | 601 | for l in zip(self.legs, [-1, 1]): 602 | path = [ 603 | self.lander.fixtures[0].body.transform 604 | * (l[1] * ROCKET_WIDTH / 2, ROCKET_HEIGHT / 8), 605 | l[0].fixtures[0].body.transform * (l[1] * LEG_LENGTH * 0.8, 0), 606 | ] 607 | self.viewer.draw_polyline( 608 | path, color=self.ship.color1, linewidth=1 if START_HEIGHT > 500 else 2 609 | ) 610 | 611 | self.viewer.draw_polyline( 612 | ( 613 | (self.helipad_x2, self.terrainheigth + SHIP_HEIGHT), 614 | (self.helipad_x1, self.terrainheigth + SHIP_HEIGHT), 615 | ), 616 | color=rgb(206, 206, 2), 617 | linewidth=1, 618 | ) 619 | 620 | self.rockettrans.set_translation(*self.lander.position) 621 | self.rockettrans.set_rotation(self.lander.angle) 622 | self.enginetrans.set_rotation(self.gimbal) 623 | self.firescale.set_scale(newx=1, newy=self.power * np.random.uniform(1, 1.3)) 624 | self.smokescale.set_scale(newx=self.force_dir, newy=1) 625 | 626 | return self.viewer.render(return_rgb_array=mode == "rgb_array") 627 | 628 | 629 | def rgb(r, g, b): 630 | return float(r) / 255, float(g) / 255, float(b) / 255 631 | -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | from setuptools import setup 2 | 3 | setup( 4 | name="gym_rocketlander", 5 | version="0.0.1", 6 | install_requires=["gym", "Box2D"], 7 | ) --------------------------------------------------------------------------------