fix(PPOTrainer.py, PPOArgs.py): added hyperparam config path to cli args, fixed missing arguments error in _step call

experiments/PPOTrainer.py

@@ -500,12 +500,14 @@ class PPOTrainer:
 
         iter_bar = tqdm.tqdm(
             range(1, self.args.num_iterations + 1),
-            disable=not sys.stdout.isatty(),
+            disable=not is_tty,
         )
         for iteration in iter_bar:
             iteration_time_start = time.time()
 
-            env_state, next_obs, next_done, loss_info = self._step(env_state, next_obs, next_done)
+            env_state, next_obs, next_done, loss_info = self._step(
+                env_state, next_obs, next_done, is_tty=is_tty, iteration=iteration
+            )
 
             if not is_tty and iteration == 1:
                 print(f">>> [HPC] First rollout completed: {time.ctime()}", flush=True)

experiments/train.py

@@ -17,11 +17,14 @@ def make_env(config_path: str | None, num_envs: int) -> BrittleStarJaxEnvWrapper
     return BrittleStarJaxEnvWrapper.from_config(config_path, num_envs=num_envs)
 
 
-def parse_args() -> PPOArgs:
+def parse_args(log: bool = True) -> PPOArgs:
     temp_args = tyro.cli(PPOArgs)
 
-    if temp_args.env_config_path is not None:
-        with open(temp_args.env_config_path, "r") as f:
+    if temp_args.hyperparameter_config_path is not None:
+        if log:
+            print(f"Loading hyperparameter config from {temp_args.hyperparameter_config_path}")
+
+        with open(temp_args.hyperparameter_config_path, "r") as f:
             config = yaml.safe_load(f)
             if config:
                 # parse PPOArgs with defaults from yaml.
@@ -32,6 +35,9 @@ def parse_args() -> PPOArgs:
                 # Reparse CLI to ensure they OVERRIDE the yaml
                 args = tyro.cli(PPOArgs, default=temp_args)
     else:
+        if log:
+            print("No hyperparameter config provided, using default config")
+
         args = temp_args
     return args
 

src/brittle_star_project/dataclasses/PPOArgs.py

@@ -13,6 +13,9 @@ class PPOArgs:
     # path to environment config file, if None, use default config
     env_config_path: str | None = None
 
+    # path to hyperparameter config file (yaml), if None, use default config
+    hyperparameter_config_path: str | None = None
+
     # the name of this experiment
     exp_name: str = "brittle_star_ppo"