ShunFang-cart/ESP
1
0
Fork 1
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Refactor IMU velocity estimator to improve gravity compensation and integrate world-frame velocity calculations

Browse Source
This commit is contained in:
Brunsmeier
2026-07-07 17:08:04 +08:00
parent fd286e6a9e
commit b4739362a5
1 changed files with 92 additions and 55 deletions
Download Patch File Download Diff File Expand all files Collapse all files

147
record_episode.py
View File

@ -239,27 +239,23 @@ class D405Camera:
# ============================================================ # ============================================================
# Simple IMU velocity estimator # Simple IMU velocity estimator Updated
# ============================================================ # ============================================================
class ProviderIMUVelocityEstimator: class ProviderWorldIMUVelocityEstimator:
""" """
Provider-style roll/pitch gravity compensation + velocity integration. Provider-style gravity/support compensation + world-frame velocity integration.
This follows verify.py: Pipeline:
gravity_body_g = [ 1. raw acc_body_g - gravity_body_g -> linear_acc_body_g
-sin(pitch), 2. estimate residual body-frame acceleration bias
cos(pitch) * sin(roll), 3. linear_acc_body_ms2 -> linear_acc_world_ms2 using R_body_to_world
cos(pitch) * cos(roll), 4. integrate velocity in world frame
] * gravity_sign
Then: Output:
linear_acc_body = raw_acc_body - gravity_body velocity_ms: world-frame velocity, unit m/s
linear_acc_ms2: world-frame linear acceleration, unit m/s^2
Important: acc_bias_ms2: body-frame residual acceleration bias, unit m/s^2
This is still an estimated IMU velocity, not ground truth.
It removes the static gravity/support acceleration approximately,
but velocity integration will still drift.
""" """
def __init__( def __init__(
@ -267,21 +263,20 @@ class ProviderIMUVelocityEstimator:
gravity_sign=1.0, gravity_sign=1.0,
bias_seconds=1.0, bias_seconds=1.0,
acc_deadband=0.10, acc_deadband=0.10,
vel_decay=0.995,
): ):
self.gravity_sign = gravity_sign self.gravity_sign = gravity_sign
self.bias_seconds = bias_seconds self.bias_seconds = bias_seconds
self.acc_deadband = acc_deadband self.acc_deadband = acc_deadband
self.vel_decay = vel_decay
self.start_t_ms = None self.start_t_ms = None
self.last_t_ms = None self.last_t_ms = None
self.bias_ready = False self.bias_ready = False
self.bias_samples = [] self.bias_samples = []
self.acc_bias = np.zeros(3, dtype=float) self.acc_bias_body_ms2 = np.zeros(3, dtype=float)
self.v = np.zeros(3, dtype=float) self.v_world = np.zeros(3, dtype=float)
self.last_acc_world_ms2 = np.zeros(3, dtype=float)
def reset(self): def reset(self):
self.start_t_ms = None self.start_t_ms = None
@ -289,9 +284,10 @@ class ProviderIMUVelocityEstimator:
self.bias_ready = False self.bias_ready = False
self.bias_samples = [] self.bias_samples = []
self.acc_bias[:] = 0.0 self.acc_bias_body_ms2[:] = 0.0
self.v[:] = 0.0 self.v_world[:] = 0.0
self.last_acc_world_ms2[:] = 0.0
def gravity_body_g(self, roll_deg, pitch_deg): def gravity_body_g(self, roll_deg, pitch_deg):
roll = math.radians(roll_deg) roll = math.radians(roll_deg)
@ -303,6 +299,22 @@ class ProviderIMUVelocityEstimator:
return np.array([g_x, g_y, g_z], dtype=float) * self.gravity_sign return np.array([g_x, g_y, g_z], dtype=float) * self.gravity_sign
def rot_body_to_world(self, roll_deg, pitch_deg, yaw_deg):
roll = math.radians(roll_deg)
pitch = math.radians(pitch_deg)
yaw = math.radians(yaw_deg)
cr, sr = math.cos(roll), math.sin(roll)
cp, sp = math.cos(pitch), math.sin(pitch)
cy, sy = math.cos(yaw), math.sin(yaw)
# R_body_to_world = Rz(yaw) @ Ry(pitch) @ Rx(roll)
return np.array([
[cy * cp, cy * sp * sr - sy * cr, cy * sp * cr + sy * sr],
[sy * cp, sy * sp * sr + cy * cr, sy * sp * cr - cy * sr],
[-sp, cp * sr, cp * cr],
], dtype=float)
def update(self, packet): def update(self, packet):
t_ms = packet["t_ms"] t_ms = packet["t_ms"]
@ -313,61 +325,80 @@ class ProviderIMUVelocityEstimator:
roll_deg = packet["angle_deg"][0] roll_deg = packet["angle_deg"][0]
pitch_deg = packet["angle_deg"][1] pitch_deg = packet["angle_deg"][1]
yaw_deg = packet["angle_deg"][2]
# ---------------- gravity/support removal in body frame ----------------
gravity_g = self.gravity_body_g( gravity_g = self.gravity_body_g(
roll_deg=roll_deg, roll_deg=roll_deg,
pitch_deg=pitch_deg, pitch_deg=pitch_deg,
) )
linear_acc_g_raw = acc_body_g - gravity_g linear_acc_body_g_raw = acc_body_g - gravity_g
linear_acc_ms2_raw = linear_acc_g_raw * 9.81 linear_acc_body_ms2_raw = linear_acc_body_g_raw * 9.80665
elapsed_s = ((t_ms - self.start_t_ms) & 0xFFFFFFFF) / 1000.0 elapsed_s = ((t_ms - self.start_t_ms) & 0xFFFFFFFF) / 1000.0
# First N seconds are used as residual bias estimation. # ---------------- bias estimation in body frame ----------------
# During this period velocity stays zero. # During this period, velocity stays zero.
if not self.bias_ready: if not self.bias_ready:
self.bias_samples.append(linear_acc_ms2_raw) self.bias_samples.append(linear_acc_body_ms2_raw.copy())
if elapsed_s >= self.bias_seconds: if elapsed_s >= self.bias_seconds:
if self.bias_samples: if self.bias_samples:
self.acc_bias = np.mean( self.acc_bias_body_ms2 = np.mean(
np.array(self.bias_samples, dtype=float), np.array(self.bias_samples, dtype=float),
axis=0, axis=0,
) )
self.bias_ready = True self.bias_ready = True
print(f"[IMU] Provider acc bias estimated: {self.acc_bias}") print(
"[IMU] Provider body-frame acc bias estimated: "
f"{self.acc_bias_body_ms2}"
)
self.last_t_ms = t_ms self.last_t_ms = t_ms
self.last_acc_world_ms2[:] = 0.0
linear_acc_ms2 = linear_acc_ms2_raw - self.acc_bias
return { return {
"gravity_g": gravity_g.copy(), "gravity_g": gravity_g.copy(),
"linear_acc_ms2": linear_acc_ms2.copy(), "linear_acc_ms2": self.last_acc_world_ms2.copy(),
"velocity_ms": self.v.copy(), "velocity_ms": self.v_world.copy(),
"acc_bias_ms2": self.acc_bias.copy(), "acc_bias_ms2": self.acc_bias_body_ms2.copy(),
"bias_ready": self.bias_ready, "bias_ready": self.bias_ready,
} }
# ---------------- remove body-frame residual bias ----------------
linear_acc_body_ms2 = linear_acc_body_ms2_raw - self.acc_bias_body_ms2
# ---------------- deadband before rotation ----------------
# Norm is unchanged by pure rotation, so before/after rotation both okay.
if np.linalg.norm(linear_acc_body_ms2) < self.acc_deadband:
linear_acc_body_ms2[:] = 0.0
# ---------------- body frame -> world frame ----------------
R_bw = self.rot_body_to_world(
roll_deg=roll_deg,
pitch_deg=pitch_deg,
yaw_deg=yaw_deg,
)
linear_acc_world_ms2 = R_bw @ linear_acc_body_ms2
# ---------------- dt ----------------
delta_ms = (t_ms - self.last_t_ms) & 0xFFFFFFFF delta_ms = (t_ms - self.last_t_ms) & 0xFFFFFFFF
self.last_t_ms = t_ms self.last_t_ms = t_ms
linear_acc_ms2 = linear_acc_ms2_raw - self.acc_bias
if np.linalg.norm(linear_acc_ms2) < self.acc_deadband:
linear_acc_ms2[:] = 0.0
if 0 < delta_ms <= 200: if 0 < delta_ms <= 200:
dt = delta_ms / 1000.0 dt = delta_ms / 1000.0
self.v = self.v * self.vel_decay + linear_acc_ms2 * dt self.v_world += linear_acc_world_ms2 * dt
self.last_acc_world_ms2 = linear_acc_world_ms2.copy()
return { return {
"gravity_g": gravity_g.copy(), "gravity_g": gravity_g.copy(),
"linear_acc_ms2": linear_acc_ms2.copy(), "linear_acc_ms2": self.last_acc_world_ms2.copy(),
"velocity_ms": self.v.copy(), "velocity_ms": self.v_world.copy(),
"acc_bias_ms2": self.acc_bias.copy(), "acc_bias_ms2": self.acc_bias_body_ms2.copy(),
"bias_ready": self.bias_ready, "bias_ready": self.bias_ready,
} }
@ -407,6 +438,7 @@ def infer_scissor_state(packet):
return "invalid_both_active" return "invalid_both_active"
# ============================================================ # ============================================================
# Episode writer # Episode writer
# ============================================================ # ============================================================
@ -647,12 +679,17 @@ def run(args):
active_value=args.record_active_value, active_value=args.record_active_value,
debounce_s=args.debounce, debounce_s=args.debounce,
) )
velocity_estimator = ProviderIMUVelocityEstimator( # velocity_estimator = ProviderIMUVelocityEstimator(
gravity_sign=args.gravity_sign, # gravity_sign=args.gravity_sign,
bias_seconds=args.bias_seconds, # bias_seconds=args.bias_seconds,
acc_deadband=args.acc_deadband, # acc_deadband=args.acc_deadband,
vel_decay=args.velocity_decay, # vel_decay=args.velocity_decay,
) # )
velocity_estimator = ProviderWorldIMUVelocityEstimator(
gravity_sign=args.gravity_sign,
bias_seconds=args.bias_seconds,
acc_deadband=args.acc_deadband,
)
writer = EpisodeWriter( writer = EpisodeWriter(
root_dir=args.dataset, root_dir=args.dataset,
@ -888,12 +925,12 @@ def parse_args():
help="JPEG quality for RGB image saving.", help="JPEG quality for RGB image saving.",
) )
parser.add_argument( # parser.add_argument(
"--velocity-decay", # "--velocity-decay",
type=float, # type=float,
default=0.995, # default=0.995,
help="Decay factor for simple raw IMU velocity integration.", # help="Decay factor for simple raw IMU velocity integration.",
) # )
parser.add_argument( parser.add_argument(
"--flush-every", "--flush-every",