fix in evaluation: consistent with training losses, updated docs

rlnav · rlnav · commit 4fc88825ccef · 2024-12-17T11:47:15.000+01:00
diff --git a/monoforce/docs/DATA.md b/monoforce/docs/DATA.md
@@ -50,11 +50,7 @@ Colored point cloud            |            Front-facing camera            |
 ![](./imgs/rgb_cloud.png) | ![](./imgs/camera_fisheye_front.png) | ![](./imgs/camera_up.png)
 
 To explore the data, please follow the tutorial
-at [../examples/explore_data_rgb_rigid_terrain.ipynb](../examples/explore_data_rgb_rigid_terrain.ipynb)
-or run the module:
-```commandline
-python -m monoforce.datasets.rough
-```
+at [../examples/explore_data_rgb_rigid_terrain.ipynb](../examples/explore_data_rgb_rigid_terrain.ipynb).
 
 ### Data Sample
 
diff --git a/monoforce/docs/DPHYS.md b/monoforce/docs/DPHYS.md
@@ -6,9 +6,9 @@ Run the differentiable physics simulation with a robot model and the provided da
 python scripts/robot_control
 ```
 
-![](./imgs/hm_learning.gif)
-
 Terrain properties optimization from the ground-truth trajectories followed by the robot:
 ```commandline
 python scripts/fit_terrain
 ```
+
+![](./imgs/hm_learning.gif)
diff --git a/monoforce/docs/INSTALL.md b/monoforce/docs/INSTALL.md
@@ -26,13 +26,12 @@ catkin build
 
 ### Model Weights
 
-The pretrained weights for the LSS terrain encoder can be downloaded from:
-- RobinGas: [lss_robingas_husky.pt](https://drive.google.com/file/d/1h1VieiIdGZB1Ml3QdIlh8ZJA67sJej4m/view?usp=sharing),
-            [lss_robingas_tradr.pt](https://drive.google.com/file/d/1jpsgXN-44Bbu9hfAWd5Z3te1DWp3s8cX/view?usp=sharing),
-            [lss_robingas_husky_oru.pt](https://drive.google.com/file/d/12v6EAvaw0LqdINYFyHYr0t5mlZn-VN6c/view?usp=sharing),
-- RELLIS-3D: [lss_rellis3d.pt](https://drive.google.com/file/d/1kK75mUxHn-4GadU4k8-c43hA9t3bZxw1/view?usp=sharing).
+The pretrained weights for the terrain encoder are available at:
+[]().
 
-Once downloaded put the weights to `monoforce/config/weights/lss` folder.
+Once downloaded, please, put the weights to 
+`monoforce/config/weights/${MODEL_NAME}` folder,
+where `${MODEL_NAME}` is the name is one of the models: `lss, lidarbev, bevfusion`.
 
 ## Docker
 
diff --git a/monoforce/docs/TERRAIN_ENCODER.md b/monoforce/docs/TERRAIN_ENCODER.md
@@ -19,11 +19,3 @@ To train the LSS model, please run:
 cd scripts/
 python train
 ```
-
-### Weights
-
-The pretrained weights for the LSS terrain encoder can be downloaded from
-[Google Drive](https://drive.google.com/drive/folders/1n0KoUEQ0FIznf-qMKshk0C956CJenIhd?usp=sharing).
-
-Once downloaded put the weights to
-`monoforce/monoforce/config/weights/lss` folder.
diff --git a/monoforce/examples/bevfusion.ipynb b/monoforce/examples/bevfusion.ipynb
@@ -3,22 +3,20 @@
   {
    "metadata": {
     "ExecuteTime": {
-     "end_time": "2024-11-25T12:39:10.061779Z",
-     "start_time": "2024-11-25T12:39:08.313307Z"
+     "end_time": "2024-12-17T10:38:44.003902Z",
+     "start_time": "2024-12-17T10:38:42.182608Z"
     }
    },
    "cell_type": "code",
    "source": [
     "import sys\n",
     "sys.path.append('../src/')\n",
     "import torch\n",
-    "import numpy as np\n",
     "from torch.utils.data import DataLoader\n",
     "from monoforce.models.terrain_encoder.bevfusion import BEVFusion\n",
     "from monoforce.utils import read_yaml, position\n",
     "from monoforce.datasets.rough import ROUGH, rough_seq_paths\n",
-    "from monoforce.dphys_config import DPhysConfig\n",
-    "from monoforce.transformations import transform_cloud"
+    "from monoforce.dphys_config import DPhysConfig"
    ],
    "id": "b7a7dacb0d156951",
    "outputs": [
@@ -45,8 +43,8 @@
   {
    "metadata": {
     "ExecuteTime": {
-     "end_time": "2024-11-25T12:39:10.087675Z",
-     "start_time": "2024-11-25T12:39:10.062753Z"
+     "end_time": "2024-12-17T10:38:44.149683Z",
+     "start_time": "2024-12-17T10:38:44.004928Z"
     }
    },
    "cell_type": "code",
@@ -62,14 +60,23 @@
     "                points)"
    ],
    "id": "dacf9fbbbc81fec7",
-   "outputs": [],
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/home/ruslan/miniconda3/envs/mf/lib/python3.9/site-packages/torch/functional.py:504: UserWarning: torch.meshgrid: in an upcoming release, it will be required to pass the indexing argument. (Triggered internally at /opt/conda/conda-bld/pytorch_1666643016022/work/aten/src/ATen/native/TensorShape.cpp:3190.)\n",
+      "  return _VF.meshgrid(tensors, **kwargs)  # type: ignore[attr-defined]\n"
+     ]
+    }
+   ],
    "execution_count": 2
   },
   {
    "metadata": {
     "ExecuteTime": {
-     "end_time": "2024-11-25T12:39:10.105485Z",
-     "start_time": "2024-11-25T12:39:10.088532Z"
+     "end_time": "2024-12-17T10:38:44.193210Z",
+     "start_time": "2024-12-17T10:38:44.150482Z"
     }
    },
    "cell_type": "code",
@@ -86,7 +93,7 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "Dataset length: 335\n"
+      "Dataset length: 349\n"
      ]
     }
    ],
@@ -95,8 +102,8 @@
   {
    "metadata": {
     "ExecuteTime": {
-     "end_time": "2024-11-25T12:39:11.121161Z",
-     "start_time": "2024-11-25T12:39:10.106441Z"
+     "end_time": "2024-12-17T10:38:45.200218Z",
+     "start_time": "2024-12-17T10:38:44.194312Z"
     }
    },
    "cell_type": "code",
@@ -121,8 +128,8 @@
   {
    "metadata": {
     "ExecuteTime": {
-     "end_time": "2024-11-25T12:39:14.271963Z",
-     "start_time": "2024-11-25T12:39:11.122171Z"
+     "end_time": "2024-12-17T10:38:48.136932Z",
+     "start_time": "2024-12-17T10:38:45.200914Z"
     }
    },
    "cell_type": "code",
@@ -142,8 +149,8 @@
   {
    "metadata": {
     "ExecuteTime": {
-     "end_time": "2024-11-25T12:39:15.973490Z",
-     "start_time": "2024-11-25T12:39:14.273103Z"
+     "end_time": "2024-12-17T10:38:49.838240Z",
+     "start_time": "2024-12-17T10:38:48.137920Z"
     }
    },
    "cell_type": "code",
diff --git a/monoforce/examples/trajectory_shooting_with_torch_diff_physics.ipynb b/monoforce/examples/trajectory_shooting_with_torch_diff_physics.ipynb
diff --git a/monoforce/scripts/eval.py b/monoforce/scripts/eval.py
@@ -16,6 +16,7 @@
 from monoforce.datasets.rough import ROUGH, rough_seq_paths
 from monoforce.models.terrain_encoder.utils import ego_to_cam, get_only_in_img_mask, denormalize_img
 from monoforce.utils import read_yaml, write_to_csv, append_to_csv
+from monoforce.losses import physics_loss, hm_loss
 import matplotlib as mpl
 
 
@@ -80,45 +81,6 @@ def __init__(self,
         # self.ds = Fusion(path=self.path, lss_cfg=self.lss_config, dphys_cfg=self.dphys_cfg, is_train=False)
         self.loader = torch.utils.data.DataLoader(self.ds, batch_size=1, shuffle=False)
 
-    def hm_loss(self, height_pred, height_gt, weights=None):
-        assert height_pred.shape == height_gt.shape, 'Height prediction and ground truth must have the same shape'
-        if weights is None:
-            weights = torch.ones_like(height_gt)
-        assert weights.shape == height_gt.shape, 'Weights and height ground truth must have the same shape'
-
-        # remove nan values
-        mask_valid = ~torch.isnan(height_gt)
-        height_gt = height_gt[mask_valid]
-        height_pred = height_pred[mask_valid]
-        weights = weights[mask_valid]
-
-        # compute weighted loss
-        loss = torch.nn.functional.mse_loss(height_pred * weights, height_gt * weights, reduction='mean')
-        assert not torch.isnan(loss), 'Terrain Loss is nan'
-
-        return loss
-
-    def physics_loss(self, states_pred, states_gt, pred_ts, gt_ts):
-        # unpack the states
-        X, Xd, R, Omega = states_gt
-        X_pred, Xd_pred, R_pred, Omega_pred = states_pred
-
-        # find the closest timesteps in the trajectory to the ground truth timesteps
-        ts_ids = torch.argmin(torch.abs(pred_ts.unsqueeze(1) - gt_ts.unsqueeze(2)), dim=2)
-
-        # get the predicted states at the closest timesteps to the ground truth timesteps
-        batch_size = X.shape[0]
-        X_pred_gt_ts = X_pred[torch.arange(batch_size).unsqueeze(1), ts_ids]
-
-        # remove nan values
-        mask_valid = ~torch.isnan(X_pred_gt_ts)
-        X_pred_gt_ts = X_pred_gt_ts[mask_valid]
-        X = X[mask_valid]
-        loss = torch.nn.functional.mse_loss(X_pred_gt_ts, X)
-        assert not torch.isnan(loss), 'Physics Loss is nan'
-
-        return loss
-
     def run(self, vis=False, save=False):
         if save:
             # create output folder
@@ -168,12 +130,12 @@ def run(self, vis=False, save=False):
                 # friction_pred = torch.ones_like(terrain_pred)
 
                 # evaluation losses
-                terrain_loss = self.hm_loss(height_pred=terrain_pred[0, 0], height_gt=hm_terrain[0, 0])
+                loss_terrain = hm_loss(height_pred=terrain_pred[0, 0], height_gt=hm_terrain[0, 0], weights=hm_terrain[0, 1])
                 states_gt = [Xs, Xds, Rs, Omegas]
                 state0 = tuple([s[:, 0] for s in states_gt])
                 states_pred, _ = self.dphysics(z_grid=terrain_pred.squeeze(1), state=state0,
                                                controls=controls, friction=friction_pred.squeeze(1))
-                physics_loss = self.physics_loss(states_pred, states_gt, pred_ts=control_ts, gt_ts=traj_ts)
+                loss_physics = physics_loss(states_pred=states_pred, states_gt=states_gt, pred_ts=control_ts, gt_ts=traj_ts)
 
                 # visualizations
                 terrain_pred = terrain_pred[0, 0].cpu()
@@ -187,7 +149,7 @@ def run(self, vis=False, save=False):
                 # hm_points = hm_points[:, terrain_mask]
 
                 plt.clf()
-                plt.suptitle(f'Terrain Loss: {terrain_loss.item():.4f}, Physics Loss: {physics_loss.item():.4f}')
+                plt.suptitle(f'Terrain Loss: {loss_terrain.item():.4f}, Physics Loss: {loss_physics.item():.4f}')
                 for imgi, img in enumerate(imgs[0]):
                     cam_pts = ego_to_cam(hm_points, rots[0, imgi], trans[0, imgi], intrins[0, imgi])
                     mask = get_only_in_img_mask(cam_pts, H, W)
@@ -251,7 +213,7 @@ def run(self, vis=False, save=False):
                 if save:
                     plt.savefig(f'{self.output_folder}/{i:04d}.png')
                     append_to_csv(f'{self.output_folder}/losses.csv',
-                                  f'{i:04d}.png, {terrain_loss.item():.4f},{physics_loss.item():.4f}\n')
+                                  f'{i:04d}.png, {loss_terrain.item():.4f},{loss_physics.item():.4f}\n')
 
             plt.close(fig)
 
diff --git a/monoforce/scripts/eval.sh b/monoforce/scripts/eval.sh
@@ -6,7 +6,7 @@ ROBOT=marv
 
 source $HOME/workspaces/traversability_ws/devel/setup.bash
 # loop through data sequences
-for SEQ_I in {0..27};
+for SEQ_I in {0..18};
 do
     echo "Evaluating sequence ${SEQ_I}"
     ./eval.py --model_path ${WEIGHTS} --robot ${ROBOT} --seq_i ${SEQ_I} --vis #--save
diff --git a/monoforce/scripts/robot_control.py b/monoforce/scripts/robot_control.py
@@ -25,7 +25,7 @@ def motion():
 
     # control inputs: linear velocity and angular velocity, v in m/s, w in rad/s
     controls = torch.stack([
-        torch.tensor([[1.0, 0.0]] * int(dphys_cfg.traj_sim_time / dphys_cfg.dt)),  # [v] m/s, [w] rad/s for each time step
+        torch.tensor([[1.0, 1.0]] * int(dphys_cfg.traj_sim_time / dphys_cfg.dt)),  # [v] m/s, [w] rad/s for each time step
     ]).to(device)
     B, N_ts, _ = controls.shape
     assert controls.shape == (B, N_ts, 2), f'controls shape: {controls.shape}'
@@ -44,14 +44,12 @@ def motion():
     # heightmap defining the terrain
     x_grid, y_grid = dphys_cfg.x_grid, dphys_cfg.y_grid
     # z_grid = torch.sin(x_grid) * torch.cos(y_grid)
-    # z_grid = torch.exp(-(x_grid - 2) ** 2 / 4) * torch.exp(-(y_grid - 0) ** 2 / 2)
-    z_grid = torch.zeros_like(x_grid)
-    # z_grid = torch.from_numpy(np.load('./gen/terrain_optimization/z_grid.npy'))
-    z_grid[80:81, 0:100] = 1.0  # add a wall
+    z_grid = torch.exp(-(x_grid - 2) ** 2 / 4) * torch.exp(-(y_grid - 0) ** 2 / 2)
+    # z_grid = torch.zeros_like(x_grid)
+    # z_grid[80:81, 0:100] = 1.0  # add a wall
 
     x_grid, y_grid, z_grid = x_grid.to(device), y_grid.to(device), z_grid.to(device)
     friction = dphys_cfg.friction
-    # friction = torch.from_numpy(np.load('./gen/terrain_optimization/friction.npy'))
     friction = friction.to(device)
 
     # repeat the heightmap for each rigid body
@@ -201,7 +199,7 @@ def shoot_multiple():
     # simulate the rigid body dynamics
     with torch.no_grad():
         t0 = time()
-        states, forces = dphysics(z_grid=z_grid, controls=controls, state=state0)
+        states, forces = dphysics(z_grid=z_grid, controls=controls, state=state0, vis=False)
         t1 = time()
         Xs, Xds, Rs, Omegas = states
         print(f'Simulation took {(t1-t0):.3f} [sec] on device: {device}')
diff --git a/monoforce/src/monoforce/datasets/rough.py b/monoforce/src/monoforce/datasets/rough.py
@@ -33,18 +33,9 @@
         os.path.join(data_dir, 'ROUGH/24-08-14-monoforce-long_drive'),
         os.path.join(data_dir, 'ROUGH/marv_2024-09-26-13-46-51'),
         os.path.join(data_dir, 'ROUGH/marv_2024-09-26-13-54-43'),
-        # os.path.join(data_dir, 'ROUGH/marv_2024-10-05-12-34-53'),
-        # os.path.join(data_dir, 'ROUGH/marv_2024-10-05-13-01-40'),
-        # os.path.join(data_dir, 'ROUGH/marv_2024-10-05-13-17-08'),
-        # os.path.join(data_dir, 'ROUGH/marv_2024-10-05-13-29-39'),
-        # os.path.join(data_dir, 'ROUGH/marv_2024-10-05-13-43-21'),
-        # os.path.join(data_dir, 'ROUGH/marv_2024-10-05-13-57-57'),
-        # os.path.join(data_dir, 'ROUGH/marv_2024-10-05-14-12-29'),
-        # os.path.join(data_dir, 'ROUGH/marv_2024-10-05-14-22-10'),
-        # os.path.join(data_dir, 'ROUGH/marv_2024-10-05-14-28-15'),
         os.path.join(data_dir, 'ROUGH/marv_2024-10-31-15-16-42'),
         os.path.join(data_dir, 'ROUGH/marv_2024-10-31-15-26-47'),
-        os.path.join(data_dir, 'ROUGH/marv_2024-10-31-15-35-05'),  # some problem with numpy arrays instead of torch tensors
+        os.path.join(data_dir, 'ROUGH/marv_2024-10-31-15-35-05'),
         os.path.join(data_dir, 'ROUGH/marv_2024-10-31-15-52-07'),
         os.path.join(data_dir, 'ROUGH/marv_2024-10-31-15-56-33'),
 
