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

141
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:
gravity_body_g = [
-sin(pitch),
cos(pitch) * sin(roll),
cos(pitch) * cos(roll),
] * gravity_sign
Pipeline:
1. raw acc_body_g - gravity_body_g -> linear_acc_body_g
2. estimate residual body-frame acceleration bias
3. linear_acc_body_ms2 -> linear_acc_world_ms2 using R_body_to_world
4. integrate velocity in world frame
Then:
linear_acc_body = raw_acc_body - gravity_body
Important:
This is still an estimated IMU velocity, not ground truth.
It removes the static gravity/support acceleration approximately,
but velocity integration will still drift.
Output:
velocity_ms: world-frame velocity, unit m/s
linear_acc_ms2: world-frame linear acceleration, unit m/s^2
acc_bias_ms2: body-frame residual acceleration bias, unit m/s^2
"""
def __init__(
@ -267,21 +263,20 @@ class ProviderIMUVelocityEstimator:
gravity_sign=1.0,
bias_seconds=1.0,
acc_deadband=0.10,
vel_decay=0.995,
):
self.gravity_sign = gravity_sign
self.bias_seconds = bias_seconds
self.acc_deadband = acc_deadband
self.vel_decay = vel_decay
self.start_t_ms = None
self.last_t_ms = None
self.bias_ready = False
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):
self.start_t_ms = None
@ -289,9 +284,10 @@ class ProviderIMUVelocityEstimator:
self.bias_ready = False
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):
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
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):
t_ms = packet["t_ms"]
@ -313,61 +325,80 @@ class ProviderIMUVelocityEstimator:
roll_deg = packet["angle_deg"][0]
pitch_deg = packet["angle_deg"][1]
yaw_deg = packet["angle_deg"][2]
# ---------------- gravity/support removal in body frame ----------------
gravity_g = self.gravity_body_g(
roll_deg=roll_deg,
pitch_deg=pitch_deg,
)
linear_acc_g_raw = acc_body_g - gravity_g
linear_acc_ms2_raw = linear_acc_g_raw * 9.81
linear_acc_body_g_raw = acc_body_g - gravity_g
linear_acc_body_ms2_raw = linear_acc_body_g_raw * 9.80665
elapsed_s = ((t_ms - self.start_t_ms) & 0xFFFFFFFF) / 1000.0
# First N seconds are used as residual bias estimation.
# During this period velocity stays zero.
# ---------------- bias estimation in body frame ----------------
# During this period, velocity stays zero.
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 self.bias_samples:
self.acc_bias = np.mean(
self.acc_bias_body_ms2 = np.mean(
np.array(self.bias_samples, dtype=float),
axis=0,
)
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
linear_acc_ms2 = linear_acc_ms2_raw - self.acc_bias
self.last_acc_world_ms2[:] = 0.0
return {
"gravity_g": gravity_g.copy(),
"linear_acc_ms2": linear_acc_ms2.copy(),
"velocity_ms": self.v.copy(),
"acc_bias_ms2": self.acc_bias.copy(),
"linear_acc_ms2": self.last_acc_world_ms2.copy(),
"velocity_ms": self.v_world.copy(),
"acc_bias_ms2": self.acc_bias_body_ms2.copy(),
"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
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:
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 {
"gravity_g": gravity_g.copy(),
"linear_acc_ms2": linear_acc_ms2.copy(),
"velocity_ms": self.v.copy(),
"acc_bias_ms2": self.acc_bias.copy(),
"linear_acc_ms2": self.last_acc_world_ms2.copy(),
"velocity_ms": self.v_world.copy(),
"acc_bias_ms2": self.acc_bias_body_ms2.copy(),
"bias_ready": self.bias_ready,
}
@ -407,6 +438,7 @@ def infer_scissor_state(packet):
return "invalid_both_active"
# ============================================================
# Episode writer
# ============================================================
@ -647,12 +679,17 @@ def run(args):
active_value=args.record_active_value,
debounce_s=args.debounce,
)
velocity_estimator = ProviderIMUVelocityEstimator(
# velocity_estimator = ProviderIMUVelocityEstimator(
# gravity_sign=args.gravity_sign,
# bias_seconds=args.bias_seconds,
# acc_deadband=args.acc_deadband,
# vel_decay=args.velocity_decay,
# )
velocity_estimator = ProviderWorldIMUVelocityEstimator(
gravity_sign=args.gravity_sign,
bias_seconds=args.bias_seconds,
acc_deadband=args.acc_deadband,
vel_decay=args.velocity_decay,
)
)
writer = EpisodeWriter(
root_dir=args.dataset,
@ -888,12 +925,12 @@ def parse_args():
help="JPEG quality for RGB image saving.",
)
parser.add_argument(
"--velocity-decay",
type=float,
default=0.995,
help="Decay factor for simple raw IMU velocity integration.",
)
# parser.add_argument(
# "--velocity-decay",
# type=float,
# default=0.995,
# help="Decay factor for simple raw IMU velocity integration.",
# )
parser.add_argument(
"--flush-every",