Flow-Planner代码阅读(2)：数据加载

一、数据加载 dataset

代码在flow_planner/data/dataset/nuplan.py中，类名NuPlanDataset()，通过__getitem__()函数随机获取某一帧数据。

def__getitem__(self,idx)->NuPlanDataSample:data=np.load(os.path.join(self.data_dir,self.data_list[idx]))#加载数据ego_agent_past=torch.from_numpy(data['ego_agent_past'])# 历史时刻信息 [21,14]ego_current_state=torch.from_numpy(data['ego_current_state'])#当前时刻状态[16]ego_agent_future=torch.from_numpy(data['ego_agent_future']).to(torch.float32)#未来时刻信息[80,3]neighbor_agents_past=torch.from_numpy(data['neighbor_agents_past'][:self._past_neighbor_num])#agent历史时刻信息[32,21,11]neighbor_agents_future=torch.from_numpy(data['neighbor_agents_future'][:self._predicted_neighbor_num])#flow_planner里面设置为0neighbor_future_observed=torch.from_numpy(data['neighbor_agents_future'])#agent未来时刻信息，[32,80,3]lanes=torch.from_numpy(data['lanes'])#车道线信息 [70,20,12]lanes_speed_limit=torch.from_numpy(data['lanes_speed_limit'])#车道限速[70,1]lanes_has_speed_limit=torch.from_numpy(data['lanes_has_speed_limit'])#车道是否限速[70,1]route_lanes=torch.from_numpy(data['route_lanes'])#route所在车道信息 [25,20,12]route_lanes_speed_limit=torch.from_numpy(data['route_lanes_speed_limit'])# route所在车道限速[25,1]route_lanes_has_speed_limit=torch.from_numpy(data['route_lanes_has_speed_limit'])#route所在车道是否有限速[25,1static_objects=torch.from_numpy(data['static_objects'])#静态目标信息,[5,10]

二、数据增强

2.1、自车状态扰动

defaugment(self,data):# Only aug current stateself._device=data.ego_current.device self.move_to(self._device)ego_current_state=data.ego_current.clone()#自车状态,[batch_size,16],分别是B=ego_current_state.shape[0]#batchsize#以30%概率做自车扰动，自车速度小于2m/s的不做自车扰动aug_flag=(torch.rand(B)<=self._augment_prob).bool().to(self._device)&~(abs(ego_current_state[:,4])<2.0)#自车每个状态的扰动范围,在[low,high]范围内随机取值random_tensor=torch.rand(B,len(self._low)).to(self._device)scaled_random_tensor=self._low+(self._high-self._low)*random_tensor new_state=torch.zeros((B,9),dtype=torch.float32).to(self._device)new_state[:,3:]=ego_current_state[:,4:10]# x, y, h is 0 because of ego-centric, update vx, vy, ax, ay, steering angle, yaw ratenew_state=new_state+scaled_random_tensor#添加扰动值new_state[:,3]=torch.max(new_state[:,3],torch.tensor(0.0,device=new_state.device))#自车纵向速度不能为负，禁止倒车new_state[:,-1]=torch.clip(new_state[:,-1],-0.85,0.85)#最后一维是 yaw rate#替换扰动后的自车状态ego_current_state[:,:2]=new_state[:,:2]ego_current_state[:,2]=torch.cos(new_state[:,2])ego_current_state[:,3]=torch.sin(new_state[:,2])ego_current_state[:,4:8]=new_state[:,3:7]ego_current_state[:,8:10]=new_state[:,-2:]# steering angle, yaw rate# 重新计算steer和yaw ratecur_velocity=ego_current_state[:,4]yaw_rate=ego_current_state[:,9]steering_angle=torch.zeros_like(cur_velocity)new_yaw_rate=torch.zeros_like(yaw_rate)mask=torch.abs(cur_velocity)<0.2not_mask=~mask steering_angle[not_mask]=torch.atan(yaw_rate[not_mask]*self._wheel_base/torch.abs(cur_velocity[not_mask]))steering_angle[not_mask]=torch.clamp(steering_angle[not_mask],-2/3*np.pi,2/3*np.pi)new_yaw_rate[not_mask]=yaw_rate[not_mask]ego_current_state[:,8]=steering_angle ego_current_state[:,9]=new_yaw_ratereturnaug_flag,ego_current_state

2.2、自车未来轨迹插值

由于自车状态做了扰动，自车当前的位置、heading等信息和自车未来下一个时刻的状态就可能会不符合运动学约束，refine_future_trajectory()函数以扰动后的自车状态为起始点，对未来20个时刻的自车状态进行插值，剩余60个不变，这样既可以在短时间内符合运动学约束，在长时间内保证模型学习的真值趋势不变。

defrefine_future_trajectory(self,aug_current_state,ego_future):""" refine future trajectory with quintic spline interpolation Args: aug_current_state: (B, 16) current state of the ego vehicle after augmentation ego_future: (B, 80, 3) future trajectory of the ego vehicle Returns: ego_future: refined future trajectory of the ego vehicle with velocity and acceleration """#利用当前增强后的自车状态 + 第 P 个未来点作为终端约束，用五次样条插值细化未来轨迹的前 P 个点，保证位置、速度、加速度连续P=self.num_refine#需要插值的20个点dt=self.time_interval#0.1s，100ms，对应10hzT=self.refine_horizon B=aug_current_state.shape[0]M_t=self.t_matrix.unsqueeze(0).expand(B,-1,-1)#时间基矩阵（t⁰~t⁵）A=self.coeff_matrix.unsqueeze(0).expand(B,-1,-1)#五次多项式边界条件 → 系数矩阵device=ego_future.device# 解析当前状态（起点约束） 包括: [x, y, heading, velocity, acceleration, yaw_rate]x0,y0,theta0,v0,a0,omega0=(aug_current_state[:,0],aug_current_state[:,1],torch.atan2((ego_future[:,int(P/2),1]-aug_current_state[:,1]),(ego_future[:,int(P/2),0]-aug_current_state[:,0])),torch.norm(aug_current_state[:,4:6],dim=-1),torch.norm(aug_current_state[:,6:8],dim=-1),aug_current_state[:,9])#解析终点状态（第 P 个点，P=20）xT,yT,thetaT,vT,aT,omegaT=(ego_future[:,P,0],ego_future[:,P,1],ego_future[:,P,2],torch.norm(ego_future[:,P,:2]-ego_future[:,P-1,:2],dim=-1)/dt,torch.norm(ego_future[:,P,:2]-2*ego_future[:,P-1,:2]+ego_future[:,P-2,:2],dim=-1)/dt**2,self.normalize_angle(ego_future[:,P,2]-ego_future[:,P-1,2])/dt)# 构造五次多项式边界条件（x 方向）sx=torch.stack([x0,v0*torch.cos(theta0),a0*torch.cos(theta0)-v0*torch.sin(theta0)*omega0,xT,vT*torch.cos(thetaT),aT*torch.cos(thetaT)-vT*torch.sin(thetaT)*omegaT],dim=-1)#构造五次多项式边界条件（y 方向）sy=torch.stack([y0,v0*torch.sin(theta0),a0*torch.sin(theta0)+v0*torch.cos(theta0)*omega0,yT,vT*torch.sin(thetaT),aT*torch.sin(thetaT)+vT*torch.cos(thetaT)*omegaT],dim=-1)ax=A @ sx[:,:,None].to(torch.float32)ay=A @ sy[:,:,None].to(torch.float32)# 轨迹插值（位置）traj_x=M_t @ ax traj_y=M_t @ ay# 根据速度方向计算航向角traj_heading=torch.cat([torch.atan2(traj_y[:,:1,0]-y0.unsqueeze(-1),traj_x[:,:1,0]-x0.unsqueeze(-1)),torch.atan2(traj_y[:,1:,0]-traj_y[:,:-1,0],traj_x[:,1:,0]-traj_x[:,:-1,0])],dim=1)returntorch.concatenate([torch.cat([traj_x,traj_y,traj_heading[...,None]],axis=-1),ego_future[:,P+1:,:3]],axis=1)