loss(deltaflow): add deltaflow loss.

Author: Kin-Zhang

README.md

@@ -1,4 +1,4 @@
-DeltaFlow
+DeltaFlow: An Efficient Multi-frame Scene Flow Estimation Method
 ---
 
 [![arXiv](https://img.shields.io/badge/arXiv-2508.17054-b31b1b?logo=arxiv&logoColor=white)](https://arxiv.org/abs/2508.17054)
@@ -14,9 +14,8 @@ Note (2025/09/18): We got accepted by NeurIPS 2025 and it's **spotlighted**! 
 - [x] 2025/08/24: Updating train data augmentation as illustrated in the DeltaFlow paper.
 - [x] 2025/08/25: Updating paper preprint link.
 - [x] 2025/09/05: Merged the latest commit from OpenSceneFlow codebase to DeltaFlow for afterward unified merged.
-- [x] 2025/09/25: DeltaFlow Model python file and config file.
+- [x] 2025/09/25: DeltaFlow Model file, config file and loss function. Update quick training example.
 - [ ] pre-trained weights upload.
-- [ ] DeltaFlow Loss fn.
 - [ ] Merged into [OpenSceneFlow](https://github.com/KTH-RPL/OpenSceneFlow)
 
 ## Quick Run
@@ -34,7 +33,7 @@ unzip demo-data-v2.zip -d /home/kin/data/av2/h5py # to your data path
 
 3. Run the training with the following command (modify the data path accordingly):
 ```bash
-python train.py model=deltaflow batch_size=4 num_frames=5 voxel_size="[0.15,0.15,0.15]" point_cloud_range="[-38.4,-38.4,-3,38.4,38.4,3]" optimizer.lr=2e-4 train_data=${demo_train_data_path} val_data=${demo_val_data_path}
+python train.py model=deltaflow loss_fn=deltaflowLoss batch_size=4 num_frames=5 voxel_size="[0.15,0.15,0.15]" point_cloud_range="[-38.4,-38.4,-3,38.4,38.4,3]" optimizer.lr=2e-4 train_data=${demo_train_data_path} val_data=${demo_val_data_path}
 ```
 ### Evaluation
 

src/lossfuncs/supervise.py

@@ -16,7 +16,49 @@
 import os, sys
 BASE_DIR = os.path.abspath(os.path.join( os.path.dirname( __file__ ), '../..' ))
 sys.path.append(BASE_DIR)
+from src.utils.av2_eval import CATEGORY_TO_INDEX, BUCKETED_METACATAGORIES
 
+# check: https://arxiv.org/abs/2508.17054
+def deltaflowLoss(res_dict):
+    pred = res_dict['est_flow']
+    gt = res_dict['gt_flow']
+    classes = res_dict['gt_classes']
+    instances = res_dict['gt_instance']
+    
+    reassign_meta = torch.zeros_like(classes, dtype=torch.int, device=classes.device)
+    for i, cats in enumerate(BUCKETED_METACATAGORIES):
+        selected_classes_ids = [CATEGORY_TO_INDEX[cat] for cat in BUCKETED_METACATAGORIES[cats]]
+        reassign_meta[torch.isin(classes, torch.tensor(selected_classes_ids, device=classes.device))] = i
+
+    pts_loss = torch.linalg.vector_norm(pred - gt, dim=-1)
+    speed = torch.linalg.vector_norm(gt, dim=-1) / 0.1
+    
+    weight_loss = deflowLoss(res_dict)['loss']
+
+    classes_loss = 0.0
+    weight = [0.1, 1.0, 2.0, 2.5, 1.5] # BACKGROUND, CAR, PEDESTRIAN, WHEELED, OTHER
+    for class_id in range(len(BUCKETED_METACATAGORIES)):
+        mask = reassign_meta == class_id
+        for loss_ in [0.1 * pts_loss[(speed < 0.4) & mask].mean(), 
+                      0.4 * pts_loss[(speed >= 0.4) & (speed <= 1.0) & mask].mean(), 
+                      0.5 * pts_loss[(speed > 1.0) & mask].mean()]:
+            classes_loss += torch.nan_to_num(loss_, nan=0.0) * weight[class_id]
+
+    instance_loss, cnt = 0.0, 0
+    if instances is not None:
+        for instance_id in torch.unique(instances):
+            mask = instances == instance_id
+            reassign_meta_instance = reassign_meta[mask]
+            class_id = torch.mode(reassign_meta_instance, 0).values.item()
+            loss_ = pts_loss[mask].mean()
+            if speed[mask].mean() <= 0.4:
+                continue
+            instance_loss += (loss_ * torch.exp(loss_) * weight[class_id])
+            cnt += 1
+        instance_loss /= (cnt if cnt > 0 else 1)
+    return {'loss': weight_loss + classes_loss + instance_loss}
+
+# check: https://arxiv.org/abs/2401.16122
 def deflowLoss(res_dict):
     pred = res_dict['est_flow']
     gt = res_dict['gt_flow']

src/trainer.py

@@ -132,7 +132,8 @@ def training_step(self, batch, batch_idx):
 
             dict2loss = {'est_flow': est_flow[batch_id], 
                         'gt_flow': None if 'flow' not in batch else batch['flow'][batch_id][pc0_valid_from_pc2res] - pose_flow_, 
-                        'gt_classes': None if 'flow_category_indices' not in batch else batch['flow_category_indices'][batch_id][pc0_valid_from_pc2res]}
+                        'gt_classes': None if 'flow_category_indices' not in batch else batch['flow_category_indices'][batch_id][pc0_valid_from_pc2res],
+                        'gt_instance': None if 'flow_instance_id' not in batch else batch['flow_instance_id'][batch_id][pc0_valid_from_pc2res],}
 
             if 'pc0_dynamic' in batch:
                 dict2loss['pc0_labels'] = batch['pc0_dynamic'][batch_id][pc0_valid_from_pc2res]