Allowed weighted objectives

nucon/rl.py

@@ -7,8 +7,6 @@ from .core import Nucon, BreakerStatus, PumpStatus, PumpDryStatus, PumpOverloadS
 
 Objectives = {
     "null": lambda obs: 0,
-    "coeff": lambda obj, coeff: lambda obs: obj(obs) * coeff,
-
     "max_power": lambda obs: obs["GENERATOR_0_KW"] + obs["GENERATOR_1_KW"] + obs["GENERATOR_2_KW"],
     "episode_time": lambda obs: obs["EPISODE_TIME"],
 }
@@ -20,11 +18,14 @@ Parameterized_Objectives = {
 class NuconEnv(gym.Env):
     metadata = {'render_modes': ['human']}
 
-    def __init__(self, render_mode=None, seconds_per_step=5, objectives=['null'], terminators=['null'], terminate_above=0):
+    def __init__(self, render_mode=None, seconds_per_step=5, objectives=['null'], terminators=['null'], objective_weights=None, terminate_above=0):
         super().__init__()
 
         self.render_mode = render_mode
         self.seconds_per_step = seconds_per_step
+        if objective_weights is None:
+            objective_weights = [1.0 for objective in objectives]
+        self.objective_weights = objective_weights
         self.terminate_at = terminate_at
         
         # Define observation space
@@ -92,9 +93,11 @@ class NuconEnv(gym.Env):
         return obs
 
     def _get_info(self):
-        info = {'objectives': {}}
-        for objective in self.objectives:
-            info['objectives'][objective.__name__] = objective(self._get_obs())
+        info = {'objectives': {}, 'objectives_weighted': {}}
+        for objective, weight in zip(self.objectives, self.objective_weights):
+            obj = objective(self._get_obs())
+            info['objectives'][objective.__name__] = obj
+            info['objectives_weighted'][objective.__name__] = obj * weight 
         return info
     
     def reset(self, seed=None, options=None):
@@ -120,7 +123,7 @@ class NuconEnv(gym.Env):
         terminated = np.sum([terminator(observation) for terminator in self.terminators]) > self.terminate_above
         truncated = False
         info = self._get_info()
-        reward = sum(obj for obj in info['objectives'].values())
+        reward = sum(obj for obj in info['objectives_weighted'].values())
 
         self._total_steps += 1
         time.sleep(self.seconds_per_step)