InteractionImitation/src/baselines/rule_policies.py

from stable_baselines3.common.base_class import BaseAlgorithm
from intersim.envs.intersimple import Intersimple, NormalizedActionSpace
from typing import Tuple, Optional, List
import numpy as np

class PControllerPolicy(BaseAlgorithm):


    def __init__(self, env):
        """
        Initialize policy with pointer to environment it will run on
        """
        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

        (But actually generate next action from underlying environment state)
        
        Args:
            observation (np.ndarray): instantaneous observation from environment

        Returns
            action (np.ndarray): action for controlled agent to take
            state (np.ndarray): hidden state for use in next prediction (null)
        """
        agent = self._env._agent
        ego_state = self._env._env.projected_state[agent].numpy() # (5,) tensor
        
        
        # relative_state = np.delete(self._env._env.relative_state[agent].numpy(), agent, axis=0) #(nv-1, 6) tensor
        
        # calculate front and left distances from ego

        # calculate relative speed in direction of position difference vector
        
        # calculate angle alpha and distance d of vehicle i from ego heading
        
        # attn[i] ~= exp( -(alpha[i])^2 - .01 * d[i] - .1 * vrel[i]


        # Proportional controller
        # action = (self.target_v  - self.attn_weight * attn.sum()) - ego_state[2]
        action = self.target_v  - ego_state[2]
        return action, None

class IDMRulePolicy(BaseAlgorithm):
    """
    IDMRulePolicy returns action predictions based on an IDM policy.

    The front car is chosen as the closer of:
        - closest car within a 45 degree half angle cone of the ego's heading
        - ''' after propagating the environment forward by `t_future' seconds with 
            current headings and velocities
    
    """

    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 (List[float]): list of future time at which to compare closest vehicle
            half_angle (float): half angle to look inside for closest vehicle
        """

        self._env = env
        self.t_future = t_future
        self.half_angle = half_angle

        # Default IDM parameters
        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

        # 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, None]:
        """
        Predict action, state from observation

        (But actually generate next action from underlying environment state)
        
        Args:
            observation (np.ndarray): instantaneous observation from environment

        Returns
            action (np.ndarray): action for controlled agent to take
            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
        """
        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[:,3:4] # (nv, 1)

        d, r, i = self.get_ego_dr(agent, xy, v, psi)

        # propagate environment forward at constant velocity
        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 = self.d_min
        else:
            d_des = self.d_min + self.tau * s + s * r / (2* (self.a_max*self.b_pref)**0.5 )
            d_des = max(d_des, self.d_min)

        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[float, float, Optional[int]]:
        """
        Return distance and relative speed of closest car within half angle from heading
        
        Args:
            agent (int): agent index
            xy (np.ndarray): (nv, 2) x and y positions
            v (np.ndarray): (nv, 1) velocity
            psi (np.ndarray): (nv, 1) heading angle

        Returns:
            d (float): distance to closest vehicle in cone
            r (float): relative speed between the two vehicles
            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

        dxys = xy - xy[agent] # (nv, 2)
        ds = np.linalg.norm(dxys,axis=1) # (nv,)
        df = (dxys*np.hstack((np.cos(psi),np.sin(psi)))).sum(-1) # (nv, )
        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)]
        
        if len(val_idx)==0:
            i = None
            d = float('inf')
            r = float('inf')
        else:
            idx = np.argmin(ds[val_idx]) # closest car which meets requirements
            i = int(val_idx[idx])
            d = ds[i]
            r = v[i,0]-v[agent,0] 
            
        return d, r, i
    
def to_circle(x: np.ndarray) -> np.ndarray:
    """
    Casts x (in rad) to [-pi, pi)

    Args:
        x (np.ndarray): (*) input angle (radians)

    Returns:
        y (np.ndarray): (*) x cast to [-pi, pi)
    """
    y = np.remainder(x + np.pi, 2*np.pi) - np.pi
    return y