direct and residual movement inconsistency fixed

This commit is contained in:
2026-08-27 21:40:57 +02:00
parent cd870a4afc
commit 4cc2d37d94
8 changed files with 109 additions and 73 deletions
+25 -22
View File
@@ -21,7 +21,8 @@ def collect_kinematics_dataset(num_samples: int = 100_000, use_gui: bool = False
"""Collects (Observation, Action) pairs directly from Kinematics Teacher."""
print(f"\n[Pretrain] Collecting {num_samples} samples from Kinematics Teacher (GUI={use_gui})...")
env = JackBotEnv(use_gui=use_gui, robot_mode="kinematics")
# Disable random command resampling inside env.step so manual command locks persist
env = JackBotEnv(use_gui=use_gui, robot_mode="kinematics", random_command=False)
env.curriculum_phase = CurriculumPhase.FULL_COMMAND
observations = []
@@ -32,39 +33,41 @@ def collect_kinematics_dataset(num_samples: int = 100_000, use_gui: bool = False
# --- PROGRESS BAR: Data Collection ---
pbar = tqdm(range(num_samples), desc=" Collecting Data", unit="step")
for i in pbar:
# 1. Sample random movement command
# 1. Update command and vector targets every 120 steps
if i % 120 == 0:
env.command = env.sample_command()
cmd_vx, cmd_vy, cmd_omega = env.command
env.robot.robot_state = (
"walking"
if (abs(cmd_vx) > 0.01 or abs(cmd_vy) > 0.01 or abs(cmd_omega) > 0.01)
else "idle"
)
env.robot.vector_dirmov = [float(cmd_vx), float(cmd_vy), float(cmd_omega)]
# 2. Let Robot compute Kinematics target angles
cmd_vx, cmd_vy, cmd_omega = env.command
env.robot.robot_state = "walking" if (abs(cmd_vx) > 0.01 or abs(cmd_vy) > 0.01 or abs(cmd_omega) > 0.01) else "idle"
env.robot.vector_dirmov = [float(cmd_vx), float(cmd_vy), float(cmd_omega)]
# Advance internal IK tick to calculate joint angles
env.robot.tick()
# Extract .data and flatten (6, 3) matrix to 18-dim 1D array
# 2. Capture observation BEFORE stepping environment
current_obs = env._get_obs()
# 3. Step environment ONCE (updates kinematics solver, PyBullet physics, and computes IK)
obs, _, terminated, truncated, _ = env.step(np.zeros(18, dtype=np.float32))
# 4. Extract procedural IK joint targets computed during this step
target_ik_rad = env.robot.current_rad.data.flatten().copy()
# 3. Convert target angles back to normalized [-1, 1] action space
normalized_action = np.where(
target_ik_rad >= env.default_joint_angles,
(target_ik_rad - env.default_joint_angles) / np.maximum(1e-5, env.max_joint_limits - env.default_joint_angles),
(target_ik_rad - env.default_joint_angles) / np.maximum(1e-5, env.default_joint_angles - env.min_joint_limits)
)
# Step 5: Convert target radians directly to [-1, 1] relative to joint limits
min_lim = env.min_joint_limits.flatten()
max_lim = env.max_joint_limits.flatten()
normalized_action = 2.0 * (target_ik_rad - min_lim) / (max_lim - min_lim) - 1.0
normalized_action = np.clip(normalized_action, -1.0, 1.0)
# 4. Save sample
observations.append(obs.copy())
# Step 6: Store matching input (obs) and target ground truth (normalized_action)
observations.append(current_obs.copy())
actions.append(normalized_action.copy())
# Step simulation environment
obs, _, terminated, truncated, _ = env.step(normalized_action)
if use_gui:
time.sleep(1.0 / 60.0)
# 7. Handle episode boundaries using terminated and truncated
if terminated or truncated:
obs, _ = env.reset()