Merge 2631c9b into 201fe7a

agents/pink_balancer/config/bullet.gin

@@ -16,6 +16,9 @@
 
 import agents.pink_balancer.wheel_balancer
 
+BalancingPlaneInIMUFrame.sagittal_axis = 0
+BalancingPlaneInIMUFrame.vertical_axis = 2
+
 Gains.pitch_damping = 10.0
 Gains.pitch_stiffness = 20.0
 Gains.position_damping = 5.0

agents/pink_balancer/config/pi3hat.gin

@@ -16,6 +16,9 @@
 
 import agents.pink_balancer.wheel_balancer
 
+BalancingPlaneInIMUFrame.sagittal_axis = 0
+BalancingPlaneInIMUFrame.vertical_axis = 2
+
 # Smooth wheels with radius=5 cm, wood flooring
 # ---------------------------------------------
 #

agents/pink_balancer/wheel_balancer.py

@@ -19,13 +19,15 @@
 Keep Upkie up! Using its wheels.
 """
 
+from dataclasses import dataclass
 from typing import Any, Dict, Tuple
 
 import gin
 import numpy as np
 
 from observers.base_pitch import compute_base_pitch_from_imu
 from utils.clamp import clamp, clamp_abs
+# from utils.exceptions import FallDetected
 from utils.filters import abs_bounded_derivative_filter, low_pass_filter
 
 
@@ -131,6 +133,14 @@ def __repr__(self):
                 f"position_stiffness={self.position_stiffness}, "
             )
 
+    @gin.configurable
+    @dataclass
+    class BalancingPlaneInIMUFrame:
+        """Axes of the balancing plane expressed in the IMU frame."""
+
+        sagittal_axis: int
+        vertical_axis: int
+
     air_return_period: float
     error: np.ndarray
     gains: Gains
@@ -194,6 +204,7 @@ def __init__(
         """
         assert 0.0 <= turning_deadband <= 1.0
         self.air_return_period = air_return_period
+        self.balancing_plane = WheelBalancer.BalancingPlaneInIMUFrame()
         self.error = np.zeros(2)
         self.fall_pitch = fall_pitch
         self.gains = WheelBalancer.Gains()  # type: ignore
@@ -317,12 +328,17 @@ def cycle(self, observation: Dict[str, Any], dt: float) -> None:
         self.update_target_ground_velocity(observation, dt)
         self.update_target_yaw_velocity(observation, dt)
 
-        pitch = compute_base_pitch_from_imu(observation["imu"]["orientation"])
+        pitch = compute_base_pitch_from_imu(
+            observation["imu"]["orientation"],
+            sagittal_axis=self.balancing_plane.sagittal_axis,
+            vertical_axis=self.balancing_plane.vertical_axis,
+        )
         self.pitch = pitch
         if abs(pitch) > self.fall_pitch:
             self.integral_error_velocity = 0.0  # [m] / [s]
             self.ground_velocity = 0.0  # [m] / [s]
             return
+            # raise FallDetected(f"Base angle {pitch=:.3} rad denotes a fall")
 
         ground_position = observation["wheel_odometry"]["position"]
         floor_contact = observation["floor_contact"]["contact"]

observers/base_pitch/base_pitch.py

@@ -24,6 +24,8 @@
 
 def compute_pitch_frame_in_parent(
     orientation_frame_in_parent: np.ndarray,
+    sagittal_axis: int = 0,
+    vertical_axis: int = 2,
 ) -> float:
     """
     Get pitch angle of a given frame relative to the parent vertical.
@@ -48,6 +50,8 @@ def compute_pitch_frame_in_parent(
     Args:
         orientation_frame_in_parent: Rotation matrix from the target frame to
             the parent frame.
+        sagittal_axis: Index of the sagittal axis in vector representation.
+        vertical_axis: Index of the vertical axis in vector representation.
 
     Returns:
         Angle from the parent z-axis (gravity) to the frame z-axis.
@@ -60,13 +64,14 @@ def compute_pitch_frame_in_parent(
         which is ``(x)`` in the above schematic (pointing away) and not ``(.)``
         (poiting from screen to the reader).
     """
-    sagittal = orientation_frame_in_parent[:, 0]
+    sagittal = orientation_frame_in_parent[:, sagittal_axis]
     sagittal /= np.linalg.norm(sagittal)  # needed for prec around cos_pitch=0.
-    heading_in_parent = sagittal - sagittal[2] * np.array([0.0, 0.0, 1.0])
+    e_z = np.array([0.0, 0.0, 1.0])
+    heading_in_parent = sagittal - sagittal[vertical_axis] * e_z
     heading_in_parent /= np.linalg.norm(heading_in_parent)
-    if orientation_frame_in_parent[2, 2] < 0:
+    if orientation_frame_in_parent[vertical_axis, vertical_axis] < 0:
         heading_in_parent *= -1.0
-    sign = +1 if sagittal[2] < 0 else -1
+    sign = +1 if sagittal[vertical_axis] < 0 else -1
     cos_pitch = clamp(np.dot(sagittal, heading_in_parent), -1.0, 1.0)
     pitch: float = sign * np.arccos(cos_pitch)
     return pitch
@@ -95,48 +100,54 @@ def rotation_matrix_from_quaternion(
     return np.array(
         [
             [
-                1 - 2 * (qy**2 + qz**2),
+                1 - 2 * (qy ** 2 + qz ** 2),
                 2 * (qx * qy - qz * qw),
                 2 * (qw * qy + qx * qz),
             ],
             [
                 2 * (qx * qy + qz * qw),
-                1 - 2 * (qx**2 + qz**2),
+                1 - 2 * (qx ** 2 + qz ** 2),
                 2 * (qy * qz - qx * qw),
             ],
             [
                 2 * (qx * qz - qy * qw),
                 2 * (qy * qz + qx * qw),
-                1 - 2 * (qx**2 + qy**2),
+                1 - 2 * (qx ** 2 + qy ** 2),
             ],
         ]
     )
 
 
 def compute_base_pitch_from_imu(
-    quat_imu_in_world: Tuple[float, float, float, float]
+    quat_imu_in_ars: Tuple[float, float, float, float],
+    sagittal_axis: int = 0,
+    vertical_axis: int = 2,
 ) -> float:
     """
-    Get pitch angle of the *base* relative to the world vertical.
+    Get pitch angle of the *base* relative to the attitude reference frame.
 
     Args:
-        quat_imu_in_world: Quaternion representing the rotation matrix from
-            the *IMU* frame to the world frame, in ``[w, x, y, z]`` format.
+        quat_imu_in_ars: Quaternion representing the rotation matrix from
+            the *IMU* frame to the ARS frame, in ``[w, x, y, z]`` format.
+        sagittal_axis: Index of the sagittal axis in vector representation.
+        vertical_axis: Index of the vertical axis in vector representation.
 
     Returns:
-        Angle from the world z-axis (gravity) to the base z-axis.
-        Equivalently, angle that the sagittal vector of the base frame makes
-        with the heading vector.
+        Angle from the world z-axis (unit vector opposite to gravity) to the
+        base z-axis. Equivalently, angle that the sagittal vector of the base
+        frame makes with the heading vector.
 
     Note:
         The output pitch angle is for the base (x-axis pointing forward), but
         the input orientation is that of the IMU. Keep in mind that the IMU
         frame is turned 180 degrees around the yaw axis of the base frame.
+
     """
-    orientation_imu_in_world = rotation_matrix_from_quaternion(
-        quat_imu_in_world
-    )
+    rotation_imu_to_ars = rotation_matrix_from_quaternion(quat_imu_in_ars)
+    # https://github.com/mjbots/pi3hat/blob/master/docs/reference.md#orientation
+    rotation_ars_to_world = np.diag([1.0, -1.0, -1.0])
+    rotation_imu_to_world = rotation_ars_to_world @ rotation_imu_to_ars
     pitch_imu_in_world = compute_pitch_frame_in_parent(
-        orientation_imu_in_world
+        rotation_imu_to_world, sagittal_axis, vertical_axis
     )
-    return -pitch_imu_in_world  # 180 degrees yaw rotation
+    return pitch_imu_in_world

observers/base_pitch/tests/base_pitch_test.py

@@ -24,6 +24,7 @@
 import numpy as np
 
 from observers.base_pitch import (
+    compute_base_pitch_from_imu,
     compute_pitch_frame_in_parent,
 )
 
@@ -73,6 +74,31 @@ def test_orientation_not_neatly_normalized(self, theta=1e-3):
         imu_pitch = compute_pitch_frame_in_parent(orientation_imu_in_world)
         self.assertTrue(np.isclose(imu_pitch, theta))
 
+    def test_balancing_plane_lateral_in_imu_frame(self):
+        quat_imu_in_ars = np.array(
+            [
+                0.008472769239730098,
+                -0.9953038144146671,
+                -0.09639792825405252,
+                -0.002443076206500708,
+            ]
+        )
+        rotation_imu_to_ars = np.array(
+            [
+                [0.98140294, 0.19193185, 0.00322969],
+                [0.19184905, -0.9812713, 0.01733697],
+                [0.00649672, -0.01639494, -0.99984449],
+            ]
+        )
+        pitch_imu_in_ars = compute_pitch_frame_in_parent(
+            rotation_imu_to_ars, sagittal_axis=1, vertical_axis=2
+        )
+        self.assertAlmostEqual(pitch_imu_in_ars, 3.125, places=3)
+        pitch_imu_in_world = compute_base_pitch_from_imu(
+            quat_imu_in_ars, sagittal_axis=1, vertical_axis=2
+        )
+        self.assertAlmostEqual(pitch_imu_in_world, -0.016, places=3)
+
 
 if __name__ == "__main__":
     unittest.main()