fix: mypy complaints

scripts/train.py

@@ -10,7 +10,7 @@ from brittle_star_project.dataclasses import PPOArgs
 from brittle_star_project.trainers.PPOTrainer import PPOTrainer
 from brittle_star_project.environment.BrittleStarJaxEnvWrapper import BrittleStarJaxEnvWrapper
 
-from experiment_logger.unified_logger import get_logger
+from experiment_logger import UnifiedLogger
 from experiment_logger.config_utils import merge_config_with_cli, print_config
 
 
@@ -37,7 +37,7 @@ def get_git_hash() -> str:
         return (
             subprocess.check_output(["git", "rev-parse", "--short", "HEAD"]).decode("ascii").strip()
         )
-    except subprocess.CalledProcessError | UnicodeDecodeError:
+    except (subprocess.CalledProcessError, UnicodeDecodeError):
         return "none"
 
 
@@ -59,8 +59,8 @@ if __name__ == "__main__":
     os.makedirs(run_dir, exist_ok=True)
 
     # Initialize Global Logger
-    logger = get_logger()
-    logger.init(
+    logger = UnifiedLogger(
+        config=vars(args),
         project_name=args.wandb_project_name,  # or default PPO-Modularity if missing
         run_name=run_name,
         base_dir=os.path.dirname(run_dir),

src/brittle_star_project/dataclasses/PPOArgs.py

@@ -22,9 +22,6 @@ class PPOArgs:
     # the directory to save the experiment results
     run_dir: str | None = None
 
-    # how often to save checkpoints (0 to disable)
-    checkpoint_frequency: int = 0
-
     # seed of the experiment
     seed: int = 1
 

src/brittle_star_project/render/renderer.py

@@ -36,6 +36,9 @@ def simulate_policy(
 
     import mujoco.viewer
 
+    if state is None:
+        raise ValueError("A valid environment state must be provided.")
+
     model = state.mj_model
     data = state.mj_data
 

src/brittle_star_project/trainers/PPOTrainer.py

@@ -312,14 +312,14 @@ class PPOTrainer:
     def _init_episode_stats(self) -> EpisodeStatistics:
         self.logger.info("[EPISODE STATS]: Initializing episode stats...")
 
-        return EpisodeStatistics(
+        return EpisodeStatistics(  # type: ignore[call-arg]
             episode_returns=jnp.zeros(self.args.num_envs, dtype=jnp.float32),
             episode_lengths=jnp.zeros(self.args.num_envs, dtype=jnp.int32),
             returned_episode_returns=jnp.zeros(self.args.num_envs, jnp.float32),
             returned_episode_lengths=jnp.zeros(self.args.num_envs, dtype=jnp.int32),
         )
 
-    def _rollout(self, env_state, next_obs, next_done) -> tuple[Storage, ...]:
+    def _rollout(self, env_state, next_obs, next_done) -> tuple[Any, ...]:
         return self._rollout_jit(
             self.agent_state,
             self.episode_stats,

src/experiment_logger/README.md

@@ -9,17 +9,20 @@ This library is designed to be a standalone package that decouples the logging l
 The recommended way to use the logger is through the `get_logger()` singleton:
 
 ```python
-from experiment_logger import get_logger
+from experiment_logger import UnifiedLogger, get_logger
 
-logger = get_logger()
 # Initialize at the start of your script (e.g., in train.py)
-logger.init(
-    project_name="MyProject",
+logger = UnifiedLogger(
     run_name="my_experiment_run",
+    config={"learning_rate": 3e-4},
+    project_name="MyProject",
     base_dir="runs",
     use_wandb=True
 )
 
+# In other files, retrieve the initialized singleton:
+# logger = get_logger()
+
 # Log metrics (Scalar values, numpy scalars, or JAX types)
 logger.log({"loss": 0.5, "accuracy": 0.98}, step=100)
 

src/experiment_logger/config_utils.py

@@ -44,10 +44,10 @@ def dataclass_from_dict(cls: Type[T], config_dict: Dict[str, Any]) -> T:
         raise ValueError(f"{cls} is not a dataclass")
 
     # Get field names and types
-    field_map = {f.name: f for f in fields(cls)}
+    field_map = {f.name: f for f in fields(cls)}  # type: ignore
 
     # Filter config to only include valid fields
-    filtered_config = {}
+    filtered_config: Dict[str, Any] = {}
     for key, value in config_dict.items():
         if key in field_map:
             field = field_map[key]
@@ -64,7 +64,7 @@ def dataclass_from_dict(cls: Type[T], config_dict: Dict[str, Any]) -> T:
                     if field.type is bool and isinstance(value, str):
                         filtered_config[key] = value.lower() in ("true", "1", "yes", "on")
                     else:
-                        filtered_config[key] = field.type(value) if value is not None else None
+                        filtered_config[key] = field.type(value) if value is not None else None  # type: ignore
             except (ValueError, TypeError) as e:
                 log.warning(f"Could not convert {key}={value} to {field.type}: {e}")
                 filtered_config[key] = value
@@ -74,7 +74,7 @@ def dataclass_from_dict(cls: Type[T], config_dict: Dict[str, Any]) -> T:
     return cls(**filtered_config)
 
 
-def merge_config_with_cli(config_class: Type[T], config_file: str = None) -> T:
+def merge_config_with_cli(config_class: Type[T], config_file: str | None = None) -> T:
     """Merge YAML config with CLI arguments, with CLI taking precedence.
 
     Args:
@@ -107,16 +107,16 @@ def merge_config_with_cli(config_class: Type[T], config_file: str = None) -> T:
 
     # Create default instance to know what the defaults are
     default_instance = config_class()
-    default_dict = {f.name: getattr(default_instance, f.name) for f in fields(config_class)}
+    default_dict = {f.name: getattr(default_instance, f.name) for f in fields(config_class)}  # type: ignore
 
     # Parse CLI args
     cli_instance = tyro.cli(config_class)
-    cli_dict = {f.name: getattr(cli_instance, f.name) for f in fields(config_class)}
+    cli_dict = {f.name: getattr(cli_instance, f.name) for f in fields(config_class)}  # type: ignore
 
     # Merge configs: YAML as base, CLI overrides non-default values
     final_config = {}
 
-    for field in fields(config_class):
+    for field in fields(config_class):  # type: ignore
         field_name = field.name
         default_value = default_dict[field_name]
         yaml_value = yaml_config.get(field_name, default_value)

src/experiment_logger/unified_logger.py

@@ -224,7 +224,7 @@ class UnifiedLogger:
         self._log_to_stdout(metrics_with_metadata)
 
         # Log to WandB
-        if self.wandb_available:
+        if self.wandb_run is not None:
             try:
                 self.wandb_run.log(metrics, step=step, commit=commit)
             except Exception as e:
@@ -306,7 +306,7 @@ class UnifiedLogger:
             self.info(f"Checkpoint saved: {checkpoint_path}")
 
             # Log to WandB as artifact
-            if self.wandb_available:
+            if self.wandb_run is not None:
                 try:
                     import wandb
 
@@ -342,7 +342,7 @@ class UnifiedLogger:
             self.info(f"Final model saved: {final_model_path}")
 
             # Log to WandB
-            if self.wandb_available:
+            if self.wandb_run is not None:
                 try:
                     import wandb