adjusted rewards and requirements

ml/env.py

@@ -46,12 +46,14 @@ class JackBotEnv(gym.Env):
         self.start_pose = start_pose
         self.max_episode_steps = max_episode_steps
         self.urdf_path = urdf_path
+        self.max_robot_speed = 0.8
 
         self.episode_count = 0
         self.step_count = 0
         self.total_steps = 0
         self.cumulative_reward = 0.0
         self.robot_rewards = [0.0]
+        self.consecutive_still_steps = [0]
         self.failed_robots_mask = [False]
         self.episode_height_sum = 0.0
         self._first_reset = True
@@ -106,7 +108,6 @@ class JackBotEnv(gym.Env):
         self.curriculum_stage_requirements = {
             CurriculumPhase.FORWARD: {
                 "survival_steps": 400,          # Must survive ~8 seconds
-                "max_displacement": 0.25,       # Must remain within 0.25m radius
                 "min_avg_height_ratio": 0.90,   # Average height >= 90% of target
                 "stability_roll_pitch": 0.18,   # Max ~10 degrees tilt
             },
@@ -150,17 +151,24 @@ class JackBotEnv(gym.Env):
             vz = 0.0
             omega = 0.0
         elif phase == CurriculumPhase.FORWARD:
-            # Phase 1: Straight Forward Walking
+            # Phase 1: Straight Forward Walking + Standing Checks
-            vx = np.random.uniform(0.5, 1.0)
+            if np.random.random() < 0.30:
-            vy = 0.0
+                # 30% chance: Stand Still Command (vx = 0)
-            vz = 0.0
+                vx = 0.0
-            omega = 0.0
+            else:
+                # 70% chance: Forward Walk Command
+                vx = np.random.uniform(0.15, 0.50)
+            
+            vy, vz, omega = 0.0, 0.0, 0.0
+
         elif phase == CurriculumPhase.TURN_AND_DIRECTION:
-            # Phase 2: Forward/Backward + Turning
+            # Phase 2: Forward/Backward + Turning + Standing
+            if np.random.random() < 0.20:
+                vx, omega = 0.0, 0.0
+            else:
                 vx = np.random.uniform(-1.0, 1.0)
-            vy = 0.0
-            vz = 0.0
                 omega = np.random.uniform(-0.8, 0.8)
+            vy, vz = 0.0, 0.0
         elif phase == CurriculumPhase.OMNI_DIRECTION:
             # Phase 3: Full Omnidirectional Movement
             vx = np.random.uniform(-1.0, 1.0)
@@ -195,6 +203,7 @@ class JackBotEnv(gym.Env):
             if self.use_gui:
                 self.sim_manager.set_robot_color(pb_id, [1.0, 1.0, 1.0, 1.0])
 
+        self.consecutive_still_steps = [0 for _ in range(len(self.pb_robots))]
         self.last_action = np.zeros(self.action_space.shape[0], dtype=np.float32)
         self.max_distance_from_start = [0.0]
         self.max_survival_steps = 0
@@ -406,56 +415,78 @@ class JackBotEnv(gym.Env):
             linear_vel, angular_vel = self.sim_manager.get_robot_velocity(pb_id)
             command = self.commands[idx]
 
-            cmd_vx = command[0]
+            cmd_vx = command[0]  # Normalized -1.0 to 1.0
-            cmd_vy = command[1]
+            cmd_vy = command[1]  # Normalized -1.0 to 1.0
             cmd_yaw = command[3]
 
-            # 1. LINEAR VECTOR SPEED MAXIMIZATION
             cmd_dir = np.array([cmd_vx, cmd_vy], dtype=np.float32)
             cmd_norm = np.linalg.norm(cmd_dir)
 
-            if cmd_norm > 0.05:
-                unit_cmd_dir = cmd_dir / cmd_norm
             actual_vel_2d = np.array([linear_vel[0], linear_vel[1]], dtype=np.float32)
 
+            # -----------------------------------------------------------------
+            # 1. NORMALIZED SPEED TRACKING (1.0 = v_max)
+            # -----------------------------------------------------------------
+            if cmd_norm > 0.05:
+                unit_cmd_dir = cmd_dir / cmd_norm
                 aligned_speed = float(np.dot(actual_vel_2d, unit_cmd_dir))
-                linear_speed_reward = 2.5 * aligned_speed
 
+                target_speed = cmd_norm * self.max_robot_speed  # e.g., 0.35 * 0.8 = 0.28 m/s
+
+                if aligned_speed <= 0.0:
+                    # Zero or backward movement gets ZERO reward
+                    linear_speed_reward = 0.0
+                elif aligned_speed <= target_speed:
+                    # Smoothly scales from 0.0 to +5.0 as speed approaches target
+                    linear_speed_reward = 5.0 * (aligned_speed / target_speed)
+                else:
+                    # Gently penalize overspeeding beyond target
+                    overspeed_ratio = (aligned_speed - target_speed) / target_speed
+                    linear_speed_reward = max(0.0, 5.0 - 2.5 * overspeed_ratio)
+
+                # Drift penalty for sideways sliding
                 perp_vel = actual_vel_2d - aligned_speed * unit_cmd_dir
                 drift_penalty = 0.35 * float(np.dot(perp_vel, perp_vel))
             else:
+                aligned_speed = 0.0
                 linear_speed_reward = 0.0
                 drift_penalty = 1.0 * (linear_vel[0]**2 + linear_vel[1]**2)
 
-            # 2. TURNING SPEED MAXIMIZATION
+            # -----------------------------------------------------------------
+            # 2. TIME-ACCUMULATING STILLNESS PENALTY
+            # -----------------------------------------------------------------
             actual_yaw_rate = angular_vel[2]
+            is_moving = abs(aligned_speed) > 0.03 or abs(actual_yaw_rate) > 0.08
 
+            if (cmd_norm > 0.05 or abs(cmd_yaw) > 0.05) and not is_moving:
+                # Increment counter every step the robot stalls under command
+                self.consecutive_still_steps[idx] += 1
+            else:
+                # Reset counter as soon as robot makes a valid move!
+                self.consecutive_still_steps[idx] = 0
+
+            # Penalty grows by 0.05 every step frozen, capped at -3.5 per step
+            still_penalty = min(3.5, 0.05 * self.consecutive_still_steps[idx])
+
+            # -----------------------------------------------------------------
+            # 3. TURNING & POSTURE PENALTIES
+            # -----------------------------------------------------------------
             if abs(cmd_yaw) > 0.05:
                 turning_reward = 1.5 * (actual_yaw_rate * cmd_yaw)
             else:
                 turning_reward = -0.6 * (actual_yaw_rate ** 2)
 
-            # 3. SMALL ALIVE BONUS + AGE-RAMPED STILLNESS PENALTY
+            alive_reward = 0.20
-            alive_reward = 0.5
-
-            age_ratio = min(1.0, self.step_count / max(1, self.curriculum_episode_limit))
-            still_penalty = 0.0
-            if (cmd_norm > 0.1 or abs(cmd_yaw) > 0.1) and (abs(linear_vel[0]) < 0.02 and abs(actual_yaw_rate) < 0.05):
-                still_penalty = 0.35 + 0.85 * age_ratio
-
-            # 4. POSTURE & STABILITY PENALTIES
-            height_penalty : float
-            min_valid_h = self.target_height * 0.90  # 90% threshold
 
+            # Height drop penalty
+            min_valid_h = self.target_height * 0.90
             if pos[2] < min_valid_h:
-                # Normalized drop below the 90% mark
                 drop = (min_valid_h - pos[2]) / self.target_height
-                # Linear + quadratic penalty that rapidly outweighs the +0.50 alive bonus
                 height_penalty = 4.0 * drop + 20.0 * (drop ** 2)
             else:
-                # Zero penalty inside the valid 90% - 110% zone!
                 height_penalty = 0.0
-            stability_penalty = 1.5 * (roll**2 + pitch**2)
+
+            stability_penalty = 8 * (roll**2 + pitch**2)
 
             control_delta = np.abs(current_actions[idx] - previous_actions[idx])
             large_delta_mask = control_delta > 0.12

ml/train.py

@@ -12,14 +12,25 @@ class MilestoneCheckpointCallback(BaseCallback):
     """
     Saves a model checkpoint the FIRST time total_timesteps 
     crosses every multiple of step_interval (e.g., 100,000).
+    Saves inside the matching PPO_X subfolder as created by TensorBoard.
     """
     def __init__(self, save_path: str, name_prefix: str = "ppo_jackbot", step_interval: int = 100_000, verbose: int = 1):
         super().__init__(verbose)
-        self.save_path = save_path
+        self.base_save_path = save_path
+        self.run_save_path = save_path
         self.name_prefix = name_prefix
         self.step_interval = step_interval
         self.last_milestone = 0
-        os.makedirs(self.save_path, exist_ok=True)
+
+    def _on_training_start(self) -> None:
+        """Executed right before training loop starts. Resolves TensorBoard's run folder name (e.g. PPO_1)."""
+        if self.logger and self.logger.dir:
+            run_folder_name = Path(self.logger.dir).name  # Extracts "PPO_1", "PPO_2", etc.
+            self.run_save_path = os.path.join(self.base_save_path, run_folder_name)
+        else:
+            self.run_save_path = self.base_save_path
+
+        os.makedirs(self.run_save_path, exist_ok=True)
 
     def _on_step(self) -> bool:
         current_milestone = self.num_timesteps // self.step_interval
@@ -29,7 +40,7 @@ class MilestoneCheckpointCallback(BaseCallback):
             milestone_step = current_milestone * self.step_interval
             
             save_file = os.path.join(
-                self.save_path, 
+                self.run_save_path, 
                 f"{self.name_prefix}_{milestone_step}_steps.zip"
             )
             self.model.save(save_file)
@@ -39,6 +50,7 @@ class MilestoneCheckpointCallback(BaseCallback):
                 
         return True
 
+
 class JackBotMetricsCallback(BaseCallback):
     """
     Tracks the best current alive-robot performance for the most recent rollout,
@@ -53,11 +65,9 @@ class JackBotMetricsCallback(BaseCallback):
         self.best_alive_reward = -float('inf')
 
     def _on_step(self) -> bool:
-        """Required by SB3 BaseCallback; no-op here because the rollout summary is emitted at rollout end."""
         return True
 
     def _on_rollout_end(self) -> bool:
-        """Executed right before PPO outputs the log table to console."""
         try:
             vec_env = self.training_env
             alive_metrics = vec_env.env_method("get_current_robot_metrics")
@@ -90,7 +100,6 @@ class JackBotMetricsCallback(BaseCallback):
             self.logger.record("custom/best_alive_survival_steps", float(self.best_alive_survival_steps))
             self.logger.record("custom/best_alive_reward", float(self.best_alive_reward) if np.isfinite(self.best_alive_reward) else 0.0)
 
-            # Backward-compatible aliases so old dashboards keep a stable field name.
             self.logger.record("custom/max_speed_mps", float(self.best_alive_speed))
             self.logger.record("custom/max_yaw_rate_rads", float(self.best_alive_yaw_rate))
             self.logger.record("custom/max_distance_from_start_m", float(self.best_alive_distance))
@@ -101,6 +110,7 @@ class JackBotMetricsCallback(BaseCallback):
 
         return True
 
+
 def parse_args():
     parser = argparse.ArgumentParser(description="Train a joint-command policy for JackBot.")
     parser.add_argument("--timesteps", type=int, default=500_000, help="Total training timesteps")
@@ -117,7 +127,7 @@ def parse_args():
 def make_env(robot_spacing, start_pose, use_gui, rank, seed=0):
     def _init():
         env = JackBotEnv(
-            use_gui=use_gui if rank == 0 else False,  # Only rank 0 gets GUI if requested
+            use_gui=use_gui if rank == 0 else False,
             random_command=True,
             robot_spacing=robot_spacing,
             start_pose=start_pose,
@@ -142,7 +152,6 @@ def train(
     except ImportError as exc:
         raise ImportError("stable-baselines3 is required. Install with: pip install stable-baselines3") from exc
 
-    # Create multi-process vector environment
     if num_workers > 1:
         env_fns = [
             make_env(robot_spacing, start_pose, use_gui, rank=i, seed=seed)
@@ -179,24 +188,24 @@ def train(
     device = resolve_device(device)
 
     policy_kwargs = dict(
-    log_std_init=-1.5,  # Sets initial std ~ 0.22 instead of 1.0
+        log_std_init=-1.5,
         net_arch=dict(pi=[256, 256], vf=[256, 256])
-)
+    )
 
     model = PPO(
         "MlpPolicy",
         env,
         verbose=1,
         seed=seed,
-        learning_rate=1.5e-4,      # Cut LR in half (from 3e-4) to smooth out updates
+        learning_rate=1.5e-4,
-        n_steps=256,               # Larger rollout buffer per env for stable gradients
+        n_steps=256,
-        batch_size=256,            # Larger minibatches reduce noise
+        batch_size=256,
-        n_epochs=10,               # Number of epoch updates per rollout
+        n_epochs=10,
-        gamma=0.99,                # Discount factor
+        gamma=0.99,
-        gae_lambda=0.95,           # GAE smoothing
+        gae_lambda=0.95,
-        clip_range=0.2,            # Standard PPO clipping
+        clip_range=0.2,
-        target_kl=0.03,           # EARLY STOPPING: Halts policy update if KL > 0.015!
+        target_kl=0.03,
-        ent_coef=0.03,            # Entropy coefficient to encourage exploration
+        ent_coef=0.03,
         vf_coef=0.5,
         max_grad_norm=0.5,
         device=device,
@@ -204,6 +213,7 @@ def train(
         tensorboard_log=str(Path(__file__).resolve().parent / "tensorboard"),
     )
 
+    # Base folder where model runs will be stored
     save_dir = str(Path(model_path).parent)
     model_prefix = Path(model_path).stem
 
@@ -217,8 +227,19 @@ def train(
 
     model.learn(total_timesteps=total_timesteps, callback=[milestone_cb, metrics_callback])
 
-    Path(model_path).parent.mkdir(parents=True, exist_ok=True)
+    # Save the final model inside the matching PPO_X directory as well
-    model.save(model_path)
+    if model.logger and model.logger.dir:
+        run_folder_name = Path(model.logger.dir).name
+        final_dir = Path(model_path).parent / run_folder_name
+    else:
+        final_dir = Path(model_path).parent
+
+    final_dir.mkdir(parents=True, exist_ok=True)
+    final_save_path = final_dir / f"{model_prefix}_final.zip"
+    model.save(str(final_save_path))
+    if model.verbose > 0:
+        print(f"[Training Complete] Saved final model to -> {final_save_path}")
+
     env.close()