Enhance modify_scene_fn argument handling

Add conditional handling for modify_scene_fn arguments.

Author: Caedael

mujoco_playground/_src/mjx_env.py

@@ -343,7 +343,10 @@ def get_image(state, modify_scn_fn=None) -> np.ndarray:
     mujoco.mj_forward(mj_model, d)
     renderer.update_scene(d, camera=camera, scene_option=scene_option)
     if modify_scn_fn is not None:
-      modify_scn_fn(renderer.scene)
+      if modify_scene_fn.__code__.co_argcount == 1:
+        modify_scn_fn(renderer.scene)
+      elif modify_scene_fn.__code__.co_argcount == 2:
+        modify_scn_fn(renderer.scene, state)
     return renderer.render()
 
   if isinstance(trajectory, list):