minor fix

Author: rithwik-db

compose_rl/utils/mlflow_utils.py

@@ -28,7 +28,7 @@ def get_valid_mlflow_experiment_name(config: Any) -> str:
 
 def get_mlflow_relative_path_for_save_folder(save_folder: str) -> str:
     """Returns the relative path for the given save folder
-    
+
     Relative in mlflow need to be of the format: `artifacts/{relative_path}`
     """
     return os.path.join('artifacts', save_folder.lstrip('/'))
@@ -131,4 +131,4 @@ def setup_mlflow(config: Any):
     os.environ['MLFLOW_EXPERIMENT_ID'] = experiment_id
     os.environ['MLFLOW_TRACKING_URI'] = 'databricks'
 
-    dist.destroy_process_group()
+    dist.destroy_process_group()

test_single_controller_ppo.py

@@ -477,7 +477,7 @@ def __init__(
         self.max_gen_len = self.prompt_handler_config['max_gen_len']
         self.iter_num = 0
 
-        self.local_save_folder = config.save_folder
+        self.local_save_folder = os.path.join(config.save_folder, 'RolloutAgent')
         self.mlflow_absolute_save_folder = get_mlflow_absolute_path_for_save_folder(self.local_save_folder)
         self.mlflow_relative_save_folder = get_mlflow_relative_path_for_save_folder(self.local_save_folder)