working with the new setup

montrealrobotics · Nov 16, 2023 · 787f088 · 787f088
1 parent a03db8d
commit 787f088
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/cleanrl/ddpg_continuous_action.py b/cleanrl/ddpg_continuous_action.py
@@ -310,7 +310,7 @@ def convert_dict_to_tensor(data, device):
     np.random.seed(args.seed)
     torch.manual_seed(args.seed)
     torch.backends.cudnn.deterministic = args.torch_deterministic
-
+    torch.set_num_threads(4)
     device = torch.device("cuda" if torch.cuda.is_available() and args.cuda else "cpu")
 
     # env setup
@@ -467,9 +467,9 @@ def convert_dict_to_tensor(data, device):
                 success_rate.append(int(infos[0]["is_success"]))
                 score.append(info["episode"]['r'])
                 if args.fine_tune_risk != "None":
-                    print(f"global_step={global_step}, episodic_return={info['episode']['r']}, Replay buffer size = {len(risk_rb)}, Total cost={total_cost}")
+                    print(f"global_step={global_step}, episodic_return={info['episode']['r']}, Replay buffer size = {len(risk_rb)}, Total cost={total_cost}, Success Rate={np.mean(success_rate[-100:])}")
                 else:
-                    print(f"global_step={global_step}, episodic_return={info['episode']['r']}, Ep Cost = {info['cum_cost']}")
+                    print(f"global_step={global_step}, episodic_return={info['episode']['r']}, Ep Cost = {info['cost']}, Success Rate={np.mean(success_rate[-100:])}")
                 writer.add_scalar("charts/episodic_return", info["episode"]["r"], global_step)
                 writer.add_scalar("charts/episodic_length", info["episode"]["l"], global_step)
                 writer.add_scalar("charts/Total Success", num_successes, global_step)