fixing the overcounting of cummulative cost

montrealrobotics · Sep 25, 2023 · e435e10 · e435e10
1 parent ef1bb58
commit e435e10
Showing 1 changed file with 7 additions and 2 deletions.
diff --git a/cleanrl/dqn.py b/cleanrl/dqn.py
@@ -418,8 +418,13 @@ def get_random_action():
                 if "episode" not in info:
                     continue
 
-                ep_cost = torch.sum(f_costs[last_step:global_step]).item()
-                cum_cost += torch.sum(f_costs).item()
+                if args.fine_tune_risk:
+                    ep_cost = torch.sum(f_costs)
+                    cum_cost += ep_cost
+                else:
+                    ep_cost = torch.sum(f_costs[last_step:global_step]).item()
+                    cum_cost = torch.sum(f_costs).item()
+
                 last_step = global_step
                 print(f"global_step={global_step}, episodic_return={info['episode']['r']}")
                 writer.add_scalar("charts/episodic_return", info["episode"]["r"], global_step)