lss train: height map difference penalty

Author: rlnav

Diff for: scripts/train_lss

@@ -14,7 +14,7 @@ from monoforce.datasets.data import TravData, explore_data
 from monoforce.config import Config
 from monoforce.utils import read_yaml, write_to_yaml, normalize
 from monoforce.datasets import seq_paths, sim_seq_paths
-from monoforce.losses import RMSE
+from monoforce.losses import RMSE, total_variation
 from tqdm import tqdm
 from torch.utils.tensorboard import SummaryWriter
 from datetime import datetime
@@ -222,7 +222,11 @@ class Trainer:
                 loss_geom = self.lidar_hm_loss(height_pred_geom, height_lidar, weights_lidar)
                 loss_rigid = self.traj_hm_loss(height_pred_rigid, height_traj, weights_traj)
 
-                loss = loss_geom + 100.*loss_rigid
+                # add height difference loss
+                # loss_hdiff = height_pred_diff.abs().mean()
+                loss_hdiff = total_variation(height_pred_diff)
+
+                loss = loss_geom + 100.*loss_rigid + 0.1*loss_hdiff
                 if self.map_consistency and len(height_pred_rigid) > 1:
                     loss_map = self.map_consistency_loss(height_pred_rigid, map_pose)
                     loss += 0.1*loss_map
@@ -239,6 +243,7 @@ class Trainer:
                 counter += 1
                 self.writer.add_scalar(f"{'train' if train else 'val'}/iter_loss_geom", loss_geom, counter)
                 self.writer.add_scalar(f"{'train' if train else 'val'}/iter_loss_rigid", loss_rigid, counter)
+                self.writer.add_scalar(f"{'train' if train else 'val'}/iter_loss_hdiff", loss_hdiff, counter)
                 self.writer.add_scalar(f"{'train' if train else 'val'}/iter_loss_map", loss_map, counter)
                 self.writer.add_scalar(f"{'train' if train else 'val'}/iter_loss", loss, counter)