add the nearby goal

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

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

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
 
 

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):