code/submodules/Pplan/Sampling/spline_planner.py

from logging import raiseExceptions
import numpy as np
import torch
import pdb
from ..utils import geometry_utils as GeoUtils
import matplotlib.pyplot as plt
from scipy.interpolate import interp1d
import random

from .forward_sampler import ForwardSampler

STATE_INDEX = [0, 1, 2, 4]

device = "cuda" if torch.cuda.is_available() else "cpu"


def mean_control_effort_coefficients(x0, dx0, xf, dxf):
    """Returns `(c4, c3, c2)` corresponding to `c4 * tf**-4 + c3 * tf**-3 + c2 * tf**-2`."""
    return (12 * (x0 - xf) ** 2, 12 * (dx0 + dxf) * (x0 - xf), 4 * dx0 ** 2 + 4 * dx0 * dxf + 4 * dxf ** 2)


def cubic_spline_coefficients(x0, dx0, xf, dxf, tf):
    return (x0, dx0, -2 * dx0 / tf - dxf / tf - 3 * x0 / tf ** 2 + 3 * xf / tf ** 2,
            dx0 / tf ** 2 + dxf / tf ** 2 + 2 * x0 / tf ** 3 - 2 * xf / tf ** 3)


def compute_interpolating_spline(state_0, state_f, tf):
    dx0, dy0 = state_0[..., 2] * \
               torch.cos(state_0[..., 3]), state_0[..., 2] * \
               torch.sin(state_0[..., 3])
    dxf, dyf = state_f[..., 2] * \
               torch.cos(state_f[..., 3]), state_f[..., 2] * \
               torch.sin(state_f[..., 3])
    tf = tf * torch.ones_like(state_0[..., 0])
    return (
        torch.stack(cubic_spline_coefficients(
            state_0[..., 0], dx0, state_f[..., 0], dxf, tf), -1),
        torch.stack(cubic_spline_coefficients(
            state_0[..., 1], dy0, state_f[..., 1], dyf, tf), -1),
        tf,
    )


def compute_spline_xyvaqrt(x_coefficients, y_coefficients, tf, N=10):
    t = torch.arange(N).unsqueeze(0).to(tf.device) * tf.unsqueeze(-1) / (N - 1)
    tp = t[..., None] ** torch.arange(4).to(tf.device)
    dtp = t[..., None] ** torch.tensor([0, 0, 1, 2]
                                       ).to(tf.device) * torch.arange(4).to(tf.device)
    ddtp = t[..., None] ** torch.tensor([0, 0, 0, 1]).to(
        tf.device) * torch.tensor([0, 0, 2, 6]).to(tf.device)
    x_coefficients = x_coefficients.unsqueeze(-1)
    y_coefficients = y_coefficients.unsqueeze(-1)
    vx = dtp @ x_coefficients
    vy = dtp @ y_coefficients
    v = torch.hypot(vx, vy)
    v_pos = torch.clip(v, min=1e-4)
    ax = ddtp @ x_coefficients
    ay = ddtp @ y_coefficients
    a = (ax * vx + ay * vy) / v_pos
    r = (-ax * vy + ay * vx) / (v_pos ** 2)
    yaw = torch.atan2(vy, vx)
    return torch.cat((
        tp @ x_coefficients,
        tp @ y_coefficients,
        v,
        a,
        yaw,
        r,
        t.unsqueeze(-1),
    ), -1)


def patch_yaw_low_speed(traj):
    idx = traj[...,]


def mean_control_effort_coefficients(x0, dx0, xf, dxf):
    """Returns `(c4, c3, c2)` corresponding to `c4 * tf**-4 + c3 * tf**-3 + c2 * tf**-2`."""
    return (12 * (x0 - xf) ** 2, 12 * (dx0 + dxf) * (x0 - xf), 4 * dx0 ** 2 + 4 * dx0 * dxf + 4 * dxf ** 2)


class SplinePlanner(object):
    def __init__(self, device, dx_grid=None, dy_grid=None, acce_grid=None, dyaw_grid=None, max_steer=0.5, max_rvel=8,
                 acce_bound=[-6, 4], vbound=[-2.0, 30], spline_order=3, N_seg=10, low_speed_threshold=2.0, seed=0):
        self.spline_order = spline_order
        self.device = device
        assert spline_order == 3
        if dx_grid is None:
            # self.dx_grid = torch.tensor([-4., 0, 4.]).to(self.device)
            self.dx_grid = torch.tensor([0.]).to(self.device)
        else:
            self.dx_grid = torch.tensor(dx_grid).to(self.device)
        if dy_grid is None:
            self.dy_grid = torch.tensor([-3., -1.5, 0, 1.5, 3.]).to(self.device)
        else:
            self.dy_grid = torch.tensor(dy_grid).to(self.device)
        self.dy_grid_lane = torch.tensor([-2., 0, 2., ]).to(self.device)
        if acce_grid is None:
            # self.acce_grid = torch.tensor([-1., -0.5, 0., 0.5, 1.]).to(self.device)
            self.acce_grid = torch.tensor([-1., 0., 1.]).to(self.device)
        else:
            self.acce_grid = torch.tensor(acce_grid).to(self.device)
        if dyaw_grid is None:
            self.dyaw_grid = torch.tensor(
                [-np.pi / 6, 0, np.pi / 6]).to(self.device)
        else:
            self.dyaw_grid = torch.tensor(dyaw_grid).to(self.device)
        self.max_steer = max_steer
        self.max_rvel = max_rvel
        self.psi_bound = [-np.pi * 0.75, np.pi * 0.75]
        self.acce_bound = acce_bound
        self.vbound = vbound
        self.N_seg = N_seg
        self.low_speed_threshold = low_speed_threshold
        self.forward_sampler = ForwardSampler(acce_grid=self.acce_grid, dhm_grid=torch.linspace(-0.7, 0.7, 9),
                                              dhf_grid=[-0.4, 0, 0.4], dt=0.1, device=self.device)
        torch.manual_seed(seed)

    def calc_trajectories(self, x0, tf, xf, N=None):
        if N is None:
            N = self.N_seg
        if x0.ndim == 1:
            x0_tile = x0.tile(xf.shape[0], 1)
            xc, yc, tf = compute_interpolating_spline(x0_tile, xf, tf)
        elif x0.ndim == xf.ndim:
            xc, yc, tf = compute_interpolating_spline(x0, xf, tf)
        else:
            raise ValueError("wrong dimension for x0")
        traj = compute_spline_xyvaqrt(xc, yc, tf, N)
        return traj

    def gen_terminals_lane(self, x0, tf, lanes):
        if lanes is None or len(lanes) == 0:
            return self.gen_terminals(x0, tf)

        gs = [self.dx_grid.shape[0], self.dy_grid_lane.shape[0], self.acce_grid.shape[0]]
        dx = self.dx_grid[:, None, None, None].repeat(1, gs[1], gs[2], 1).flatten()
        dy = self.dy_grid_lane[None, :, None, None].repeat(gs[0], 1, gs[2], 1).flatten()
        dv = self.acce_grid[None, None, :, None].repeat(
            gs[0], gs[1], 1, 1).flatten() * tf
        xf = list()
        if x0.ndim == 1:
            for lane in lanes:
                f, p_start = lane
                if isinstance(p_start, np.ndarray):
                    p_start = torch.from_numpy(p_start).to(x0.device)
                elif isinstance(p_start, torch.Tensor):
                    p_start = p_start.to(x0.device)
                offset = x0[:2] - p_start[:2]
                s_offset = offset[0] * \
                           torch.cos(p_start[2]) + offset[1] * torch.sin(p_start[2])
                ds = dx + dv / 2 * tf + x0[2:3] * tf
                ss = ds + s_offset
                xyyaw = torch.from_numpy(f(ss.cpu().numpy())).type(
                    torch.float).to(x0.device)
                xyyaw[..., 1] += dy
                xf.append(
                    torch.cat((xyyaw[:, :2], dv.reshape(-1, 1) + x0[2:3], xyyaw[:, 2:]), -1))
            # adding the end points not fixated on lane
            xf_straight = torch.stack([ds, dy, dv + x0[2], x0[3].tile(ds.shape[0])], -1)
            xf.append(xf_straight)
        elif x0.ndim == 2:
            for lane in lanes:
                f, p_start = lane
                if isinstance(p_start, np.ndarray):
                    p_start = torch.from_numpy(p_start).to(x0.device)
                elif isinstance(p_start, torch.Tensor):
                    p_start = p_start.to(x0.device)
                offset = x0[:, :2] - p_start[None, :2]
                s_offset = offset[:, 0] * torch.cos(p_start[2]) + offset[:, 1] * torch.sin(p_start[2])

                ds = (dx + dv / 2 * tf).unsqueeze(0) + x0[:, 2:3] * tf
                ss = ds + s_offset.unsqueeze(-1)
                xyyaw = torch.from_numpy(f(ss.cpu().numpy())).type(
                    torch.float).to(x0.device)
                xyyaw[..., 1] += dy
                xf.append(torch.cat((xyyaw[..., :2], dv.tile(
                    x0.shape[0], 1).unsqueeze(-1) + x0[:, None, 2:3], xyyaw[..., 2:]), -1))
            # adding the end points not fixated on lane
            xf_straight = torch.stack([ds, dy.tile(x0.shape[0], 1), dv.tile(x0.shape[0], 1) + x0[:, None, 2],
                                       x0[:, None, 3].tile(1, ds.shape[1])], -1)
            xf.append(xf_straight)
        else:
            raise ValueError("x0 must have dimension 1 or 2")
        xf = torch.cat(xf, -2)
        return xf

    def gen_terminals(self, x0, tf):
        gs = [self.dx_grid.shape[0], self.dy_grid.shape[0],
              self.acce_grid.shape[0], self.dyaw_grid.shape[0]]
        dx = self.dx_grid[:, None, None, None].repeat(
            1, gs[1], gs[2], gs[3]).flatten()
        dy = self.dy_grid[None, :, None, None].repeat(
            gs[0], 1, gs[2], gs[3]).flatten()
        dv = tf * self.acce_grid[None, None, :,
                  None].repeat(gs[0], gs[1], 1, gs[3]).flatten()
        dyaw = self.dyaw_grid[None, None, None, :].repeat(
            gs[0], gs[1], gs[2], 1).flatten()
        delta_x = torch.stack([dx, dy, dv, dyaw], -1)

        if x0.ndim == 1:
            xy = torch.cat(
                (delta_x[:, 0:1] + delta_x[:, 2:3] / 2 * tf + x0[2:3] * tf, delta_x[:, 1:2]), -1)
            rotated_delta_xy = GeoUtils.batch_rotate_2D(xy, x0[3])
            refpsi = torch.arctan2(rotated_delta_xy[..., 1], rotated_delta_xy[..., 0])
            rotated_xy = rotated_delta_xy + x0[:2]
            return torch.cat((rotated_xy, delta_x[:, 2:3] + x0[2:3], delta_x[:, 3:] + refpsi.unsqueeze(-1)), -1)
        elif x0.ndim == 2:
            delta_x = torch.tile(delta_x, [x0.shape[0], 1, 1])
            xy = torch.cat(
                (delta_x[:, :, 0:1] + delta_x[:, :, 2:3] / 2 * tf + x0[:, None, 2:3] * tf, delta_x[:, :, 1:2]), -1)
            rotated_delta_xy = GeoUtils.batch_rotate_2D(xy, x0[:, 3:4])
            refpsi = torch.arctan2(rotated_delta_xy[..., 1], rotated_delta_xy[..., 0])
            rotated_xy = rotated_delta_xy + x0[:, None, :2]
            return torch.cat(
                (rotated_xy, delta_x[:, :, 2:3] + x0[:, None, 2:3], delta_x[:, :, 3:] + refpsi.unsqueeze(-1)), -1)
        else:
            raise ValueError("x0 must have dimension 1 or 2")

    def feasible_flag(self, traj, xf):
        diff = traj[..., -1, STATE_INDEX] - xf

        feas_flag = ((traj[..., 2] >= self.vbound[0]) & (traj[..., 2] < self.vbound[1]) &
                     (traj[..., 4] >= self.psi_bound[0]) & (traj[..., 4] < self.psi_bound[1]) &
                     (traj[..., 3] >= self.acce_bound[0]) & (traj[..., 3] <= self.acce_bound[1]) &
                     (torch.abs(traj[..., 5] * traj[..., 2]) <= self.max_rvel) & (
                             torch.clip(torch.abs(traj[..., 2]), min=0.5) * self.max_steer >= torch.abs(
                         traj[..., 5]))).all(1) & (
                            diff.abs() < 5e-3).all(-1)

        return feas_flag

    def gen_trajectories(self, x0, tf, lanes=None, dyn_filter=True, N=None, lane_only=False):
        if N is None:
            N = self.N_seg
        if lanes is not None:
            if isinstance(lanes, torch.Tensor):
                lanes = lanes.cpu().numpy()
            lane_interp = [GeoUtils.interp_lanes(lane) for lane in lanes]
            xf_lane = self.gen_terminals_lane(
                x0, tf, lane_interp)
        else:
            xf_lane = None
        if lane_only:
            assert xf_lane is not None
            xf_set = xf_lane
        else:
            xf_set = self.gen_terminals(x0, tf)
            if xf_lane is not None:
                xf_set = torch.cat((xf_lane, xf_set), 0)
        x0[..., 2] = torch.clip(x0[..., 2], min=1e-3)
        xf_set[..., 2] = torch.clip(xf_set[..., 2], min=1e-3)

        # x, y, v, a, yaw,r, t
        traj = self.calc_trajectories(x0, tf, xf_set, N)
        if dyn_filter:
            feas_flag = self.feasible_flag(traj, xf_set)
            traj = traj[feas_flag]
            xf = xf_set[feas_flag]
        traj = traj[..., 1:, :]  # remove the first time step
        if x0[2] < self.low_speed_threshold:
            # call forward sampler when velocity is low
            extra_traj = self.forward_sampler.sample_trajs(x0.unsqueeze(0), int(tf / self.forward_sampler.dt)).squeeze(
                0)
            f = interp1d(np.arange(1, extra_traj.shape[-2] + 1) * self.forward_sampler.dt, extra_traj.cpu().numpy(),
                         axis=-2)
            extra_traj = torch.from_numpy(f(np.arange(1, N) * tf / N)).to(self.device)
            traj = torch.cat((traj, extra_traj), 0)
        return traj, traj[..., -1, STATE_INDEX]

    @staticmethod
    def get_similarity_flag(x0, x1, thres=[2.0, 0.5, 2.0, np.pi / 12]):
        thres = torch.tensor(thres, device=x0.device)
        diff = x0.unsqueeze(-3) - x1.unsqueeze(-2)
        flag = diff.abs() < thres
        flag = flag.all(-1).any(-1)
        return flag

    def gen_terminals_hardcoded(self, x0_set, tf):
        X0, Y0, v0, psi0 = x0_set[..., 0:1], x0_set[..., 1:2], x0_set[..., 2:3], x0_set[..., 3:]
        xf_set = list()
        # drive straight
        xf_straight = torch.cat((X0 + v0 * tf * torch.cos(psi0), Y0 + v0 * tf * torch.sin(psi0), v0, psi0),
                                -1).unsqueeze(1)
        xf_set.append(xf_straight)
        # hard brake
        decel = torch.clip(-v0 / tf, min=self.acce_bound[0])
        xf_brake = torch.cat((X0 + (v0 + decel * 0.5 * tf) * tf * torch.cos(psi0),
                              Y0 + (v0 + decel * 0.5 * tf) * tf * torch.sin(psi0), v0 + decel * tf, psi0),
                             -1).unsqueeze(1)
        xf_set.append(xf_brake)
        xf_set = torch.cat(xf_set, 1)
        return xf_set

    def gen_trajectory_batch(self, x0_set, tf, lanes=None, dyn_filter=True, N=None, max_children=None):
        if N is None:
            N = self.N_seg
        device = x0_set.device
        xf_set_sample = self.gen_terminals(x0_set, tf)
        importance_score = torch.rand(xf_set_sample.shape[:2], device=device)
        xf_set_hardcoded = self.gen_terminals_hardcoded(x0_set, tf)
        xf_set = torch.cat((xf_set_sample, xf_set_hardcoded), 1)
        importance_score = torch.cat((importance_score, 2 * torch.ones(xf_set_hardcoded.shape[:2], device=device)), 1)
        if lanes is not None:
            lane_interp = [GeoUtils.interp_lanes(lane) for lane in lanes]
            xf_set_lane = self.gen_terminals_lane(x0_set, tf, lane_interp)
            xf_set = torch.cat((xf_set, xf_set_lane), -2)
            importance_score = torch.cat((importance_score, torch.ones(xf_set_lane.shape[:2], device=x0_set.device)), 1)
        x0_set[..., 2] = torch.clip(x0_set[..., 2], min=1e-3)
        xf_set[..., 2] = torch.clip(xf_set[..., 2], min=1e-3)
        num_node = x0_set.shape[0]
        num = xf_set.shape[1]
        x0_tiled = x0_set.repeat_interleave(num, 0)
        xf_tiled = xf_set.reshape(-1, xf_set.shape[-1])
        traj = self.calc_trajectories(x0_tiled, tf, xf_tiled, N)
        if dyn_filter:
            feas_flag = self.feasible_flag(traj, xf_tiled)
        else:
            feas_flag = torch.ones(
                num * num_node, dtype=torch.bool).to(x0_set.device)
        feas_flag = feas_flag.reshape(num_node, num)
        traj = traj.reshape(num_node, num, *traj.shape[1:])
        if (x0_set[:, 2] < self.low_speed_threshold).any():
            extra_traj = self.forward_sampler.sample_trajs(x0_set, int(tf / self.forward_sampler.dt))
            f = interp1d(np.arange(1, extra_traj.shape[-2] + 1) * self.forward_sampler.dt, extra_traj.cpu().numpy(),
                         axis=-2, bounds_error=False, fill_value="extrapolate")
            extra_traj = torch.from_numpy(f(np.arange(0, N) * tf / N)).to(self.device)
            traj = torch.cat((traj, extra_traj), 1)
            extra_importance_score = torch.rand(extra_traj.shape[:2], device=device)
            importance_score = torch.cat((importance_score, extra_importance_score), 1)
            feas_flag = torch.cat((feas_flag, torch.ones(extra_traj.shape[:2], device=device)), 1)

        importance_score = importance_score * feas_flag
        chosen_idx = [torch.where(importance_score[i])[0].tolist() for i in range(num_node)]
        if max_children is not None:
            chosen_idx = [idx if len(idx) <= max_children else torch.topk(importance_score[i], max_children)[1] for
                          i, idx in enumerate(chosen_idx)]
        traj_batch = [traj[i, chosen_idx[i], 1:] for i in range(num_node)]
        return traj_batch

    def gen_trajectory_tree(self, x0, tf, n_layers, dyn_filter=True, N=None):
        if N is None:
            N = self.N_seg
        trajs = list()
        nodes = [x0[None, :]]
        for i in range(n_layers):
            xf = self.gen_terminals(nodes[i], tf)
            x0i = torch.tile(nodes[i], [xf.shape[1], 1])
            xf = xf.reshape(-1, xf.shape[-1])

            traj = self.calc_trajectories(x0i, tf, xf, N)
            if dyn_filter:
                feas_flag = self.feasible_flag(traj, xf)
                traj = traj[feas_flag]
                xf = xf[feas_flag]

            trajs.append(traj)

            nodes.append(xf.reshape(-1, xf.shape[-1]))
        return trajs, nodes[1:]


if __name__ == "__main__":
    planner = SplinePlanner("cuda")
    x0 = torch.tensor([1., 2., 1., 0.]).cuda()
    tf = 5
    traj, xf = planner.gen_trajectories(x0, tf)
    trajs = planner.gen_trajectory_batch(xf, tf)

    # x, y, v, a, yaw,r, t = traj
    msize = 12
    trajs, nodes = planner.gen_trajectory_tree(x0, tf, 2)
    x0 = x0.cpu().numpy()
    traj = traj.cpu().numpy()
    plt.figure(figsize=(20, 10))
    plt.plot(x0[0], x0[1], marker="o", color="b", markersize=msize)
    for node, traj in zip(nodes, trajs):
        node = node.cpu().numpy()
        traj = traj.cpu().numpy()
        x = traj[..., 0]
        y = traj[..., 1]
        plt.plot(x.T, y.T, color="k")
        for p in node:
            plt.plot(p[0], p[1], marker="o", color="b", markersize=msize)
    plt.show()