CoolFace
Datasetpublic

piyushkumar-tiwari/RRC_Experiments

sourceHugging Faceupdated 3mo agoView on Hugging Face
0likes5.6kdownloads
mpc_controller.py124 linesDownload Raw Back to root
1import casadi as ca2import numpy as np3 4class MPCController:5    def __init__(self, horizon=10, dt=0.1, max_vel=0.25, min_safe_dist=0.20):6        self.N = horizon7        self.dt = dt8        self.max_vel = max_vel9        self.min_safe_dist = min_safe_dist  # Hard minimum distance radius (meters)10        11        # State: [x, y, z, yaw], Input: [vx, vz, yaw_rate]12        self.nx = 413        self.nu = 314        15        self.init_solver()16 17    def init_solver(self):18        # Symbols19        x = ca.MX.sym('x', self.nx)20        u = ca.MX.sym('u', self.nu)21        22        # Parameters: Current State, Reference, and Obstacles23        # We simplify obstacles to the 'closest' point for the QP, 24        # or a list of points for the NLP.25        self.p = ca.MX.sym('p', self.nx + self.nx + 3) # [x0, ref, closest_obs]26 27        # Prediction horizon variables28        X = ca.MX.sym('X', self.nx, self.N + 1)29        U = ca.MX.sym('U', self.nu, self.N)30 31        obj = 032        g = []33        34        # Costs35        Q = ca.diag([500, 100, 40, 40])36        #R = ca.diag([5, 5, 5])37        R = ca.diag([18, 22, 20])  # Increased vx cost (18 instead of 15) to prefer yaw over reversing; reduced yaw cost (2) to encourage rotation38        obs_weight = 20.0  # Increased obstacle penalty to make avoidance more aggressive39 40        g.append(X[:, 0] - self.p[0:4]) # Initial condition constraint41 42        for k in range(self.N):43            # Objective: State error + Control effort44            st_err = X[:, k] - self.p[4:8]45            obj += ca.mtimes([st_err.T, Q, st_err])46            obj += ca.mtimes([U[:, k].T, R, U[:, k]])47            48            # Obstacle avoidance: Direction-aware repulsive potential49            # Penalizes obstacles ahead more to encourage yaw-based lateral avoidance50            obs_rel = self.p[8:11] - X[0:3, k]  # obstacle relative position51            dist_to_obs = ca.norm_2(obs_rel)52            53            # Compute directional penalty: obstacles ahead (positive Xb) are worse54            # Forward component (Xb) weighted heavily, side components (Yb) less so55            # forward_component = obs_rel[0]  # Xb (forward in body frame)56            # directional_cost = ca.if_else(57            #     forward_component > 0,58            #     forward_component / (dist_to_obs + 0.1),  # normalize by distance59            #     060            # )61            62            influence_radius = 1.5  # only repel within this range63            repulsion = ca.if_else(64              dist_to_obs < influence_radius,65              obs_weight * ((1/dist_to_obs - 1/influence_radius)**2),66              067              )68            obj += repulsion69            70            # Hard minimum distance constraint71            g.append(dist_to_obs - self.min_safe_dist)72 73            # Dynamics (Forward Euler)74            # x_next = x + dt * f(x, u)75            # In your model, yaw affects x and y translation76            x_next = ca.vertcat(77                X[0, k] + self.dt * (U[0, k] * ca.cos(X[3, k])),78                X[1, k] + self.dt * (U[0, k] * ca.sin(X[3, k])),79                X[2, k] + self.dt * U[1, k],80                X[3, k] + self.dt * U[2, k]81            )82            g.append(X[:, k+1] - x_next)83 84        # Solver Setup85        opt_variables = ca.reshape(X, -1, 1)86        opt_variables = ca.vertcat(opt_variables, ca.reshape(U, -1, 1))87        88        nlp_prob = {'f': obj, 'x': opt_variables, 'g': ca.vertcat(*g), 'p': self.p}89        opts = {'ipopt.print_level': 0, 'print_time': 0}90        self.solver = ca.nlpsol('solver', 'ipopt', nlp_prob, opts)91 92    def control(self, x0, ref, closest_obs):93        # Flatten initial guess94        x_init = np.zeros((self.nx * (self.N + 1) + self.nu * self.N, 1))95        96        # Parameters: current state, ref state, and one representative obstacle point97        params = np.concatenate([x0, ref, closest_obs])98        99        # Constraints (bounds)100        # Total constraints: 4 (initial) + N (distance) + 4*N (dynamics) = 4 + 5*N101        lbg = np.zeros(4 + 5 * self.N)102        ubg = np.zeros(4 + 5 * self.N)103        104        # Set bounds for distance constraints105        for k in range(self.N):106            dist_idx = 4 + k * 5  # distance constraint index for step k107            lbg[dist_idx] = 0108            ubg[dist_idx] = np.inf109        110        # Variable bounds (Velocities)111        lbx = -np.inf * np.ones(self.nx * (self.N + 1) + self.nu * self.N)112        ubx = np.inf * np.ones(self.nx * (self.N + 1) + self.nu * self.N)113        114        # U bounds115        u_start = self.nx * (self.N + 1)116        for i in range(self.N):117            lbx[u_start + i*3 : u_start + i*3 + 2] = -self.max_vel118            ubx[u_start + i*3 : u_start + i*3 + 2] = self.max_vel119            lbx[u_start + i*3 + 2] = -np.deg2rad(45) # Yaw rate120            ubx[u_start + i*3 + 2] = np.deg2rad(45)121 122        sol = self.solver(x0=x_init, p=params, lbx=lbx, ubx=ubx, lbg=lbg, ubg=ubg)123        u_out = ca.reshape(sol['x'][u_start:], self.nu, self.N)124        return np.array(u_out[:, 0]).flatten()
piyushkumar-tiwari/RRC_Experiments · CoolFace