Merge pull request #147 from StafaH:update_mujoco_light_type

PiperOrigin-RevId: 774828564
Change-Id: I00f178c389e199a39b3a29579b77a579a493e6dc

Author: copybara-github

mujoco_playground/_src/manipulation/franka_emika_panda/randomize_vision.py

@@ -72,7 +72,7 @@ def domain_randomize(
       'cam_quat': 0,
       'light_pos': 0,
       'light_dir': 0,
-      'light_directional': 0,
+      'light_type': 0,
       'light_castshadow': 0,
   })
   rng = jax.random.key(0)
@@ -151,15 +151,15 @@ def rand(rng: jax.Array, light_position: jax.Array):
     ).astype(jp.float32)
 
     # No need to randomize into specular lighting
-    light_directional = jp.ones((nlight,))
+    light_type = jp.ones((nlight,))
 
     return (
         geom_rgba,
         geom_matid,
         cam_pos,
         cam_quat,
         light_dir,
-        light_directional,
+        light_type,
         light_castshadow,
     )
 
@@ -169,7 +169,7 @@ def rand(rng: jax.Array, light_position: jax.Array):
       cam_pos,
       cam_quat,
       light_dir,
-      light_directional,
+      light_type,
       light_castshadow,
   ) = rand(jax.random.split(rng, num_worlds), light_positions)
 
@@ -180,7 +180,7 @@ def rand(rng: jax.Array, light_position: jax.Array):
       'cam_quat': cam_quat,
       'light_pos': light_positions,
       'light_dir': light_dir,
-      'light_directional': light_directional,
+      'light_type': light_type,
       'light_castshadow': light_castshadow,
   })
 

mujoco_playground/_src/wrapper.py

@@ -196,7 +196,7 @@ def _identity_vision_randomization_fn(
       'geom_size': 0,
       'light_pos': 0,
       'light_dir': 0,
-      'light_directional': 0,
+      'light_type': 0,
       'light_castshadow': 0,
       'light_cutoff': 0,
   })
@@ -218,8 +218,8 @@ def _identity_vision_randomization_fn(
       'light_dir': jp.repeat(
           jp.expand_dims(mjx_model.light_dir, 0), num_worlds, axis=0
       ),
-      'light_directional': jp.repeat(
-          jp.expand_dims(mjx_model.light_directional, 0), num_worlds, axis=0
+      'light_type': jp.repeat(
+          jp.expand_dims(mjx_model.light_type, 0), num_worlds, axis=0
       ),
       'light_castshadow': jp.repeat(
           jp.expand_dims(mjx_model.light_castshadow, 0), num_worlds, axis=0
@@ -245,7 +245,7 @@ def _supplement_vision_randomization_fn(
       'geom_size',
       'light_pos',
       'light_dir',
-      'light_directional',
+      'light_type',
       'light_castshadow',
       'light_cutoff',
   ]
@@ -292,7 +292,7 @@ def __init__(
         'geom_size',
         'light_pos',
         'light_dir',
-        'light_directional',
+        'light_type',
         'light_castshadow',
         'light_cutoff',
     ]