Add gravity and wind as kwarg in Lunar Lander

gym/envs/box2d/lunar_lander.py

@@ -26,6 +26,7 @@
 SIDE_ENGINE_POWER = 0.6
 
 INITIAL_RANDOM = 1000.0  # Set 1500 to make game harder
+WIND_POWER = 15.0
 
 LANDER_POLY = [(-14, +17), (-17, 0), (-17, -10), (+17, -10), (+17, 0), (+14, +17)]
 LEG_AWAY = 20
@@ -151,12 +152,26 @@ class LunarLander(gym.Env, EzPickle):
 
     metadata = {"render_modes": ["human", "rgb_array"], "render_fps": FPS}
 
-    def __init__(self, continuous: bool = False):
+    def __init__(
+        self,
+        continuous: bool = False,
+        gravity: float = -10.0,
+        enable_wind: bool = False,
+    ):
         EzPickle.__init__(self)
+
+        assert (
+            -12.0 < gravity and gravity < 0.0
+        ), f"gravity (current value: {gravity}) must be between -12 and 0"
+        self.gravity = gravity
+
+        self.enable_wind = enable_wind
+        self.wind_idx = np.random.randint(-9999, 9999)
+
         self.screen = None
         self.clock = None
         self.isopen = True
-        self.world = Box2D.b2World()
+        self.world = Box2D.b2World(gravity=(0, gravity))
         self.moon = None
         self.lander = None
         self.particles = []
@@ -330,6 +345,25 @@ def _clean_particles(self, all):
             self.world.DestroyBody(self.particles.pop(0))
 
     def step(self, action):
+        # Update wind
+        if self.enable_wind and not (
+            self.legs[0].ground_contact or self.legs[1].ground_contact
+        ):
+            # the function used for wind is tanh(sin(2 k x) + sin(pi k x)),
+            # which is proven to never be periodic, k = 0.01
+            wind_mag = (
+                math.tanh(
+                    math.sin(0.02 * self.wind_idx)
+                    + (math.sin(math.pi * 0.01 * self.wind_idx))
+                )
+                * WIND_POWER
+            )
+            self.wind_idx += 1
+            self.lander.ApplyForceToCenter(
+                (wind_mag, 0.0),
+                True,
+            )
+
         if self.continuous:
             action = np.clip(action, -1, +1).astype(np.float32)
         else:
@@ -352,9 +386,8 @@ def step(self, action):
                 assert m_power >= 0.5 and m_power <= 1.0
             else:
                 m_power = 1.0
-            ox = (
-                tip[0] * (4 / SCALE + 2 * dispersion[0]) + side[0] * dispersion[1]
-            )  # 4 is move a bit downwards, +-2 for randomness
+            # 4 is move a bit downwards, +-2 for randomness
+            ox = +tip[0] * (4 / SCALE + 2 * dispersion[0]) + side[0] * dispersion[1]
             oy = -tip[1] * (4 / SCALE + 2 * dispersion[0]) - side[1] * dispersion[1]
             impulse_pos = (self.lander.position[0] + ox, self.lander.position[1] + oy)
             p = self._create_particle(
@@ -618,6 +651,11 @@ def heuristic(env, s):
 
 
 def demo_heuristic_lander(env, seed=None, render=False):
+
+    # wind power must be reduced for heuristic landing
+    global WIND_POWER
+    WIND_POWER = 1.0
+
     total_reward = 0
     steps = 0
     s = env.reset(seed=seed)

tests/envs/test_lunar_lander.py

@@ -18,6 +18,16 @@ def test_lunar_lander_continuous():
     _test_lander(LunarLander(continuous=True), seed=0)
 
 
+@pytest.mark.skipif(Box2D is None, reason="Box2D not installed")
+def test_lunar_lander_wind():
+    _test_lander(LunarLander(enable_wind=True), seed=0)
+
+
+@pytest.mark.skipif(Box2D is None, reason="Box2D not installed")
+def test_lunar_lander_wind_continuous():
+    _test_lander(LunarLander(continuous=True, enable_wind=True), seed=0)
+
+
 @pytest.mark.skipif(Box2D is None, reason="Box2D not installed")
 def _test_lander(env, seed=None, render=False):
     total_reward = demo_heuristic_lander(env, seed=seed, render=render)