BUG FIXES: moving around when policy is loaded, adding BaseAlgorithm abstract classes, correcting metrics, normalizng actions if idm environment is a normalized action one, manually updating environment graph when using idm, implementing idm forward class

src/baselines/rule_policies.py

@@ -1,6 +1,6 @@
 from stable_baselines3.common.base_class import BaseAlgorithm
-from intersim.envs.intersimple import Intersimple
-from typing import Tuple, Optional
+from intersim.envs.intersimple import Intersimple, NormalizedActionSpace
+from typing import Tuple, Optional, List
 import numpy as np
 
 class PControllerPolicy(BaseAlgorithm):
@@ -10,11 +10,17 @@ class PControllerPolicy(BaseAlgorithm):
         """
         Initialize policy with pointer to environment it will run on
         """
-        assert(isinstance(env, Intersimple), 'Environment is not an intersimple environment')
+        assert isinstance(env, Intersimple), 'Environment is not an intersimple environment'
         self._env = env
         self.target_v = 8.94 # m/s
         self.attn_weight = 20
     
+    # BaseAlgorithm abstract methods
+    def _setup_model(self): 
+        return None
+    def learn(self, *args, **kwargs):
+        return self
+
     def predict(self, observation: np.ndarray, *args, **kwargs):
         """
         Generate action, state from observation
@@ -59,37 +65,45 @@ class IDMRulePolicy(BaseAlgorithm):
     
     """
 
-    def __init__(self, env: Intersimple, target_speed: float= 8.94, t_future:float=0):
+    def __init__(self, env: Intersimple, 
+        target_speed:float= 8.94, 
+        t_future:List[float]=[0., 1., 2., 3.], 
+        half_angle:float=60.):
         """
         Initialize policy with pointer to environment it will run on and target speed
 
         Args:
             env (Intersimple): intersimple environment which IDM runs on
             target_speed (float): target speed in roundabout (default: 8.94=20 mph)
-            t_future (float): future time at which to compare closest
+            t_future (List[float]): list of future time at which to compare closest vehicle
+            half_angle (float): half angle to look inside for closest vehicle
         """
-        # assert(isinstance(env, Intersimple), 'Environment is not an intersimple environment')
+
         self._env = env
-
-        
-        assert(t_future >=0, 'negative target speed')
         self.t_future = t_future
-
-        self.half_angle = 45 # degrees for finding car to follow
+        self.half_angle = half_angle
 
         # Default IDM parameters
-        assert(target_speed>0, 'negative target speed')
+        assert target_speed>0, 'negative target speed'
         self.s_max = target_speed
         self.a_max = np.array([3.]) # nominal acceleration
         self.tau = 0.5 # desired time headway
         self.b_pref = 2.5 # preferred deceleration
         self.d_min = 1 #minimum spacing
-        super().__init__()
 
+        # for np.remainder nan warnings
+        np.seterr(invalid='ignore')
+
+    # BaseAlgorithm abstract methods
+    def _setup_model(self): 
+        return None
+    def learn(self, *args, **kwargs):
+        return self
+    
     def predict(self, observation:np.ndarray, 
-        *args, **kwargs) -> Tuple[np.ndarray,Optional[np.ndarray]]:
+        *args, **kwargs) -> Tuple[np.ndarray, None]:
         """
-        Generate action, state from observation
+        Predict action, state from observation
 
         (But actually generate next action from underlying environment state)
         
@@ -98,41 +112,59 @@ class IDMRulePolicy(BaseAlgorithm):
 
         Returns
             action (np.ndarray): action for controlled agent to take
-            state (np.ndarray): the index of the chosen vehicle for IDM
+            state (None): None (hidden state for a recurrent policy)
+        """
+        return self.forward(observation, *args, **kwargs), None
+
+    def forward(self, *args, **kwargs) -> np.ndarray:
+        """
+        Generate action from underlying environment
+
+        Returns
+            action (np.ndarray): action for controlled agent to take
         """
-        import pdb
-        pdb.set_trace()
         agent = self._env._agent
         full_state = self._env._env.projected_state.numpy() #(nv, 5)
         ego_state = full_state[agent] # (5,)
         s = ego_state[2]
         xy = full_state[:,0:2] # (nv, 2)
         v = full_state[:,2:3] # (nv, 1)
-        psi = full_state[:,2:3] # (nv, 1)
-
+        psi = full_state[:,3:4] # (nv, 1)
 
         d, r, i = self.get_ego_dr(agent, xy, v, psi)
 
         # propagate environment forward at constant velocity
-        if self.t_future > 0:
-            xy2 = xy + self.t_future * v * np.vstack((np.cos(psi[:,0]), np.sin(psi[:,0])))
-            d2, r2, i2 = self.get_ego_dr(agent, xy2, v, psi)
-            
-            # choose closer vehicle (now vs imagined)
-            if d2 < d:
-                d, r, i = d2, r2, i2   
+        for t in self.t_future:
+            if t > 0:
+                xy2 = xy + t * v * np.vstack((np.cos(psi[:,0]), np.sin(psi[:,0]))).T
+                d2, r2, i2 = self.get_ego_dr(agent, xy2, v, psi)
+                
+                # choose closer vehicle (now vs imagined)
+                if d2 < d:
+                    d, r, i = d2, r2, i2
         
+        # Update environment interaction graph with i
+        if i:
+            self._env._env._graph._neighbor_dict={agent:[i]}
+
         if d == np.inf:
-            d_des = 0
+            d_des = self.d_min
         else:
             d_des = self.d_min + self.tau * s + s * r / (2* (self.a_max*self.b_pref)**0.5 )
-        action = self.a_max*(1 - (s/self.s_max)**4 - (d_des/d)**2)
+            d_des = max(d_des, self.d_min)
 
-        assert(action.shape==(1,))
-        return action, i
+        assert (d_des>= self.d_min)
+        action = self.a_max*(1 - (s/self.s_max)**4 - (d_des/d)**2)
+        
+        # normalize action to range if env is a NormalizedActionSpace
+        if isinstance(self._env, NormalizedActionSpace):
+            action = self._env._normalize(action)
+
+        assert action.shape==(1,)
+        return action
     
     def get_ego_dr(self, agent:int, xy: np.ndarray, 
-        v: np.ndarray, psi: np.ndarray) -> Tuple[Optional[np.ndarray], float, float]:
+        v: np.ndarray, psi: np.ndarray) -> Tuple[float, float, Optional[int]]:
         """
         Return distance and relative speed of closest car within half angle from heading
         
@@ -145,14 +177,14 @@ class IDMRulePolicy(BaseAlgorithm):
         Returns:
             d (float): distance to closest vehicle in cone
             r (float): relative speed between the two vehicles
-            i (Union[None,np.ndarray]): (1,) array of closest vehicle index, or None 
+            i (Optional[int]): index of closest vehicle, or None 
         """
         nv, nxy = xy.shape
         nv2, nvel = v.shape
         nv3, npsi = psi.shape
-        assert(nv==nv2==nv3)
-        assert(nxy==2)
-        assert(nvel==npsi==1)
+        assert nv==nv2==nv3
+        assert nxy==2
+        assert nvel==npsi==1
 
         dxys = xy - xy[agent] # (nv, 2)
         ds = np.linalg.norm(dxys,axis=1) # (nv,)
@@ -160,7 +192,7 @@ class IDMRulePolicy(BaseAlgorithm):
         dl = (dxys*np.hstack((-np.sin(psi), np.cos(psi)))).sum(-1) # (nv, )
         alpha = to_circle(np.arctan2(dl, df))
 
-        val_idx = np.arange(nv)[(np.abs(alpha) < self.half_angle*np.pi/180 & np.arange(nv) != agent)]
+        val_idx = np.arange(nv)[(np.abs(alpha) < self.half_angle*np.pi/180) & (np.arange(nv) != agent)]
         
         if len(val_idx)==0:
             i = None
@@ -168,7 +200,7 @@ class IDMRulePolicy(BaseAlgorithm):
             r = float('inf')
         else:
             idx = np.argmin(ds[val_idx]) # closest car which meets requirements
-            i = val_idx[idx]
+            i = int(val_idx[idx])
             d = ds[i]
             r = v[i,0]-v[agent,0]