add cost info (#528)

decisionforce · Aug 22, 2021 · 6be95b0 · 6be95b0
1 parent be692ee
commit 6be95b0
Showing 1 changed file with 13 additions and 0 deletions.
diff --git a/pgdrive/envs/safe_pgdrive_env.py b/pgdrive/envs/safe_pgdrive_env.py
@@ -29,6 +29,19 @@ def default_config(self) -> Config:
             allow_add_new_key=True
         )
         return config
+    def __init__(self,config):
+        super(SafePGDriveEnv, self).__init__(config)
+        self.episode_cost = 0
+
+    def reset(self, *args, **kwargs):
+        self.episode_cost = 0
+        return super(SafePGDriveEnv, self).reset(*args, **kwargs)
+
+    def cost_function(self, vehicle_id: str):
+        cost, step_info = super(SafePGDriveEnv, self).cost_function(vehicle_id)
+        self.episode_cost += cost
+        step_info["total_cost"] = self.episode_cost
+        return cost, step_info
 
     def _post_process_config(self, config):
         config = super(SafePGDriveEnv, self)._post_process_config(config)