Compute and plot contact forces between the floor and the foot (#5)

This PR adds a function to compute the contact forces between the floor and the right and left foot. It also plots the value of those forces computed during simulation. This a required feature to build contact detection.

Author: rdesarz

biped_walking_controller/model.py

@@ -113,5 +113,14 @@ def get_joint_id(self, name):
 
         return self.model.joints[jid].idx_q if jid < n_joints else None
 
-    def get_locked_joints_idx(self):
+    @staticmethod
+    def get_locked_joints_idx():
         return range(14, 46)
+
+    @staticmethod
+    def get_rf_link_name():
+        return "leg_right_6_link"
+
+    @staticmethod
+    def get_lf_link_name():
+        return "leg_left_6_link"

biped_walking_controller/plot.py

@@ -185,4 +185,19 @@ def traj2d(arr):  # drop last 2 samples as in your code
     ax2.legend(loc="center left", bbox_to_anchor=(1, 0.5))
     ax2.set_title(f"{title_prefix} — plan view (x–y)")
 
-    plt.show()
+
+def plot_contact_forces(t, force_rf, force_lf, title="Contact force Fx"):
+    t = np.asarray(t).ravel()
+    force_rf = np.asarray(force_rf).ravel()
+    force_lf = np.asarray(force_lf).ravel()
+    assert t.size == force_rf.size == force_lf.size
+
+    plt.figure(figsize=(12, 8), layout="constrained")
+    plt.plot(t, force_rf, label="right foot")
+    plt.plot(t, force_lf, label="left foot")
+    plt.xlabel("t [s]")
+    plt.ylabel("Normal force [N]")
+    plt.title(title)
+    plt.grid(True)
+    plt.legend()
+    plt.tight_layout()

biped_walking_controller/simulation.py

@@ -319,6 +319,27 @@ def _link_index(body_id, link_name):
     return None
 
 
+def _compute_force(
+    contact_pnts: typing.List,
+) -> typing.Tuple[float, typing.List[float]]:
+    mean_x = 0.0
+    mean_y = 0.0
+    mean_z = 0.0
+    total_force = 0.0
+    for cp in contact_pnts:
+        # tuple fields (relevant):
+        # cp[4]=posOnA (world xyz), cp[5]=posOnB, cp[6]=normalOnB (world xyz),
+        # cp[7]=distance, cp[8]=normalForce
+        posB = cp[6]
+        fN = cp[9]  # Newtons
+        mean_x += posB[0]
+        mean_y += posB[1]
+        mean_z += posB[2]
+        total_force += fN
+
+    return total_force, [mean_x, mean_y, mean_z]
+
+
 class Simulator:
     """
     Thin PyBullet wrapper for loading a robot URDF and driving it from
@@ -381,6 +402,15 @@ def __init__(self, dt, path_to_robot_urdf: Path, model, launch_gui=True):
             flags=pb.URDF_MERGE_FIXED_LINKS,
         )
 
+        # Find the joint that matches the right and left foot. It will be used to compute contact forces
+        for j in range(pb.getNumJoints(self.robot_id)):
+            info = pb.getJointInfo(self.robot_id, j)
+            link_name = info[12].decode("ascii")
+            if link_name == model.get_lf_link_name():
+                self.rf_link_id = info[0]
+            elif link_name == model.get_rf_link_name():
+                self.lf_link_id = info[0]
+
         # Build mapping between joints for both velocities and position
         self.pb_to_pin_joint_vel = _build_pb_to_pin_joint_vel_vmap(self.robot_id, model.model)
         self.pb_to_pin_joints_pos = _build_pb_to_pin_joints_map(self.robot_id, model)
@@ -472,67 +502,54 @@ def update_camera_to_follow_pos(self, x: float, y: float, z: float):
             cameraTargetPosition=[x, y, z],
         )
 
-    def draw_contact_forces(self, color=(0, 1, 0), scale=0.002):
+    def get_contact_forces(self) -> typing.Tuple[float, float]:
         """
-        Render an averaged ground-reaction normal as a debug line.
+        Return the net x-axis contact forces for the right and left foot.
 
-        Parameters
-        ----------
-        color : tuple[float, float, float], default (0, 1, 0)
-            RGB color in [0,1].
-        scale : float, default 0.002
-            Visual scale factor converting Newtons to line length.
+        The method queries PyBullet for contacts between the robot and the ground
+        on each foot link, then uses the private helper `_compute_force(...)`
+        to aggregate the force component along the world x-axis.
+
+        Requirements
+        ------------
+        - `self.robot_id` : int
+            PyBullet body id of the robot.
+        - `self.plane_id` : int
+            PyBullet body id of the ground/plane.
+        - `self.rf_link_id` : int
+            Link index of the right foot.
+        - `self.lf_link_id` : int
+            Link index of the left foot.
+        - `_compute_force(contacts) -> tuple[float, Any]`
+            Helper that returns the x-axis force [N] as its first element.
+            It should return 0.0 if `contacts` is empty.
+
+        Returns
+        -------
+        (float, float)
+            `(Fx_right, Fx_left)` in Newtons, world frame.
 
         Notes
         -----
-        - Aggregates all robot–plane contacts, averages contact locations,
-          sums normal forces, and draws a single arrow along the last
-          observed contact normal.
-        - Updates a persistent debug line if it already exists.
+        - Only contacts with `bodyA=self.robot_id`, `bodyB=self.plane_id`,
+          and `linkIndexA` equal to the corresponding foot link are considered.
+        - If a foot has no active contacts, its returned force is expected to be 0.0.
+        - This method assumes `_compute_force` already applies the correct
+          Newton third-law sign convention to yield the force exerted on the robot.
         """
-        # iterate all foot links if you want per-foot colors
-        cps_all = []
-
-        cps_all.extend(pb.getContactPoints(bodyA=self.robot_id, bodyB=self.plane_id))
-
-        mean_x = 0.0
-        mean_y = 0.0
-        mean_z = 0.0
-        total_force = 0.0
-        n_B = [0.0, 0.0, 0.0]
-        for cp in cps_all:
-            # tuple fields (relevant):
-            # cp[4]=posOnA (world xyz), cp[5]=posOnB, cp[6]=normalOnB (world xyz),
-            # cp[7]=distance, cp[8]=normalForce
-            posB = cp[6]
-            n_B = cp[7]
-            fN = cp[9]  # Newtons
-            mean_x += posB[0]
-            mean_y += posB[1]
-            mean_z += posB[2]
-            total_force += fN
-
-        if len(cps_all) is not 0:
-            start = [mean_x / len(cps_all), mean_y / len(cps_all), mean_z / len(cps_all)]
-            end = (
-                start[0] + n_B[0] * total_force * scale,
-                start[1] + n_B[1] * total_force * scale,
-                start[2] + n_B[2] * total_force * scale,
+        rf_force, _ = _compute_force(
+            pb.getContactPoints(
+                bodyA=self.robot_id, bodyB=self.plane_id, linkIndexA=self.rf_link_id
             )
+        )
 
-            # arrow for normal force
-            if self._displayed_lines is None:
-                self._displayed_lines = pb.addUserDebugLine(
-                    start, end, lineColorRGB=color, lineWidth=3
-                )
-            else:
-                pb.addUserDebugLine(
-                    start,
-                    end,
-                    lineColorRGB=color,
-                    lineWidth=3,
-                    replaceItemUniqueId=self._displayed_lines,
-                )
+        lf_force, _ = _compute_force(
+            pb.getContactPoints(
+                bodyA=self.robot_id, bodyB=self.plane_id, linkIndexA=self.lf_link_id
+            )
+        )
+
+        return rf_force, lf_force
 
     def draw_points(
         self,

examples/example_4_physics_simulation.py

@@ -5,14 +5,15 @@
 import numpy as np
 import pybullet as pb
 import pinocchio as pin
+from matplotlib import pyplot as plt
 
 from biped_walking_controller.foot import (
     compute_feet_path_and_poses,
     BezierCurveFootPathGenerator,
 )
 
 from biped_walking_controller.inverse_kinematic import InvKinSolverParams, solve_inverse_kinematics
-from biped_walking_controller.plot import plot_feet_and_com
+from biped_walking_controller.plot import plot_feet_and_com, plot_contact_forces
 
 from biped_walking_controller.preview_control import (
     PreviewControllerParams,
@@ -208,6 +209,9 @@ def main():
 
     zmp_pos = np.zeros((len(phases), 3))
 
+    rf_forces = np.zeros((len(phases), 1))
+    lf_forces = np.zeros((len(phases), 1))
+
     # We start the walking phase
     for k, _ in enumerate(phases[:-2]):
         # Get the current configuration of the robot from the simulator
@@ -285,7 +289,10 @@ def main():
             rf_pin_pos[k] = talos.data.oMf[talos.right_foot_id].translation
             rf_pb_pos[k], _ = simulator.get_robot_frame_pos("leg_right_6_link")
 
+            rf_forces[k], lf_forces[k] = simulator.get_contact_forces()
+
     if args.plot_results:
+
         zmp_ref_plot = np.zeros((zmp_ref.shape[0], 3))
         zmp_ref_plot[:, :2] = zmp_ref
 
@@ -305,6 +312,10 @@ def main():
             zmp_ref=zmp_ref_plot,
         )
 
+        plot_contact_forces(t=t, force_rf=rf_forces, force_lf=lf_forces)
+
+        plt.show()
+
     # Infinite loop to display the ending position
     while True:
         simulator.step()