Skip to content
Snippets Groups Projects
agbot_stanley_teacher.py 7.99 KiB
import gurobipy as gp
import numpy as np

from dtree_teacher_base import DTreeGurobiTeacherBase


K_P = 0.1 #the controller gain is the strength of action a controller will take at a particular point
FORWARD_VEL = 1.0  # m/s car velocity
CYCLE_TIME = 0.05  # s

ANG_VEL_LIM = 0.5  # rad/s^2

# Limits on unconstrained variables to avoid overflow and angle normalization
ANG_LIM = np.pi / 2  # radian
CTE_LIM = 0.76 * 0.5  # meter


# Bounds for intermediate variables.
# These are not necessary but can improve efficiency
K_CTE_V_LIM = K_P * CTE_LIM / FORWARD_VEL
ATAN_K_CTE_V_LIM = np.arctan(K_CTE_V_LIM)
RAW_ANG_ERR_LIM = (ANG_LIM + ATAN_K_CTE_V_LIM)
RAW_ANG_VEL_LIM = RAW_ANG_ERR_LIM / CYCLE_TIME

NEW_K_CTE_V_LIM = K_CTE_V_LIM + FORWARD_VEL * CYCLE_TIME * 1.0
NEW_ATAN_K_CTE_V_LIM = np.arctan(NEW_K_CTE_V_LIM)
NEW_RAW_ANG_ERR_LIM = ANG_LIM + FORWARD_VEL * CYCLE_TIME


def sensor(state):
    """ Assuming the lane to track is aligned with x-axis (i.e., y==0 and yaw==0)
        Hence, heading = 0-yaw = -yaw and distance = 0-y = -y."""
    x, y, yaw = state
    # TODO approximation instead of perfect perception
    # TODO is this where we do AX+b
    prcv_heading = -yaw
    prcv_distance = -y
    return prcv_distance,prcv_heading


def controller(prcv_ctd, prcv_heading):
    error = prcv_heading + np.arctan(K_P*prcv_ctd/FORWARD_VEL)
    
    # Calculate controller output
    ang_vel = error / CYCLE_TIME
    if ang_vel > ANG_VEL_LIM:
        ang_vel = ANG_VEL_LIM
    elif ang_vel < -ANG_VEL_LIM:
        ang_vel = -ANG_VEL_LIM

    # Return actuator values
    return (ang_vel,)


def dynamics(old_state, ang_vel):
    """ This dynamics for state variables x, y, yaw

        x[n+1] = x[n] + v*cos(yaw)*CYCLE_TIME
        y[n+1] = y[n] + v*sin(yaw)*CYCLE_TIME
        yaw[n+1] = yaw[n] + ang_vel*CYCLE_TIME
    """
    old_x, old_y, old_yaw = old_state
    new_x = old_x + FORWARD_VEL * np.cos(old_yaw) * CYCLE_TIME
    new_y = old_y + FORWARD_VEL * np.sin(old_yaw) * CYCLE_TIME
    new_yaw = old_yaw + ang_vel*CYCLE_TIME
    return new_x, new_y, new_yaw


class DTreeAgBotStanleyGurobiTeacher(DTreeGurobiTeacherBase):
    # Ideal perception as a linear transform from state to ground truth percept
    PERC_GT = np.array([[0., -1., 0.],
                        [0., 0., -1.]], float)

    def __init__(self, name="agbot_stanley", norm_ord=2, ultimate_bound: float=0.0) -> None:
        assert ultimate_bound >= 0.0
        self._ultimate_bound = ultimate_bound

        super().__init__(name=name,
                         state_dim=3, perc_dim=2, ctrl_dim=1, norm_ord=norm_ord)

    def is_safe_state(self, ex) -> bool:
        def v(ctd, psi) -> float:
            return abs(psi + np.arctan( (K_P*ctd) / FORWARD_VEL))

        def spec(x, y, theta, ctd, psi) -> bool:
            m_star = sensor
            f = dynamics
            g = controller
            v_old = v(*m_star((x, y, theta)))
            v_new = v(*m_star( f((x, y, theta),*g(ctd, psi)) ))
            return v_new <= max(v_old, self._ultimate_bound)
        return spec(*ex)

    def _add_system(self) -> None:
        # Bounds on all domains
        self._old_state.lb = (-np.inf, -CTE_LIM, -ANG_LIM)
        self._old_state.ub = (np.inf, CTE_LIM, ANG_LIM)
        self._percept.lb = (-CTE_LIM, -ANG_LIM)
        self._percept.ub = (CTE_LIM, ANG_LIM)
        self._control.lb = (-ANG_VEL_LIM,)
        self._control.ub = (ANG_VEL_LIM,)

        # Variable Aliases
        m = self._gp_model
        old_x, old_y, old_yaw = self._old_state.tolist()
        new_x, new_y, new_yaw = self._new_state.tolist()
        ctd, psi = self._percept.tolist()
        ang_vel, = self._control.tolist()

        # Controller
        K_ctd_V = m.addVar(name="K*dhat/Vf", lb=-K_CTE_V_LIM, ub=K_CTE_V_LIM)
        m.addConstr(K_ctd_V == K_P*ctd/FORWARD_VEL)
        atan_K_ctd_V = m.addVar(name="atan(K*dhat/Vf)",
                                lb=-ATAN_K_CTE_V_LIM, ub=ATAN_K_CTE_V_LIM)
        m.addGenConstrTan(xvar=atan_K_ctd_V, yvar=K_ctd_V)

        # Clip angular velocity
        raw_ang_vel = m.addVar(name="(φhat+atan(K*dhat/V))/T",
                               lb=-RAW_ANG_VEL_LIM, ub=RAW_ANG_VEL_LIM)
        m.addConstr(raw_ang_vel == (psi + atan_K_ctd_V) / CYCLE_TIME)

        ang_vel = m.addVar(name="ω", lb=-ANG_VEL_LIM, ub=ANG_VEL_LIM)
        m.addGenConstrPWL(name="clip", xvar=raw_ang_vel, yvar=ang_vel,
                          xpts=[-RAW_ANG_VEL_LIM, -ANG_VEL_LIM, ANG_VEL_LIM, RAW_ANG_VEL_LIM],
                          ypts=[-ANG_VEL_LIM, -ANG_VEL_LIM, ANG_VEL_LIM, ANG_VEL_LIM])

        # Dynamics
        cos_yaw = m.addVar(name="cosθ", **self.TRIGVAR)
        m.addGenConstrCos(old_yaw, cos_yaw)
        sin_yaw = m.addVar(name="sinθ", **self.TRIGVAR)
        m.addGenConstrSin(old_yaw, sin_yaw)

        new_x = old_x + FORWARD_VEL * CYCLE_TIME * cos_yaw
        new_y = m.addVar(name="y'", **self.FREEVAR)
        m.addConstr(new_y == old_y + FORWARD_VEL * CYCLE_TIME * sin_yaw)
        new_yaw = m.addVar(name="θ'", **self.FREEVAR)
        m.addConstr(new_yaw == old_yaw + ang_vel * CYCLE_TIME)
        m.update()
    
    def _add_unsafe(self) -> None:
        assert self.PERC_GT.shape == (self.perc_dim, self.state_dim)
        # Variable Aliases
        m = self._gp_model

        # Add nonincreasing constraints
        old_truth = m.addMVar(shape=(self.perc_dim,), name="m(x)", **self.FREEVAR)
        m.addConstr(old_truth == self.PERC_GT @ self._old_state)

        old_cte, old_psi = old_truth.tolist()
        K_old_cte_V = m.addVar(name="K*d/Vf", lb=-K_CTE_V_LIM, ub=K_CTE_V_LIM)
        m.addConstr(K_old_cte_V == K_P * old_cte / FORWARD_VEL)
        atan_K_old_cte_V = m.addVar(name="atan(K*d/Vf)", lb=-ATAN_K_CTE_V_LIM, ub=ATAN_K_CTE_V_LIM)
        m.addGenConstrTan(xvar=atan_K_old_cte_V, yvar=K_old_cte_V)
        old_err = m.addVar(name="ψ+atan(K*d/Vf)", **self.FREEVAR)
        m.addConstr(old_err == old_psi + atan_K_old_cte_V)

        old_lya_val = m.addVar(name="V(m(x))", **self.NNEGVAR)
        m.addConstr(old_lya_val == gp.abs_(old_err))

        # TODO What is the Lyapunov function?
        # Add ultimate bound
        if self._ultimate_bound > 0:
            bound_new_lya_val = m.addVar(name=f"max(|m(x)|,{self._ultimate_bound})",
                                         **self.NNEGVAR)
            m.addConstr(bound_new_lya_val == gp.max_(old_lya_val, constant=self._ultimate_bound))
        else:
            bound_new_lya_val = old_lya_val

        new_truth = m.addMVar(shape=(self.perc_dim,), name="m(x')", **self.FREEVAR)
        m.addConstr(new_truth == self.PERC_GT @ self._new_state)

        new_cte, new_psi = new_truth.tolist()
        K_new_cte_V = m.addVar(name="K*d'/Vf", lb=-K_CTE_V_LIM, ub=K_CTE_V_LIM)
        m.addConstr(K_new_cte_V == K_P * new_cte / FORWARD_VEL)
        atan_K_new_cte_V = m.addVar(name="atan(K*d'/Vf)", lb=-ATAN_K_CTE_V_LIM, ub=ATAN_K_CTE_V_LIM)
        m.addGenConstrTan(xvar=atan_K_new_cte_V, yvar=K_new_cte_V)
        new_err = m.addVar(name="ψ'+atan(K*d'/Vf)", **self.FREEVAR)
        m.addConstr(new_err == new_psi + atan_K_new_cte_V)
        new_lya_val = m.addVar(name="V(m(x'))", **self.NNEGVAR)
        m.addConstr(new_lya_val == gp.abs_(new_err))

        # Negation of stability spec
        m.addConstr(new_lya_val >= bound_new_lya_val, name="Increasing Error")  # Tracking error is increasing
        m.update()


    def _add_objective(self) -> None:
        # Variable Aliases
        m = self._gp_model
        x = self._old_state
        z = self._percept
        z_diff = self._percept_diff

        # Add objective
        norm_var = m.addVar(name="|z-m(x)|", **self.NNEGVAR)
        m.addConstr(z_diff == z - self.PERC_GT @ x)
        m.addConstr(norm_var == gp.norm(z_diff, float(self._norm_ord)))

        m.setObjective(norm_var, gp.GRB.MINIMIZE)
        m.update()


def test_gem_stanley_gurobi_teacher():
    teacher = DTreeAgBotStanleyGurobiTeacher(ultimate_bound=.0)
    teacher.set_old_state_bound(
        lb=(-np.inf, 0.5, 0.0625),
        ub=(np.inf, 1.0, 0.25)
    )
    teacher.dump_system_encoding()


if __name__ == "__main__":
    test_gem_stanley_gurobi_teacher()