Update loss fn

Author: Martin Schubert

src/invrs_gym/challenges/metalens/challenge.py

@@ -64,7 +64,9 @@ def loss(self, response: metalens_component.MetalensResponse) -> jnp.ndarray:
             enhancement = response.enhancement_ex
         else:
             enhancement = response.enhancement_ey
-        return -jnp.mean(enhancement)
+        # The `sqrt` is used to help avoid situtations where one wavelength has
+        # strong enhancement while the others have less.
+        return -jnp.mean(jnp.sqrt(enhancement))
 
     def distance_to_target(
         self, response: metalens_component.MetalensResponse