ballsy

dryvr_plus_plus/example/example_agent/ball_agent.py

+# Example agent.
+from typing import Tuple, List
+
+import numpy as np
+from scipy.integrate import ode
+
+from dryvr_plus_plus.scene_verifier.agents.base_agent import BaseAgent
+from dryvr_plus_plus.scene_verifier.map.lane_map import LaneMap
+from dryvr_plus_plus.scene_verifier.code_parser.pythonparser import EmptyAst
+
+
+class BallAgent(BaseAgent):
+    '''Dynamics of a frictionless billiard ball
+        on a 2D-plane'''
+    def __init__(self, id, code = None, file_name = None):
+        super().__init__(id, code, file_name)
+
+    @staticmethod
+    def dynamic(t, state, u):
+        x, y, vx, vy = state
+        x_dot = vx
+        y_dot = vy
+        vx_dot = 0
+        vy_dot = 0
+        return [x_dot, y_dot, vx_dot, vy_dot]
+
+    def action_handler(self, mode: List[str], state, lane_map:LaneMap)->Tuple[float, float]:
+        x,y,theta,v = state
+        vehicle_mode = mode[0]
+        vehicle_lane = mode[1]
+        vehicle_pos = np.array([x,y])
+        a = 0
+        if vehicle_mode == "Normal":
+            d = -lane_map.get_lateral_distance(vehicle_lane, vehicle_pos)
+        elif vehicle_mode == "SwitchLeft":
+            d = -lane_map.get_lateral_distance(vehicle_lane, vehicle_pos) + 3
+        elif vehicle_mode == "SwitchRight":
+            d = -lane_map.get_lateral_distance(vehicle_lane, vehicle_pos) - 3
+        elif vehicle_mode == "Brake":
+            d = -lane_map.get_lateral_distance(vehicle_lane, vehicle_pos)
+            a = -1
+        elif vehicle_mode == "Accel":
+            d = -lane_map.get_lateral_distance(vehicle_lane, vehicle_pos)
+            a = 1
+        elif vehicle_mode == 'Stop':
+            d = -lane_map.get_lateral_distance(vehicle_lane, vehicle_pos)
+            a = 0
+        psi = lane_map.get_lane_heading(vehicle_lane, vehicle_pos)-theta
+        steering = psi + np.arctan2(0.45*d, v)
+        steering = np.clip(steering, -0.61, 0.61)
+        return steering, a
+
+    def TC_simulate(self, mode: List[str], initialCondition, time_bound, lane_map:LaneMap=None)->np.ndarray:
+        time_step = 0.05
+        time_bound = float(time_bound)
+        number_points = int(np.ceil(time_bound/time_step))
+        t = [round(i*time_step,10) for i in range(0,number_points)]
+
+        init = initialCondition
+        trace = [[0]+init]
+        for i in range(len(t)):
+            steering, a = self.action_handler(mode, init, lane_map)
+            r = ode(self.dynamic)
+            r.set_initial_value(init).set_f_params([steering, a])
+            res:np.ndarray = r.integrate(r.t + time_step)
+            init = res.flatten().tolist()
+            if init[3] < 0:
+                init[3] = 0
+            trace.append([t[i] + time_step] + init)
+
+        return np.array(trace)