Add energy reward from benni/joy (#864)

* Add energy reward from benni/joy

* Add tests

flow/core/kernel/vehicle/aimsun.py

@@ -401,6 +401,10 @@ def add(self, veh_id, type_id, edge, pos, lane, speed):
             self.num_type[type_id] = 1
             self.total_num_type[type_id] = 1
 
+    def reset(self):
+        """See parent class."""
+        pass
+
     def remove(self, aimsun_id):
         """See parent class."""
         veh_id = self._id_aimsun2flow[aimsun_id]

flow/core/kernel/vehicle/base.py

@@ -100,6 +100,10 @@ def add(self, veh_id, type_id, edge, pos, lane, speed):
         """
         raise NotImplementedError
 
+    def reset(self):
+        """Reset any additional state that needs to be reset."""
+        raise NotImplementedError
+
     def remove(self, veh_id):
         """Remove a vehicle.
 

flow/core/kernel/vehicle/traci.py

@@ -82,6 +82,9 @@ def __init__(self,
         except AttributeError:
             self._force_color_update = False
 
+        # old speeds used to compute accelerations
+        self.previous_speeds = {}
+
     def initialize(self, vehicles):
         """Initialize vehicle state information.
 
@@ -127,8 +130,11 @@ def update(self, reset):
             specifies whether the simulator was reset in the last simulation
             step
         """
+        # copy over the previous speeds
+
         vehicle_obs = {}
         for veh_id in self.__ids:
+            self.previous_speeds[veh_id] = self.get_speed(veh_id)
             vehicle_obs[veh_id] = \
                 self.kernel_api.vehicle.getSubscriptionResults(veh_id)
         sim_obs = self.kernel_api.simulation.getSubscriptionResults()
@@ -362,6 +368,10 @@ def _add_departed(self, veh_id, veh_type):
 
         return new_obs
 
+    def reset(self):
+        """See parent class."""
+        self.previous_speeds = {}
+
     def remove(self, veh_id):
         """See parent class."""
         # remove from sumo
@@ -513,6 +523,12 @@ def get_departed_ids(self):
         else:
             return 0
 
+    def get_previous_speed(self, veh_id, error=-1001):
+        """See parent class."""
+        if isinstance(veh_id, (list, np.ndarray)):
+            return [self.get_previous_speed(vehID, error) for vehID in veh_id]
+        return self.previous_speeds.get(veh_id, 0)
+
     def get_speed(self, veh_id, error=-1001):
         """See parent class."""
         if isinstance(veh_id, (list, np.ndarray)):

flow/core/rewards.py

@@ -304,3 +304,29 @@ def punish_rl_lane_changes(env, penalty=1):
             total_lane_change_penalty -= penalty
 
     return total_lane_change_penalty
+
+
+def energy_consumption(env, gain=.001):
+    """Calculate power consumption of a vehicle.
+
+    Assumes vehicle is an average sized vehicle.
+    The power calculated here is the lower bound of the actual power consumed
+    by a vehicle.
+    """
+    power = 0
+
+    M = 1200  # mass of average sized vehicle (kg)
+    g = 9.81  # gravitational acceleration (m/s^2)
+    Cr = 0.005  # rolling resistance coefficient
+    Ca = 0.3  # aerodynamic drag coefficient
+    rho = 1.225  # air density (kg/m^3)
+    A = 2.6  # vehicle cross sectional area (m^2)
+    for veh_id in env.k.vehicle.get_ids():
+        speed = env.k.vehicle.get_speed(veh_id)
+        prev_speed = env.k.vehicle.get_previous_speed(veh_id)
+
+        accel = abs(speed - prev_speed) / env.sim_step
+
+        power += M * speed * accel + M * g * Cr * speed + 0.5 * rho * A * Ca * speed ** 3
+
+    return -gain * power

flow/envs/base.py

@@ -487,6 +487,9 @@ def reset(self):
             except (FatalTraCIError, TraCIException):
                 print("Error during start: {}".format(traceback.format_exc()))
 
+        # do any additional resetting of the vehicle class needed
+        self.k.vehicle.reset()
+
         # reintroduce the initial vehicles to the network
         for veh_id in self.initial_ids:
             type_id, edge, lane_index, pos, speed = \

flow/envs/multiagent/base.py

@@ -205,6 +205,9 @@ def reset(self, new_inflow_rate=None):
             except (FatalTraCIError, TraCIException):
                 print("Error during start: {}".format(traceback.format_exc()))
 
+        # do any additional resetting of the vehicle class needed
+        self.k.vehicle.reset()
+
         # reintroduce the initial vehicles to the network
         for veh_id in self.initial_ids:
             type_id, edge, lane_index, pos, speed = \

tests/fast_tests/test_rewards.py

@@ -7,6 +7,7 @@
 from flow.core.rewards import average_velocity, min_delay
 from flow.core.rewards import desired_velocity, boolean_action_penalty
 from flow.core.rewards import penalize_near_standstill, penalize_standstill
+from flow.core.rewards import energy_consumption
 
 os.environ["TEST_FLAG"] = "True"
 
@@ -151,6 +152,31 @@ def test_penalize_near_standstill(self):
         self.assertEqual(penalize_near_standstill(env, thresh=2), -10)
         self.assertEqual(penalize_near_standstill(env, thresh=0.5), -9)
 
+    def test_energy_consumption(self):
+        """Test the energy consumption method."""
+        vehicles = VehicleParams()
+        vehicles.add("test", num_vehicles=10)
+
+        env_params = EnvParams(additional_params={
+            "target_velocity": 10, "max_accel": 1, "max_decel": 1,
+            "sort_vehicles": False})
+
+        env, _, _ = ring_road_exp_setup(vehicles=vehicles,
+                                        env_params=env_params)
+
+        # check the penalty is zero at speed zero
+        self.assertEqual(energy_consumption(env, gain=1), 0)
+
+        # change the speed of one vehicle
+        env.k.vehicle.test_set_speed("test_0", 1)
+        self.assertEqual(energy_consumption(env), -12.059337750000001)
+
+        # check that stepping change the previous speeds and increases the energy consumption
+        env.step(rl_actions=None)
+        env.step(rl_actions=None)
+        self.assertGreater(env.k.vehicle.get_previous_speed("test_0"), 0.0)
+        self.assertLess(energy_consumption(env), -12.059337750000001)
+
     def test_boolean_action_penalty(self):
         """Test the boolean_action_penalty method."""
         actions = [False, False, False, False, False]