diff --git a/Robot.py b/Robot.py index 43807a2..5805571 100644 --- a/Robot.py +++ b/Robot.py @@ -265,10 +265,8 @@ class Robot: self.set_joint_angles(target_rad) def apply_rl_action(self, action: np.ndarray) -> None: - """Applies absolute action targets directly for Direct RL.""" action = np.asarray(action, dtype=np.float32) - scaled_action = np.clip(action, -1.0, 1.0) * (np.pi / 2.0) - new_rad = dt.RadArray(data=scaled_action.reshape(self.current_rad.data.shape)) + new_rad = dt.RadArray(data=action.reshape(self.current_rad.data.shape)) self.set_joint_angles(new_rad) def apply_rl_action_delta(self, action: np.ndarray) -> None: diff --git a/ml/env.py b/ml/env.py index 8688067..ed7e974 100644 --- a/ml/env.py +++ b/ml/env.py @@ -81,9 +81,11 @@ class JackBotEnv(gym.Env): self.action_space = spaces.Box(-1.0, 1.0, shape=(action_dim,), dtype=np.float32) self.observation_space = spaces.Box(-np.inf, np.inf, shape=(obs_dim,), dtype=np.float32) + self.min_joint_limits, self.max_joint_limits = self.sim._get_urdf_joint_limits() self.command = np.zeros(3, dtype=np.float32) # [vx, vy, omega] self.last_action = np.zeros(action_dim, dtype=np.float32) + self.last_last_action = np.zeros(action_dim, dtype=np.float32) self.target_height = 0.122 self.collapse_height_fraction = 0.55 self.tilt_failure_rad = 0.9 @@ -143,7 +145,7 @@ class JackBotEnv(gym.Env): self.last_reward_components = {} self.episode_reward_components_sum = defaultdict(float) - spawn_pos = [0.0, 0.0, 0.20] + spawn_pos = [0.0, 0.0, 0.15] spawn_orn = [0.0, 0.0, 0.0, 1.0] # 1. Reset base pose and velocities @@ -155,7 +157,9 @@ class JackBotEnv(gym.Env): if self.use_gui: self.sim.set_robot_color([1.0, 1.0, 1.0, 1.0]) - self.last_action = np.zeros(self.action_space.shape[0], dtype=np.float32) + action_dim = self.action_space.shape[0] + self.last_action = np.zeros(action_dim, dtype=np.float32) + self.last_last_action = np.zeros(action_dim, dtype=np.float32) self.max_distance_from_start = 0.0 self.max_survival_steps = 0 @@ -170,8 +174,13 @@ class JackBotEnv(gym.Env): self.start_position = list(pos) # Drop settlement - self.target_height = self.sim.settle_and_measure_height(steps=200, fallback_height=0.122) + self.target_height = self.sim.settle_and_measure_height( + target_angles=self.robot.current_rad, + steps=300, + fallback_height=0.122 + ) self.default_joint_angles = self.sim.get_robot_joint_angles() + self.default_action = np.clip(self.default_joint_angles / (np.pi / 2.0), -1.0, 1.0) if self.use_gui: self.hud.reset() @@ -183,9 +192,11 @@ class JackBotEnv(gym.Env): return self.robot.get_observation(command=self.command) def step(self, action: np.ndarray) -> Tuple[np.ndarray, float, bool, bool, Dict[str, Any]]: + previous_action = self.last_action.copy() self.step_count += 1 self.total_steps += 1 - previous_action = self.last_action.copy() + + self.last_last_action = self.last_action.copy() self.last_action = action.copy() # Command resampling @@ -200,9 +211,14 @@ class JackBotEnv(gym.Env): # Mirror main.py input resolution logic: update robot_state and vector_dirmov directly 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)] - + # Delegate execution tick to Robot instance - self.robot.tick(action=action) + target_angles = np.where( + action < 0.0, + self.default_joint_angles + action * (self.default_joint_angles - self.min_joint_limits), + self.default_joint_angles + action * (self.max_joint_limits - self.default_joint_angles) + ) + self.robot.tick(action=target_angles) if self.robot_mode != "kinematics" and self.step_count % 60 == 0: random_force = np.random.uniform(-2.0, 2.0, size=2) @@ -265,9 +281,8 @@ class JackBotEnv(gym.Env): def _update_curriculum(self): self._curriculum_advanced = False - forced_forward = (self.curriculum_phase < CurriculumPhase.FORWARD) and (self.step_count >= 2500) - if self.curriculum_phase < CurriculumPhase.FORWARD and (self._phase_progress_ready(CurriculumPhase.FORWARD) or forced_forward): + if self.curriculum_phase < CurriculumPhase.FORWARD and (self._phase_progress_ready(CurriculumPhase.FORWARD)): self.curriculum_phase = CurriculumPhase.FORWARD self._curriculum_advanced = True elif self.curriculum_phase < CurriculumPhase.TURN_AND_DIRECTION and self._phase_progress_ready(CurriculumPhase.TURN_AND_DIRECTION): @@ -289,12 +304,19 @@ class JackBotEnv(gym.Env): cmd_norm = math.hypot(cmd_vx, cmd_vy) raw_speed = math.hypot(linear_vel[0], linear_vel[1]) + # Forgiving zones: ignore noise below 0.04 m/s and 0.05 rad/s filtered_vx, filtered_vy = (linear_vel[0], linear_vel[1]) if raw_speed >= 0.04 else (0.0, 0.0) filtered_speed = raw_speed if raw_speed >= 0.04 else 0.0 - raw_yaw_rate = abs(angular_vel[2]) - filtered_yaw_rate = angular_vel[2] if raw_yaw_rate >= 0.05 else 0.0 + filtered_yaw_rate = raw_yaw_rate if raw_yaw_rate >= 0.05 else 0.0 + # --- 1. FIXED JITTER PENALTY --- + # Scaled way down (0.005) and capped so it can never dominate the reward + action_accel = action - 2.0 * self.last_action + self.last_last_action + raw_jitter = float(np.mean(np.square(action_accel))) + jitter_penalty = -0.005 * min(raw_jitter, 10.0) + + # --- Base Components --- height_error = pos[2] - self.target_height r_height = math.exp(-150.0 * (height_error ** 2)) r_stability = math.exp(-25.0 * (roll**2 + pitch**2)) @@ -302,12 +324,20 @@ class JackBotEnv(gym.Env): r_smoothness = math.exp(-0.1 * np.mean(np.square(action - previous_action))) r_lin_vel, r_ang_vel, stillness_penalty = 0.0, 0.0, 0.0 + stand_penalty = 0.0 + # --- 2. COMMAND IS ZERO: STANDING MODE --- if cmd_norm < 0.05 and abs(cmd_yaw) < 0.05: w_height, w_stability, w_pose, w_smoothness = 0.35, 0.35, 0.20, 0.10 - total_reward = (w_height * r_height) + (w_stability * r_stability) + (w_pose * r_pose) + (w_smoothness * r_smoothness) + base_reward = (w_height * r_height) + (w_stability * r_stability) + (w_pose * r_pose) + (w_smoothness * r_smoothness) + + # Soft quadratic penalty on filtered speed (gives a forgiving gap near 0) + stand_penalty = -0.5 * filtered_speed - 0.1 * filtered_yaw_rate + total_reward = base_reward + stand_penalty + + # --- 3. COMMAND IS NON-ZERO: WALKING MODE --- else: - is_moving = (filtered_speed > 0.0) or (abs(filtered_yaw_rate) > 0.0) + is_moving = (filtered_speed > 0.0) or (filtered_yaw_rate > 0.0) target_vx, target_vy = cmd_vx * self.max_robot_speed, cmd_vy * self.max_robot_speed target_speed = math.hypot(target_vx, target_vy) @@ -320,7 +350,7 @@ class JackBotEnv(gym.Env): if target_speed > 0.08 and raw_speed < 0.03: r_lin_vel = 0.0 - stillness_penalty = -0.25 + stillness_penalty = -0.1 # Softened from -0.25 w_lin_vel, w_ang_vel, w_height, w_stability, w_smoothness = 0.55, 0.15, 0.10, 0.12, 0.08 total_reward = ( @@ -328,7 +358,9 @@ class JackBotEnv(gym.Env): + (w_stability * r_stability) + (w_smoothness * r_smoothness) + stillness_penalty ) - final_reward = float(total_reward / 10.0) + step_reward = float(total_reward / 10.0) + alive_bonus = 0.01 + final_reward = max(0.0, step_reward + jitter_penalty + alive_bonus) self.last_reward_components = { "height": float(r_height), @@ -337,13 +369,15 @@ class JackBotEnv(gym.Env): "smoothness": float(r_smoothness), "lin_vel": float(r_lin_vel), "ang_vel": float(r_ang_vel), + "jitter_penalty": float(jitter_penalty), + "stand_penalty": float(stand_penalty), "total": final_reward, } for k, v in self.last_reward_components.items(): self.episode_reward_components_sum[k] += v return final_reward - + def get_reward_component_averages(self) -> Dict[str, float]: steps = max(1, self.step_count) return {k: v / steps for k, v in self.episode_reward_components_sum.items()} @@ -377,7 +411,7 @@ class JackBotEnv(gym.Env): ) def _update_robot_failure(self): - if self.is_failed or self.step_count < 15: + if self.is_failed: return position, (roll, pitch, _) = self.sim.get_robot_pose_and_rpy() diff --git a/ml/run_eval_training.py b/ml/run_eval_training.py index f1afc53..f1990f2 100644 --- a/ml/run_eval_training.py +++ b/ml/run_eval_training.py @@ -4,6 +4,9 @@ ml/run_eval_training.py - Benchmark reward system across all curriculum phases. import time import numpy as np from ml.env import JackBotEnv, CurriculumPhase +import sys +from pathlib import Path +sys.path.append(str(Path(__file__).resolve().parent.parent)) def evaluate_kinematics(episode_length: int = 1000): env = JackBotEnv( diff --git a/ml/run_train.py b/ml/run_train.py index 4043494..0c6e2ca 100644 --- a/ml/run_train.py +++ b/ml/run_train.py @@ -110,7 +110,7 @@ def main(): eval_env, best_model_save_path=best_model_path, log_path="ml/logs/results", - eval_freq=max(1, 10_000 // args.num_workers), + eval_freq=max(1, 20_000 // args.num_workers), deterministic=True, render=False, ) diff --git a/simulation.py b/simulation.py index 5735f7a..5ec2149 100644 --- a/simulation.py +++ b/simulation.py @@ -82,7 +82,7 @@ class Simulation: jointIndex=joint_index, controlMode=p.POSITION_CONTROL, targetPosition=float(target_angle), - force=500, + force=30, physicsClientId=self.physics_client ) @@ -116,7 +116,7 @@ class Simulation: p.resetBasePositionAndOrientation(self.robot_id, pos, orn, physicsClientId=self.physics_client) p.resetBaseVelocity(self.robot_id, linearVelocity=lin_v, angularVelocity=ang_v, physicsClientId=self.physics_client) - + # --- TELEMETRY (GETTERS) --- def get_robot_pose(self) -> Tuple[List[float], List[float]]: @@ -141,6 +141,23 @@ class Simulation: joint_states = p.getJointStates(self.robot_id, self.revolute_joints, physicsClientId=self.physics_client) return np.array([state[0] for state in joint_states], dtype=np.float32) + def _get_urdf_joint_limits(self) -> Tuple[np.ndarray, np.ndarray]: + """Dynamically reads lower and upper limits for all revolute joints from PyBullet.""" + lower_limits = [] + upper_limits = [] + + # Iterate through joints in PyBullet + for j_idx in range(p.getNumJoints(self.robot_id, physicsClientId=self.physics_client)): + info = p.getJointInfo(self.robot_id, j_idx, physicsClientId=self.physics_client) + joint_type = info[2] + + # Only collect limits for revolute joints + if joint_type == p.JOINT_REVOLUTE: + lower_limits.append(info[8]) # Index 8 = jointLowerLimit + upper_limits.append(info[9]) # Index 9 = jointUpperLimit + + return np.array(lower_limits, dtype=np.float32), np.array(upper_limits, dtype=np.float32) + # --- SIMULATION LIFECYCLE CONTROLS --- def step(self) -> None: @@ -153,8 +170,13 @@ class Simulation: for j in range(num_joints): p.changeVisualShape(self.robot_id, j, rgbaColor=rgba, physicsClientId=self.physics_client) - def settle_and_measure_height(self, steps: int = 200, fallback_height: float = 0.122) -> float: + def settle_and_measure_height( + self, target_angles: Optional[np.ndarray] = None, steps: int = 100, fallback_height: float = 0.122 + ) -> float: + """Settles the robot into the ground while actively holding target joint angles.""" for _ in range(steps): + if target_angles is not None: + self.set_robot_joint_angles(target_angles) self.step() pos, _ = self.get_robot_pose() return pos[2] if pos[2] > 0.0 else fallback_height