JackM323/JackBot
1
0
Fork 0
Code Issues 3 Pull Requests Actions Packages Projects 2 Releases Wiki Activity

direct and residual movement inconsistency fixed

Browse Source
This commit is contained in:
JackM323
2026-08-27 21:40:57 +02:00
parent cd870a4afc
commit 4cc2d37d94
8 changed files with 109 additions and 73 deletions
Download Patch File Download Diff File Expand all files Collapse all files
ml/env.py
+40 -13
View File
@@ -184,7 +184,19 @@ class JackBotEnv(gym.Env):
return self._get_obs(), {}
def _get_obs(self) -> np.ndarray:
return self.robot.get_observation(command=self.command)
# Read raw joint angles from backend
raw_angles = np.asarray(self.robot.backend.get_joint_angles(), dtype=np.float32).flatten()
min_lim = self.min_joint_limits.flatten()
max_lim = self.max_joint_limits.flatten()
# Map raw joint radians [min, max] -> normalized [-1, 1]
normalized_joints = 2.0 * (raw_angles - min_lim) / (max_lim - min_lim) - 1.0
normalized_joints = np.clip(normalized_joints, -1.0, 1.0)
# Concatenate normalized joints with active command vector
obs = np.concatenate([normalized_joints, self.command]).astype(np.float32)
return obs
def step(self, action: np.ndarray) -> Tuple[np.ndarray, float, bool, bool, Dict[str, Any]]:
previous_action = self.last_action.copy()
@@ -194,25 +206,36 @@ class JackBotEnv(gym.Env):
self.last_last_action = self.last_action.copy()
self.last_action = action.copy()
# Command resampling
if self.random_command and (self.step_count >= self.next_cmd_resample_step or self._curriculum_advanced):
# Command resampling ONLY if random_command is True
if self.random_command and (
self.step_count >= self.next_cmd_resample_step or self._curriculum_advanced):
self.command = self.sample_command()
random_interval = np.random.randint(self.min_cmd_hold_steps, self.max_cmd_hold_steps + 1)
self.next_cmd_resample_step = self.step_count + random_interval
# Extract [vx, vy, omega]
# Extract active [vx, vy, omega]
cmd_vx, cmd_vy, cmd_omega = self.command
# Mirror main.py input resolution logic: update robot_state and vector_dirmov directly
# Mirror input resolution logic to keep robot state synchronized
self.robot.robot_state = "walking" if (abs(cmd_vx) > 0.01 or abs(cmd_vy) > 0.01 or abs(cmd_omega) > 0.01) else "idle"
self.robot.vector_dirmov = [float(cmd_vx), float(cmd_vy), float(cmd_omega)]
joint_range = np.minimum(
self.default_joint_angles - self.min_joint_limits,
self.max_joint_limits - self.default_joint_angles
)
target_angles = self.default_joint_angles + action * joint_range
self.robot.tick(action=target_angles)
# Direct Mode: Target joint scaling
action_flat = np.asarray(action, dtype=np.float32).flatten()
action_clipped = np.clip(action_flat, -1.0, 1.0)
if self.robot_mode == "direct":
# Map [-1, 1] linearly to physical joint limits [min, max]
min_lim = self.min_joint_limits.flatten()
max_lim = self.max_joint_limits.flatten()
target_angles = min_lim + (action_clipped + 1.0) * 0.5 * (max_lim - min_lim)
else:
# Residual mode mapping logic
target_angles = self.default_joint_angles.flatten() + action_clipped * 0.20
# Apply target joint angles to physics engine
self.robot.tick(action=target_angles, physics_substeps=4)
if self.robot_mode != "kinematics" and self.step_count % 60 == 0:
random_force = np.random.uniform(-2.0, 2.0, size=2)
@@ -222,8 +245,9 @@ class JackBotEnv(gym.Env):
self._update_distance_metrics()
self._update_curriculum()
# Build next observation preserving active command
obs = self._get_obs()
reward = self._compute_reward(action, previous_action)
reward = self._compute_reward(action_flat, previous_action)
self.cumulative_reward += reward
self.robot_reward += reward
@@ -235,6 +259,9 @@ class JackBotEnv(gym.Env):
if self.step_count % 120 == 0 and self.use_gui:
self._update_hud()
if self.use_gui:
time.sleep(1.0 / self.control_freq)
return obs, reward, terminated, truncated, info
def _update_distance_metrics(self):
ml/pretrain_bc.py
+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()
ml/run_eval.py
+13 -5
View File
@@ -33,11 +33,12 @@ def main():
use_gui=args.gui,
random_command=False,
max_episode_steps=args.max_steps_per_episode,
robot_mode="direct"
)
# Multi-Phase Configurations Suite
phase_configs = [
(CurriculumPhase.STAND_ONLY, "STAND", np.array([0.0, 0.0, 0.0], dtype=np.float32)),
#(CurriculumPhase.STAND_ONLY, "STAND", np.array([0.0, 0.0, 0.0], dtype=np.float32)),
(CurriculumPhase.FORWARD, "FORWARD GAIT", np.array([1.0, 0.0, 0.0], dtype=np.float32)),
(CurriculumPhase.TURN_AND_DIRECTION, "FORWARD + YAW TURN", np.array([0.5, 0.0, 0.4], dtype=np.float32)),
(CurriculumPhase.OMNI_DIRECTION, "STRIDE LATERAL", np.array([0.5, 0.5, 0.0], dtype=np.float32)),
@@ -62,9 +63,18 @@ def main():
for ep in range(args.episodes_per_phase):
obs, _ = env.reset()
# Force environment into active curriculum phase and lock command
# Force environment into active curriculum phase and lock command BEFORE getting obs
env.curriculum_phase = phase_enum
env.command = test_cmd.copy()
cmd_vx, cmd_vy, cmd_omega = test_cmd
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)]
# Get correct observation with test_cmd attached
obs = env._get_obs()
done = False
@@ -72,12 +82,10 @@ def main():
steps = 0
while not done:
# Enforce locked command each step
env.command = test_cmd.copy()
# Predict deterministic action from policy
action, _ = model.predict(obs, deterministic=True)
# Step environment
obs, reward, terminated, truncated, _ = env.step(action)
done = terminated or truncated
ml/run_eval_training.py
+1 -1
View File
@@ -21,7 +21,7 @@ def evaluate_kinematics(episode_length: int = 1000):
print("=" * 70 + "\n")
phase_configs = [
(CurriculumPhase.STAND_ONLY, "STAND", np.array([1.0, 0.0, 0.0], dtype=np.float32)),
(CurriculumPhase.STAND_ONLY, "STAND", np.array([1.0, 0.0, 0.0], dtype=np.float32)),
(CurriculumPhase.FORWARD, "FORWARD GAIT", np.array([1.0, 0.0, 0.0], dtype=np.float32)),
(CurriculumPhase.TURN_AND_DIRECTION, "FORWARD + YAW TURN", np.array([0.5, 0.0, 0.4], dtype=np.float32)),
(CurriculumPhase.OMNI_DIRECTION, "STRIDE LATERAL", np.array([0.5, 0.5, 0.0], dtype=np.float32)),
ml/run_train.py
+10 -5
View File
@@ -86,7 +86,12 @@ def main():
model = PPO.load(
args.pretrained_model,
env=vec_env,
learning_rate=1e-4, # Lower learning rate so RL fine-tunes without destroying base gait
learning_rate=5e-5, # Lower learning rate so RL fine-tunes without destroying base gait
ent_coef=0.001,
target_kl=0.05,
vf_coef=0.5,
max_grad_norm=0.5,
verbose=2,
tensorboard_log=args.log_dir,
device="cpu",
)
@@ -95,18 +100,18 @@ def main():
model = PPO(
policy="MlpPolicy",
env=vec_env,
learning_rate=1e-4,
learning_rate=5e-5,
n_steps=256,
batch_size=256,
n_epochs=10,
gamma=0.99,
gae_lambda=0.95,
clip_range=0.2,
ent_coef=0.01,
ent_coef=0.001,
target_kl=0.05,
vf_coef=0.5,
max_grad_norm=0.5,
verbose=1,
verbose=2,
tensorboard_log=args.log_dir,
device="cpu",
)
@@ -127,7 +132,7 @@ def main():
eval_env,
best_model_save_path=best_model_path,
log_path="ml/logs/results",
eval_freq=max(1, 20_000 // args.num_workers),
eval_freq=max(1, 50_000 // args.num_workers),
deterministic=True,
render=False,
)