Merge pull request #8426 from initrv/add-sb3-learn-progress-bar

Add sb3 learn progress bar
2023-04-17 10:22:49 +02:00
parent 45193127e3 9caa74c796
commit daabc8ffbe
4 changed files with 6 additions and 1 deletions
@@ -600,6 +600,7 @@ CONF_SCHEMA = {
                        "policy_type": {"type": "string", "default": "MlpPolicy"},
                        "net_arch": {"type": "array", "default": [128, 128]},
                        "randomize_starting_position": {"type": "boolean", "default": False},
+                        "progress_bar": {"type": "boolean", "default": True},
                        "model_reward_parameters": {
                            "type": "object",
                            "properties": {
@@ -71,7 +71,8 @@ class ReinforcementLearner(BaseReinforcementLearningModel):

        model.learn(
            total_timesteps=int(total_timesteps),
-            callback=[self.eval_callback, self.tensorboard_callback]
+            callback=[self.eval_callback, self.tensorboard_callback],
+            progress_bar=self.rl_config.get('progress_bar', False)
        )

        if Path(dk.data_path / "best_model.zip").is_file():