diff --git a/README.md b/README.md index ea44d7e..8b907fd 100644 --- a/README.md +++ b/README.md @@ -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: @@ -307,5 +297,4 @@ 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. \ No newline at end of file +* **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. \ No newline at end of file diff --git a/Robot.py b/Robot.py index a046289..9b5bf8f 100644 --- a/Robot.py +++ b/Robot.py @@ -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,12 +200,12 @@ 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) - self.step_sim() + # 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: """Procedural Tripod Gait solver.""" @@ -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.""" diff --git a/ml/env.py b/ml/env.py index f5f3fd7..87ff77e 100644 --- a/ml/env.py +++ b/ml/env.py @@ -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): diff --git a/ml/pretrain_bc.py b/ml/pretrain_bc.py index 51c2dac..3bbd3f4 100644 --- a/ml/pretrain_bc.py +++ b/ml/pretrain_bc.py @@ -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() diff --git a/ml/run_eval.py b/ml/run_eval.py index 2188c56..15e9509 100644 --- a/ml/run_eval.py +++ b/ml/run_eval.py @@ -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 diff --git a/ml/run_eval_training.py b/ml/run_eval_training.py index 5d95508..8128d51 100644 --- a/ml/run_eval_training.py +++ b/ml/run_eval_training.py @@ -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)), diff --git a/ml/run_train.py b/ml/run_train.py index ddb3551..bb9f317 100644 --- a/ml/run_train.py +++ b/ml/run_train.py @@ -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, ) diff --git a/simulation.py b/simulation.py index 7c264f8..1851ac2 100644 --- a/simulation.py +++ b/simulation.py @@ -85,7 +85,7 @@ class Simulation: force=30, physicsClientId=self.physics_client ) - + def hard_reset_joint_angles(self, target_angles: Union[np.ndarray, dt.RadArray]) -> None: """Instantly teleports joint angles to target positions, clearing velocity state.""" radflat = target_angles.data.flatten() if isinstance(target_angles, dt.RadArray) else target_angles.flatten()