diff --git a/scripts/demos/raycast_sensor.py b/scripts/demos/raycast_sensor.py
index 26b74ef14..c744e6352 100644
--- a/scripts/demos/raycast_sensor.py
+++ b/scripts/demos/raycast_sensor.py
@@ -113,7 +113,7 @@ def create_env_cfg(
         border_width=0.25,
       ),
     },
-    add_lights=False,
+    add_lights=True,
   )
 
   terrain_cfg = TerrainImporterCfg(
diff --git a/src/mjlab/terrains/config.py b/src/mjlab/terrains/config.py
index 24124f26e..06dc4f405 100644
--- a/src/mjlab/terrains/config.py
+++ b/src/mjlab/terrains/config.py
@@ -10,59 +10,60 @@
   num_rows=10,
   num_cols=20,
   sub_terrains={
-    "flat": terrain_gen.BoxFlatTerrainCfg(proportion=0.4),
+    "flat": terrain_gen.BoxFlatTerrainCfg(proportion=0.2),
     "pyramid_stairs": terrain_gen.BoxPyramidStairsTerrainCfg(
-      proportion=0.3,
+      proportion=0.2,
       step_height_range=(0.0, 0.1),
       step_width=0.3,
       platform_width=3.0,
       border_width=1.0,
     ),
     "pyramid_stairs_inv": terrain_gen.BoxInvertedPyramidStairsTerrainCfg(
-      proportion=0.3,
+      proportion=0.2,
       step_height_range=(0.0, 0.1),
       step_width=0.3,
       platform_width=3.0,
       border_width=1.0,
     ),
-    # NOTE: Heightfield terrains are currently disabled due to compilation issues
-    # in mujoco-warp.
-    # "hf_pyramid_slope": terrain_gen.HfPyramidSlopedTerrainCfg(
-    #   proportion=0.1,
-    #   slope_range=(0.0, 1.0),
-    #   platform_width=2.0,
-    #   border_width=0.25,
-    # ),
-    # "hf_pyramid_slope_inv": terrain_gen.HfPyramidSlopedTerrainCfg(
-    #   proportion=0.1,
-    #   slope_range=(0.0, 1.0),
-    #   platform_width=2.0,
-    #   border_width=0.25,
-    #   inverted=True,
-    # ),
-    # "random_rough": terrain_gen.HfRandomUniformTerrainCfg(
-    #   proportion=0.2,
-    #   noise_range=(0.02, 0.10),
-    #   noise_step=0.02,
-    #   border_width=0.25,
-    # ),
-    # "wave_terrain": terrain_gen.HfWaveTerrainCfg(
-    #   proportion=0.2,
-    #   amplitude_range=(0.0, 0.2),
-    #   num_waves=4,
-    #   border_width=0.25,
-    # ),
+    "hf_pyramid_slope": terrain_gen.HfPyramidSlopedTerrainCfg(
+      proportion=0.1,
+      slope_range=(0.0, 1.0),
+      platform_width=2.0,
+      border_width=0.25,
+    ),
+    "hf_pyramid_slope_inv": terrain_gen.HfPyramidSlopedTerrainCfg(
+      proportion=0.1,
+      slope_range=(0.0, 1.0),
+      platform_width=2.0,
+      border_width=0.25,
+      inverted=True,
+    ),
+    "random_rough": terrain_gen.HfRandomUniformTerrainCfg(
+      proportion=0.1,
+      noise_range=(0.02, 0.10),
+      noise_step=0.02,
+      border_width=0.25,
+    ),
+    "wave_terrain": terrain_gen.HfWaveTerrainCfg(
+      proportion=0.1,
+      amplitude_range=(0.0, 0.2),
+      num_waves=4,
+      border_width=0.25,
+    ),
   },
-  add_lights=False,
+  add_lights=True,
 )
 
 
 if __name__ == "__main__":
   import mujoco.viewer
+  import torch
+
+  device = "cuda" if torch.cuda.is_available() else "cpu"
 
   terrain_cfg = TerrainImporterCfg(
     terrain_type="generator",
     terrain_generator=ROUGH_TERRAINS_CFG,
   )
-  terrain = TerrainImporter(terrain_cfg, device="cuda:0")
+  terrain = TerrainImporter(terrain_cfg, device=device)
   mujoco.viewer.launch(terrain.spec.compile())
diff --git a/src/mjlab/terrains/terrain_generator.py b/src/mjlab/terrains/terrain_generator.py
index 049593c47..137def285 100644
--- a/src/mjlab/terrains/terrain_generator.py
+++ b/src/mjlab/terrains/terrain_generator.py
@@ -253,25 +253,10 @@ def _add_grid_lights(self, spec: mujoco.MjSpec) -> None:
 
     total_width = self.cfg.size[0] * self.cfg.num_rows
     total_height = self.cfg.size[1] * self.cfg.num_cols
-
     light_height = max(total_width, total_height) * 0.6
 
-    positions = [
-      (0, 0),  # Center.
-      (-total_width * 0.5, -total_height * 0.5),  # Bottom-left.
-      (-total_width * 0.5, total_height * 0.5),  # Top-left.
-      (total_width * 0.5, -total_height * 0.5),  # Bottom-right.
-      (total_width * 0.5, total_height * 0.5),  # Top-right.
-    ]
-
-    for i, (x, y) in enumerate(positions):
-      intensity = 0.4 if i == 0 else 0.2
-
-      spec.body("terrain").add_light(
-        pos=(x, y, light_height),
-        type=mujoco.mjtLightType.mjLIGHT_SPOT,
-        diffuse=(intensity, intensity, intensity * 0.95),
-        specular=(0.1, 0.1, 0.1),
-        cutoff=70,
-        exponent=2,
-      )
+    spec.body("terrain").add_light(
+      pos=(0, 0, light_height),
+      type=mujoco.mjtLightType.mjLIGHT_DIRECTIONAL,
+      dir=(0, 0, -1),
+    )