refactor(log): improved logging workflow

src/brittle_star_project/environment/BrittleStarJaxEnvWrapper.py

@@ -34,6 +34,13 @@ class BrittleStarJaxEnvWrapper:
 
         self._action_rng = None
 
+        from experiment_logger import get_logger
+
+        self.logger = get_logger()
+        self.logger.info(
+            f"Initialized BrittleStarJaxEnvWrapper with {num_envs} envs on {backend.value}"
+        )
+
     @property
     def backend(self):
         return self._backend
@@ -51,6 +58,7 @@ class BrittleStarJaxEnvWrapper:
         return self._env.observation_space
 
     def reset(self, seed: int = 0):
+        self.logger.info(f"Resetting vectorized environment environments with seed {seed}")
         self._action_rng, env_rng = jax.random.split(jax.random.PRNGKey(seed), 2)
         env_rngs = jnp.array(jax.random.split(env_rng, self._num_envs))
         return self._vectorized_reset(rng=env_rngs)

src/brittle_star_project/environment/factory.py

@@ -98,9 +98,15 @@ class BrittleStarEnvFactory:
             case _:
                 raise ValueError(f"Unsupported task: {env_config.task}")
 
-        return env_class.from_morphology_and_arena(
+        env = env_class.from_morphology_and_arena(
             morphology=morphology,
             arena=arena,
             configuration=env_configuration,
             backend=backend.value,
         )
+
+        from experiment_logger import get_logger
+
+        get_logger().info(f"Created {env_config.task.value} env on backend {backend.value}")
+
+        return env