From 75c6d5bacda319a39a45f5aed3acbee49c3435b6 Mon Sep 17 00:00:00 2001 From: Tibo De Peuter Date: Wed, 8 Apr 2026 23:37:01 +0200 Subject: [PATCH] fix: mypy complaints --- scripts/train.py | 8 ++++---- src/brittle_star_project/dataclasses/PPOArgs.py | 3 --- src/brittle_star_project/render/renderer.py | 3 +++ src/brittle_star_project/trainers/PPOTrainer.py | 4 ++-- src/experiment_logger/README.md | 11 +++++++---- src/experiment_logger/config_utils.py | 14 +++++++------- src/experiment_logger/unified_logger.py | 6 +++--- 7 files changed, 26 insertions(+), 23 deletions(-) diff --git a/scripts/train.py b/scripts/train.py index 39b062d..944956b 100644 --- a/scripts/train.py +++ b/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), diff --git a/src/brittle_star_project/dataclasses/PPOArgs.py b/src/brittle_star_project/dataclasses/PPOArgs.py index 68265fa..036b44f 100644 --- a/src/brittle_star_project/dataclasses/PPOArgs.py +++ b/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 diff --git a/src/brittle_star_project/render/renderer.py b/src/brittle_star_project/render/renderer.py index 4d97ef1..91e669c 100644 --- a/src/brittle_star_project/render/renderer.py +++ b/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 diff --git a/src/brittle_star_project/trainers/PPOTrainer.py b/src/brittle_star_project/trainers/PPOTrainer.py index bd97b9e..d4297de 100644 --- a/src/brittle_star_project/trainers/PPOTrainer.py +++ b/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, diff --git a/src/experiment_logger/README.md b/src/experiment_logger/README.md index ce31d69..4e13b6b 100644 --- a/src/experiment_logger/README.md +++ b/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) diff --git a/src/experiment_logger/config_utils.py b/src/experiment_logger/config_utils.py index 7a92c0e..4c5b79f 100644 --- a/src/experiment_logger/config_utils.py +++ b/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) diff --git a/src/experiment_logger/unified_logger.py b/src/experiment_logger/unified_logger.py index 40229c5..051fe22 100644 --- a/src/experiment_logger/unified_logger.py +++ b/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