limiting max forces by m*g for stability reasons, using trajectory discount factor gamma=0.9 by default

Author: rlnav

Diff for: monoforce/scripts/train.sh

@@ -2,16 +2,16 @@
 
 MODEL=lss  # lss, lidarbev, bevfusion
 ROBOT=marv
-DEBUG=False
+DEBUG=True
 VIS=False
-BSZ=32  # 32, 32, 8
+BSZ=4  # 32, 32, 8
 WEIGHTS=$HOME/workspaces/traversability_ws/src/monoforce/monoforce/config/weights/${MODEL}/val.pth
 
 source $HOME/workspaces/traversability_ws/devel/setup.bash
 ./train.py --bsz $BSZ --nepochs 1000 --lr 1e-3 --weight_decay 1e-7 \
            --debug $DEBUG --vis $VIS \
            --geom_weight 1.0 --terrain_weight 2.0 --phys_weight 1.0 \
-           --traj_sim_time 4.0 \
+           --traj_sim_time 5.0 \
            --robot $ROBOT \
            --model $MODEL \
            --pretrained_model_path ${WEIGHTS}

Diff for: monoforce/src/monoforce/datasets/rough.py

@@ -238,6 +238,7 @@ def get_traj(self, i, T_horizon=None):
             # repeat the last pose to fill the trajectory
             poses = np.concatenate([poses, np.tile(poses[-1:], (n_frames - len(poses), 1, 1))], axis=0)
             stamps = np.concatenate([stamps, stamps[-1] + np.arange(1, n_frames - len(stamps) + 1) * dt], axis=0)
+            assert len(poses) == n_frames, f'Poses and stamps have different lengths {len(poses)} != {n_frames}'
         # truncate the trajectory
         poses = poses[:n_frames]
         stamps = stamps[:n_frames]

Diff for: monoforce/src/monoforce/losses.py

@@ -99,7 +99,7 @@ def hm_loss(height_pred, height_gt, weights=None, h_max=None):
     return loss
 
 
-def physics_loss(states_pred, states_gt, pred_ts, gt_ts, gamma=0., rotation_loss=False):
+def physics_loss(states_pred, states_gt, pred_ts, gt_ts, gamma=0.9, rotation_loss=False):
     """
     Compute the physics loss between predicted and ground truth states.
     :param states_pred: predicted states [N x T1 x 3]

Diff for: monoforce/src/monoforce/models/dphysics.py

@@ -170,16 +170,17 @@ def forward_kinematics(self, t, state):
         assert in_contact.shape == (B, n_pts, 1)
 
         # reaction at the contact points as spring-damper forces
+        m, g = self.dphys_cfg.robot_mass, self.dphys_cfg.gravity
         xd_points_n = (xd_points * n).sum(dim=2).unsqueeze(2)  # normal velocity
         assert xd_points_n.shape == (B, n_pts, 1)
         F_spring = -torch.mul((stiffness_points * dh_points + damping_points * xd_points_n), n)  # F_s = -k * dh - b * v_n
         F_spring = torch.mul(F_spring, in_contact) / n_pts  # apply forces only at the contact points
+        F_spring = torch.clamp(F_spring, min=-m*g, max=m*g)
         assert F_spring.shape == (B, n_pts, 3)
 
         # static and dynamic friction forces: https://en.wikipedia.org/wiki/Friction
         thrust_dir = normalized(R[..., 0])  # direction of the thrust
         N = torch.norm(F_spring, dim=2)  # normal force magnitude at the contact points
-        m, g = self.dphys_cfg.robot_mass, self.dphys_cfg.gravity
         track_vels = vw_to_track_vels(v=controls_t[:, 0], w=controls_t[:, 1],
                                       robot_size=self.dphys_cfg.robot_size, n_tracks=len(self.dphys_cfg.driving_parts))
         assert track_vels.shape == (B, len(self.dphys_cfg.driving_parts))
@@ -195,6 +196,7 @@ def forward_kinematics(self, t, state):
         vels_diff_tau = vels_diff - vels_diff_n * n  # tangential velocity difference
         F_friction = N.unsqueeze(2) * vels_diff_tau  # F_f = mu * N * v
         # F_friction = N.unsqueeze(2) * vels_diff
+        F_friction = torch.clamp(F_friction, min=-m*g, max=m*g)
         assert F_friction.shape == (B, n_pts, 3)
 
         # rigid body rotation: M = sum(r_i x F_i)