direct and residual movement inconsistency fixed
This commit is contained in:
+25
-22
@@ -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()
|
||||
|
||||
|
||||
Reference in New Issue
Block a user