fix: mypy complaints

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,