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
README.md
+1 -12
View File
@@ -144,17 +144,7 @@ The machine learning system trains the hexapod to move using a combination of Be
### ML Workflow Steps
The policy input contains:
* 18 joint angles from the robot, and
* 4 command dimensions describing the desired motion direction and yaw rate.
The policy output is:
* 18 continuous values, one per joint, describing how much each joint should change.
This means the learning method is not controlling a discrete gait state directly. Instead, it learns a continuous control policy for the whole body.
The machine learning system trains the hexapod to move using a combination of Behavioral Cloning (BC) (optional) and Proximal Policy Optimization (PPO) driven by a multi-phase curriculum.
#### Methods implemented
The current ML stack uses:
@@ -308,4 +298,3 @@ In multi-robot vectorized training (`JackBotEnv`), individual robot failures are
When training with `--gui`, `JackBotEnv` includes dynamic visual feedback mechanisms:
* **Metrics HUD:** A live on-screen text overlay tracking active episode count, total step rate (FPS), average base height, roll/pitch angles, and cumulative per-robot rewards.
* **Leader Crown ($\text{👑}$):** A floating crown debug indicator tracks and positions itself directly above the base of whichever robot is achieving the highest cumulative reward in the multi-robot grid.
Robot.py
+16 -12
View File
@@ -188,11 +188,7 @@ class Robot:
self.current_state = STATE_REGISTRY[next_state_key]
self.current_state.enter(self)
def tick(self, action: Optional[np.ndarray] = None) -> None:
"""
Unified control loop tick.
Processes commands through direct RL, residual RL, or State Machine kinematics.
"""
def tick(self, action: Optional[np.ndarray] = None, physics_substeps: int = 4) -> None:
vx, vy, omega = self.vector_dirmov
if self.mode == "direct":
@@ -204,11 +200,11 @@ class Robot:
if action is not None:
self.apply_rl_action_delta(action)
else: # "kinematics" / standard State Machine execution
next_state_key = self.current_state.execute(self)
if next_state_key:
self.transition_to(next_state_key)
else: # kinematics mode
self.step_kinematic_gait(vx, vy, omega)
# Step PyBullet engine sub-steps to allow physics actuation
for _ in range(physics_substeps):
self.step_sim()
def step_kinematic_gait(self, vx: float, vy: float, omega: float) -> None:
@@ -265,9 +261,17 @@ class Robot:
self.set_joint_angles(target_rad)
def apply_rl_action(self, action: np.ndarray) -> None:
action = np.asarray(action, dtype=np.float32)
new_rad = dt.RadArray(data=action.reshape(self.current_rad.data.shape))
self.set_joint_angles(new_rad)
"""
Applies direct RL joint action deltas to the current joint positions.
"""
action_flat = np.asarray(action, dtype=np.float32).flatten()
# Keep internal memory updated in (6, 3) format for Kinematic/IK math
self.current_rad = dt.RadArray(data=action_flat.reshape(6, 3))
# Send target positions to the PyBullet backend
if self.backend:
self.backend.send_angles(self.current_rad)
def apply_rl_action_delta(self, action: np.ndarray) -> None:
"""Applies action deltas on top of joint state for Residual RL."""
ml/env.py
+39 -12
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
+22 -19
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()
# 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.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()
# 2. Capture observation BEFORE stepping environment
current_obs = env._get_obs()
# Extract .data and flatten (6, 3) matrix to 18-dim 1D array
# 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_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,
)