diff --git a/YOPO/config/traj_opt.yaml b/YOPO/config/traj_opt.yaml index 63ddfd1..4728396 100644 --- a/YOPO/config/traj_opt.yaml +++ b/YOPO/config/traj_opt.yaml @@ -41,8 +41,7 @@ az_mean_unit: 0.0 ax_std_unit: 0.4 ay_std_unit: 0.4 az_std_unit: 0.3 -goal_length: 10.0 -goal_pitch_std: 10.0 +goal_pitch_std: 10.0 # clip goal_length to: 2 * radio_range; 10% probability [0, 2 * radio_range] goal_yaw_std: 20.0 # extend the map to avoid trajectory exceeding bounds (especially in z-axis) and treating the sky outside the map as obstacle diff --git a/YOPO/loss/guidance_loss.py b/YOPO/loss/guidance_loss.py index 5fa08cc..3094670 100644 --- a/YOPO/loss/guidance_loss.py +++ b/YOPO/loss/guidance_loss.py @@ -1,15 +1,10 @@ -import os import torch.nn as nn import torch as th -from ruamel.yaml import YAML class GuidanceLoss(nn.Module): def __init__(self): super(GuidanceLoss, self).__init__() - base_dir = os.path.dirname(os.path.abspath(__file__)) - cfg = YAML().load(open(os.path.join(base_dir, "../config/traj_opt.yaml"), 'r')) - self.max_similarity = 2.0 * cfg['radio_range'] def forward(self, Df, Dp, goal): """ @@ -19,14 +14,19 @@ class GuidanceLoss(nn.Module): goal: (batch_size, 3) Returns: similarity: (batch_size) → guidance loss + + GuidanceLoss: Projection length of the trajectory onto the goal direction: higher cosine similarity and longer trajectory are preferred """ cur_pos = Df[:, :, 0] end_pos = Dp[:, :, 0] traj_dir = end_pos - cur_pos # [B, 3] goal_dir = goal - cur_pos # [B, 3] - goal_dir = goal_dir / (goal_dir.norm(dim=1, keepdim=True) + 1e-8) # [B, 3] - similarity = self.max_similarity - th.sum(traj_dir * goal_dir, dim=1) # [B] + goal_length = goal_dir.norm(dim=1) - return similarity + goal_dir_norm = goal_dir / (goal_length.unsqueeze(1) + 1e-8) # [B, 3] + similarity = th.sum(traj_dir * goal_dir_norm, dim=1) # [B] + + similarity_loss = th.abs(goal_length - similarity) + return similarity_loss diff --git a/YOPO/policy/state_transform.py b/YOPO/policy/state_transform.py index 8392449..79a89c6 100644 --- a/YOPO/policy/state_transform.py +++ b/YOPO/policy/state_transform.py @@ -10,6 +10,7 @@ class StateTransform: base_dir = os.path.dirname(os.path.abspath(__file__)) cfg = YAML().load(open(os.path.join(base_dir, "../config/traj_opt.yaml"), 'r')) self.lattice_primitive = LatticePrimitive.get_instance(cfg) + self.goal_length = 2.0 * cfg['radio_range'] def pred_to_endstate(self, endstate_pred: torch.Tensor) -> torch.Tensor: """ @@ -112,7 +113,12 @@ class StateTransform: def normalize_obs(self, vel_acc_goal): vel_acc_goal[:, 0:3] = vel_acc_goal[:, 0:3] / self.lattice_primitive.vel_max vel_acc_goal[:, 3:6] = vel_acc_goal[:, 3:6] / self.lattice_primitive.acc_max - vel_acc_goal[:, 6:9] = vel_acc_goal[:, 6:9] / (vel_acc_goal[:, 6:9].norm(p=2, dim=1, keepdim=True) + 1e-8) + vel_acc_goal[:, 6:9] = vel_acc_goal[:, 6:9] / self.goal_length + + # Clamp the goal direction to unit length + goal_norm_length = vel_acc_goal[:, 6:9].norm(dim=1, keepdim=True) + scaling = goal_norm_length.clamp(max=1.0) / (goal_norm_length + 1e-8) + vel_acc_goal[:, 6:9] = vel_acc_goal[:, 6:9] * scaling return vel_acc_goal diff --git a/YOPO/policy/yopo_dataset.py b/YOPO/policy/yopo_dataset.py index b430dd7..1c67a29 100644 --- a/YOPO/policy/yopo_dataset.py +++ b/YOPO/policy/yopo_dataset.py @@ -26,7 +26,7 @@ class YOPODataset(Dataset): self.v_std = np.array([cfg["vx_std_unit"], cfg["vy_std_unit"], cfg["vz_std_unit"]]) self.a_mean = np.array([cfg["ax_mean_unit"], cfg["ay_mean_unit"], cfg["az_mean_unit"]]) self.a_std = np.array([cfg["ax_std_unit"], cfg["ay_std_unit"], cfg["az_std_unit"]]) - self.goal_length = cfg["goal_length"] + self.goal_length = 2.0 * cfg['radio_range'] self.goal_pitch_std = cfg["goal_pitch_std"] self.goal_yaw_std = cfg["goal_yaw_std"] if mode == 'train': self.print_data() @@ -127,6 +127,10 @@ class YOPODataset(Dataset): goal_pitch_angle, goal_yaw_angle = np.radians(goal_pitch_angle), np.radians(goal_yaw_angle) goal_w_dir = np.array([np.cos(goal_yaw_angle) * np.cos(goal_pitch_angle), np.sin(goal_yaw_angle) * np.cos(goal_pitch_angle), np.sin(goal_pitch_angle)]) + # 10% probability to generate a nearby goal (× goal_length is actual length) + random_near = np.random.rand() + if random_near < 0.1: + goal_w_dir = random_near * 10 * goal_w_dir return self.goal_length * goal_w_dir def print_data(self):