From 402c20dfb5764ccd8ccc74d88988dfa0a8fc25db Mon Sep 17 00:00:00 2001 From: JackM323 Date: Tue, 4 Aug 2026 21:03:53 +0200 Subject: [PATCH] env cleanup rewards adjustment pybullet logic contained in SimManager --- Robot.py | 25 +++--- ml/SimManager.py | 83 +++++++++++++++++++- ml/env.py | 197 +++++++++++++++-------------------------------- ml/train.py | 10 ++- 4 files changed, 162 insertions(+), 153 deletions(-) diff --git a/Robot.py b/Robot.py index b0aa496..d8fd537 100644 --- a/Robot.py +++ b/Robot.py @@ -172,9 +172,7 @@ class Robot: # --- RL METHODS --- def apply_rl_action(self, action: np.ndarray) -> None: - """ - Applies continuous RL action deltas [-1, 1] to current joint angles. - """ + """Applies continuous RL action deltas [-1, 1] to current joint angles.""" action = np.asarray(action, dtype=np.float32) scaled_action = np.clip(action, -1.0, 1.0) * self.action_scale @@ -191,20 +189,19 @@ class Robot: def get_observation(self, command: Optional[np.ndarray] = None) -> np.ndarray: """ Returns observation vector [18 joint angles] + [optional 4 command dimensions]. - Queries PyBullet if backend is PyBulletBackend; otherwise falls back to internal state. + Queries SimManager helper if PyBulletBackend is used; falls back to internal state otherwise. """ - if isinstance(self.backend, PyBulletBackend) and self.backend.sim and hasattr(self.backend.sim, 'physics_client'): - physics_client = self.backend.sim.physics_client + if isinstance(self.backend, PyBulletBackend) and self.backend.sim: body_id = self.backend.body_id if self.backend.body_id is not None else 0 - - # Retrieve joint mapping from SimManager/Simulation if available - if hasattr(self.backend.sim, 'robot_joints') and body_id in self.backend.sim.robot_joints: - joint_indices = self.backend.sim.robot_joints[body_id] + if hasattr(self.backend.sim, 'get_robot_joint_angles'): + joint_angles = self.backend.sim.get_robot_joint_angles(body_id) + elif hasattr(self.backend.sim, 'physics_client'): + physics_client = self.backend.sim.physics_client + joint_indices = self.backend.sim.robot_joints.get(body_id, list(range(18))) if hasattr(self.backend.sim, 'robot_joints') else list(range(18)) + joint_states = p.getJointStates(body_id, joint_indices, physicsClientId=physics_client) + joint_angles = np.array([state[0] for state in joint_states], dtype=np.float32) else: - joint_indices = list(range(18)) - - joint_states = p.getJointStates(body_id, joint_indices, physicsClientId=physics_client) - joint_angles = np.array([state[0] for state in joint_states], dtype=np.float32) + joint_angles = self.current_rad.data.flatten().astype(np.float32) else: joint_angles = self.current_rad.data.flatten().astype(np.float32) diff --git a/ml/SimManager.py b/ml/SimManager.py index c13e52e..3a96701 100644 --- a/ml/SimManager.py +++ b/ml/SimManager.py @@ -1,9 +1,10 @@ """ ml/SimManager.py - PyBullet Simulation & Multi-Body Manager """ -from typing import Dict, List, Tuple +from typing import Dict, List, Tuple, Optional import pybullet as p import pybullet_data +import numpy as np import DataTypes as dt @@ -84,4 +85,82 @@ class SimManager: def disconnect(self): if self.physics_client is not None and p.isConnected(self.physics_client): p.disconnect(self.physics_client) - self.physics_client = None \ No newline at end of file + self.physics_client = None + + # --- ROBOT GETTERS AND SETTERS --- + + def reset_robot_base( + self, + body_id: int, + position: List[float], + orientation: Optional[List[float]] = None, + linear_velocity: Optional[List[float]] = None, + angular_velocity: Optional[List[float]] = None + ) -> None: + """Resets a robot body's base position, orientation, and velocities.""" + if orientation is None: + orientation = [0.0, 0.0, 0.0, 1.0] + if linear_velocity is None: + linear_velocity = [0.0, 0.0, 0.0] + if angular_velocity is None: + angular_velocity = [0.0, 0.0, 0.0] + + p.resetBasePositionAndOrientation( + body_id, position, orientation, physicsClientId=self.physics_client + ) + p.resetBaseVelocity( + body_id, linearVelocity=linear_velocity, angularVelocity=angular_velocity, + physicsClientId=self.physics_client + ) + + def get_robot_pose(self, body_id: int) -> Tuple[List[float], List[float]]: + """Returns base position (x, y, z) and orientation quaternion (x, y, z, w).""" + pos, orn = p.getBasePositionAndOrientation(body_id, physicsClientId=self.physics_client) + return list(pos), list(orn) + + def get_robot_rpy(self, body_id: int) -> Tuple[float, float, float]: + """Returns roll, pitch, yaw angles in radians for the given robot body.""" + _, orn = p.getBasePositionAndOrientation(body_id, physicsClientId=self.physics_client) + roll, pitch, yaw = p.getEulerFromQuaternion(orn) + return float(roll), float(pitch), float(yaw) + + def get_robot_pose_and_rpy(self, body_id: int) -> Tuple[List[float], Tuple[float, float, float]]: + """Returns base position and (roll, pitch, yaw) tuple in radians.""" + pos, orn = p.getBasePositionAndOrientation(body_id, physicsClientId=self.physics_client) + roll, pitch, yaw = p.getEulerFromQuaternion(orn) + return list(pos), (float(roll), float(pitch), float(yaw)) + + def get_robot_velocity(self, body_id: int) -> Tuple[List[float], List[float]]: + """Returns linear velocity (vx, vy, vz) and angular velocity (wx, wy, wz).""" + lin_v, ang_v = p.getBaseVelocity(body_id, physicsClientId=self.physics_client) + return list(lin_v), list(ang_v) + + def get_robot_joint_angles(self, body_id: int, joint_indices: Optional[List[int]] = None) -> np.ndarray: + """Returns joint angles as a 1D numpy array float32 for specified or registered joint indices.""" + if joint_indices is None: + joint_indices = self.robot_joints.get(body_id, list(range(18))) + joint_states = p.getJointStates(body_id, joint_indices, physicsClientId=self.physics_client) + return np.array([state[0] for state in joint_states], dtype=np.float32) + + def set_robot_color(self, body_id: int, rgba: List[float]) -> None: + """Changes visual color RGBA of base link and all joints of the specified robot body.""" + num_joints = p.getNumJoints(body_id, physicsClientId=self.physics_client) + p.changeVisualShape(body_id, -1, rgbaColor=rgba, physicsClientId=self.physics_client) + for j in range(num_joints): + p.changeVisualShape(body_id, j, rgbaColor=rgba, physicsClientId=self.physics_client) + + def measure_robot_heights(self, robot_ids: List[int]) -> List[float]: + """Gets current Z height for all specified robot body IDs.""" + heights = [] + for body_id in robot_ids: + pos, _ = p.getBasePositionAndOrientation(body_id, physicsClientId=self.physics_client) + heights.append(pos[2]) + return heights + + def settle_and_measure_height(self, robot_ids: List[int], steps: int = 200, fallback_height: float = 0.14) -> float: + """Steps simulation for designated steps so robot settles, then calculates target standing height.""" + for _ in range(steps): + self.step() + heights = self.measure_robot_heights(robot_ids) + mean_height = float(np.mean(heights)) if heights else fallback_height + return mean_height if mean_height > 0.0 else fallback_height \ No newline at end of file diff --git a/ml/env.py b/ml/env.py index cacfd4c..62078ba 100644 --- a/ml/env.py +++ b/ml/env.py @@ -8,7 +8,6 @@ from typing import Optional, Tuple, Dict, Any, List import gymnasium as gym from gymnasium import spaces import numpy as np -import pybullet as p from config import cfg from Robot import Robot, PyBulletBackend @@ -79,9 +78,7 @@ class JackBotEnv(gym.Env): self.collapse_height_fraction = 0.55 self.tilt_failure_rad = 0.9 - # Small random exploration pulse to break local optima. - # The bonus is only granted on a rare step and only when the robot - # is still stable enough that a bold move is likely to remain safe. + # Small random exploration pulse settings self.exploration_bonus_prob = 0.03 self.exploration_bonus_interval = 120 self.exploration_bonus_scale = 0.08 @@ -93,7 +90,7 @@ class JackBotEnv(gym.Env): self.curriculum_phase = 0 self.curriculum_episode_limit = 150 self.curriculum_stage_requirements = { - 1: {"survival_steps": 200, "distance": 0.00, "stability_roll_pitch": 0.35}, + 1: {"survival_steps": 1000, "distance": 0.00, "stability_roll_pitch": 0.35}, 2: {"survival_steps": 350, "distance": 10.0, "stability_roll_pitch": 0.30}, 3: {"survival_steps": 550, "distance": 15.00, "stability_roll_pitch": 0.25}, } @@ -103,13 +100,6 @@ class JackBotEnv(gym.Env): self.leader_crown = LeaderCrown(physics_client_id=self.sim_manager.physics_client) self.last_time = time.time() - def _set_robot_color(self, pb_id: int, rgba: List[float]): - """Helper to change the visual color of a robot body and all its links.""" - num_joints = p.getNumJoints(pb_id, physicsClientId=self.sim_manager.physics_client) - p.changeVisualShape(pb_id, -1, rgbaColor=rgba, physicsClientId=self.sim_manager.physics_client) - for j in range(num_joints): - p.changeVisualShape(pb_id, j, rgbaColor=rgba, physicsClientId=self.sim_manager.physics_client) - def _robot_base_position(self, robot_id: int, spacing: float = 0.5) -> list[float]: return [0.0, 0.0, 0.14] @@ -148,20 +138,13 @@ class JackBotEnv(gym.Env): for idx, (pb_id, robot_obj) in enumerate(zip(self.pb_robots, self.robots)): spawn_pos = self._robot_base_position(idx, self.robot_spacing) - spawn_orn = [0, 0, 0, 1] - - p.resetBasePositionAndOrientation( - pb_id, spawn_pos, spawn_orn, physicsClientId=self.sim_manager.physics_client - ) - p.resetBaseVelocity( - pb_id, linearVelocity=[0, 0, 0], angularVelocity=[0, 0, 0], - physicsClientId=self.sim_manager.physics_client - ) + spawn_orn = [0.0, 0.0, 0.0, 1.0] + self.sim_manager.reset_robot_base(pb_id, spawn_pos, spawn_orn) robot_obj.reset_to_init() if self.use_gui: - self._set_robot_color(pb_id, [1.0, 1.0, 1.0, 1.0]) + self.sim_manager.set_robot_color(pb_id, [1.0, 1.0, 1.0, 1.0]) self.last_action = np.zeros(self.action_space.shape[0], dtype=np.float32) self.max_distance_from_start = [0.0] @@ -178,18 +161,14 @@ class JackBotEnv(gym.Env): else: self.commands = np.zeros((1, 4), dtype=np.float32) - for idx, (pb_id, robot_obj) in enumerate(zip(self.pb_robots, self.robots)): - pos, _ = p.getBasePositionAndOrientation( - pb_id, physicsClientId=self.sim_manager.physics_client - ) + for idx, pb_id in enumerate(self.pb_robots): + pos, _ = self.sim_manager.get_robot_pose(pb_id) self.start_positions[idx] = [float(pos[0]), float(pos[1]), float(pos[2])] - for _ in range(200): - self.sim_manager.step() - - self.target_height = self._measure_settled_height() - if self.target_height <= 0.0: - self.target_height = 0.14 + # Settle for 200 steps after dropping in a standing position, then measure target height + self.target_height = self.sim_manager.settle_and_measure_height( + self.pb_robots, steps=200, fallback_height=0.14 + ) if self.use_gui: self.hud.reset() @@ -199,83 +178,63 @@ class JackBotEnv(gym.Env): return self._get_obs(), {} def get_robot_velocities(self) -> list: - """Exposes velocities for the SB3 metrics callback.""" - vels = [] - for pb_id in self.pb_robots: - lin_v, ang_v = p.getBaseVelocity(pb_id, physicsClientId=self.sim_manager.physics_client) - vels.append((lin_v, ang_v)) - return vels + """Exposes velocities for the SB3 metrics callback.""" + vels = [] + for pb_id in self.pb_robots: + lin_v, ang_v = self.sim_manager.get_robot_velocity(pb_id) + vels.append((lin_v, ang_v)) + return vels def get_robot_distance_metrics(self) -> list: - """Exposes distance-from-start metrics for logging without affecting reward.""" - metrics = [] - for idx, pb_id in enumerate(self.pb_robots): - pos, _ = p.getBasePositionAndOrientation( - pb_id, physicsClientId=self.sim_manager.physics_client - ) - start_x, start_y, _ = self.start_positions[idx] - dist = float(np.linalg.norm(np.array([pos[0] - start_x, pos[1] - start_y], dtype=np.float32))) - self.max_distance_from_start[idx] = max(self.max_distance_from_start[idx], dist) - metrics.append((dist, self.max_distance_from_start[idx])) + """Exposes distance-from-start metrics for logging without affecting reward.""" + metrics = [] + for idx, pb_id in enumerate(self.pb_robots): + pos, _ = self.sim_manager.get_robot_pose(pb_id) + start_x, start_y, _ = self.start_positions[idx] + dist = float(np.linalg.norm(np.array([pos[0] - start_x, pos[1] - start_y], dtype=np.float32))) + self.max_distance_from_start[idx] = max(self.max_distance_from_start[idx], dist) + metrics.append((dist, self.max_distance_from_start[idx])) - self.max_survival_steps = max(self.max_survival_steps, self.step_count) - return metrics + self.max_survival_steps = max(self.max_survival_steps, self.step_count) + return metrics def get_current_robot_metrics(self) -> list: - """Returns current per-robot reward, distance, and velocity summaries for alive robots only.""" - metrics = [] - for idx, pb_id in enumerate(self.pb_robots): - if self.failed_robots_mask[idx]: - continue + """Returns current per-robot reward, distance, and velocity summaries for alive robots only.""" + metrics = [] + for idx, pb_id in enumerate(self.pb_robots): + if self.failed_robots_mask[idx]: + continue - pos, _ = p.getBasePositionAndOrientation( - pb_id, physicsClientId=self.sim_manager.physics_client - ) - linear_vel, angular_vel = p.getBaseVelocity( - pb_id, physicsClientId=self.sim_manager.physics_client - ) - start_x, start_y, _ = self.start_positions[idx] - dist = float(np.linalg.norm(np.array([pos[0] - start_x, pos[1] - start_y], dtype=np.float32))) - speed = float(np.linalg.norm(np.array([linear_vel[0], linear_vel[1]], dtype=np.float32))) - yaw_rate = float(abs(angular_vel[2])) - metrics.append({ - "reward": float(self.robot_rewards[idx]), - "distance_from_start": dist, - "speed": speed, - "yaw_rate": yaw_rate, - "alive": True, - "survival_steps": int(self.step_count), - }) + pos, _ = self.sim_manager.get_robot_pose(pb_id) + linear_vel, angular_vel = self.sim_manager.get_robot_velocity(pb_id) + start_x, start_y, _ = self.start_positions[idx] + dist = float(np.linalg.norm(np.array([pos[0] - start_x, pos[1] - start_y], dtype=np.float32))) + speed = float(np.linalg.norm(np.array([linear_vel[0], linear_vel[1]], dtype=np.float32))) + yaw_rate = float(abs(angular_vel[2])) + metrics.append({ + "reward": float(self.robot_rewards[idx]), + "distance_from_start": dist, + "speed": speed, + "yaw_rate": yaw_rate, + "alive": True, + "survival_steps": int(self.step_count), + }) - return metrics + return metrics def get_survival_steps(self) -> int: - """Returns the current survival length for the environment's current episode.""" - return int(self.step_count) + """Returns the current survival length for the environment's current episode.""" + return int(self.step_count) def _get_obs(self) -> np.ndarray: obs_list = [] for idx, (pb_id, joint_indices) in enumerate(zip(self.pb_robots, self.robot_joint_indices)): - joint_states = p.getJointStates( - pb_id, - joint_indices, - physicsClientId=self.sim_manager.physics_client - ) - joint_angles = np.array([state[0] for state in joint_states], dtype=np.float32) + joint_angles = self.sim_manager.get_robot_joint_angles(pb_id, joint_indices) robot_obs = np.concatenate([joint_angles, self.commands[idx]]) obs_list.append(robot_obs) return np.concatenate(obs_list).astype(np.float32) - def _measure_settled_height(self) -> float: - heights = [] - for pb_id in self.pb_robots: - pos, _ = p.getBasePositionAndOrientation( - pb_id, physicsClientId=self.sim_manager.physics_client - ) - heights.append(pos[2]) - return float(np.mean(heights)) if heights else 0.14 - def step(self, action: np.ndarray) -> Tuple[np.ndarray, float, bool, bool, Dict[str, Any]]: self.step_count += 1 self.total_steps += 1 @@ -322,10 +281,7 @@ class JackBotEnv(gym.Env): survival_ok = self.max_survival_steps >= req["survival_steps"] distance_ok = self.max_distance_from_start[0] >= req["distance"] - position, orientation = p.getBasePositionAndOrientation( - self.pb_robots[0], physicsClientId=self.sim_manager.physics_client - ) - roll, pitch, _ = p.getEulerFromQuaternion(orientation) + position, (roll, pitch, _) = self.sim_manager.get_robot_pose_and_rpy(self.pb_robots[0]) stability_ok = abs(roll) <= req["stability_roll_pitch"] and abs(pitch) <= req["stability_roll_pitch"] height_ok = position[2] >= max(0.09, self.target_height * 0.85) @@ -363,78 +319,55 @@ class JackBotEnv(gym.Env): previous_actions = previous_action.reshape(1, 18) for idx, pb_id in enumerate(self.pb_robots): - pos, orientation = p.getBasePositionAndOrientation( - pb_id, physicsClientId=self.sim_manager.physics_client - ) - linear_vel, angular_vel = p.getBaseVelocity( - pb_id, physicsClientId=self.sim_manager.physics_client - ) - roll, pitch, _ = p.getEulerFromQuaternion(orientation) + pos, (roll, pitch, _) = self.sim_manager.get_robot_pose_and_rpy(pb_id) + linear_vel, angular_vel = self.sim_manager.get_robot_velocity(pb_id) command = self.commands[idx] cmd_vx = command[0] cmd_vy = command[1] cmd_yaw = command[3] - # ------------------------------------------------------------- - # 1. LINEAR VECTOR SPEED MAXIMIZATION (Magnitude + Direction) - # ------------------------------------------------------------- + # 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: - # Normalize target direction vector unit_cmd_dir = cmd_dir / cmd_norm actual_vel_2d = np.array([linear_vel[0], linear_vel[1]], dtype=np.float32) - # Speed aligned with target direction (m/s) aligned_speed = float(np.dot(actual_vel_2d, unit_cmd_dir)) - - # Moderate movement reward to encourage directional motion linear_speed_reward = 2.5 * aligned_speed - # Penalize sideways drift (perpendicular velocity to commanded direction) perp_vel = actual_vel_2d - aligned_speed * unit_cmd_dir drift_penalty = 0.35 * float(np.dot(perp_vel, perp_vel)) else: - # If no linear command given, penalize all horizontal movement linear_speed_reward = 0.0 drift_penalty = 1.0 * (linear_vel[0]**2 + linear_vel[1]**2) - # ------------------------------------------------------------- # 2. TURNING SPEED MAXIMIZATION - # ------------------------------------------------------------- - actual_yaw_rate = angular_vel[2] # rad/s in PyBullet Z-axis + actual_yaw_rate = angular_vel[2] if abs(cmd_yaw) > 0.05: - # Reward turning in the commanded direction, but less aggressively turning_reward = 1.5 * (actual_yaw_rate * cmd_yaw) else: - # Penalize unwanted rotation when joystick turn is centered turning_reward = -0.6 * (actual_yaw_rate ** 2) - # ------------------------------------------------------------- # 3. SMALL ALIVE BONUS + AGE-RAMPED STILLNESS PENALTY - # ------------------------------------------------------------- - alive_reward = 0.02 + 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 = 6.0 * ((self.target_height - pos[2]) ** 2) if pos[2] < self.target_height else 0.0 + height_penalty = 10.0 * ((self.target_height - pos[2]) ** 2) if pos[2] < self.target_height else 0.0 stability_penalty = 1.5 * (roll**2 + pitch**2) - - # Penalize only large control jumps; allow continuous low-amplitude motion + print(-height_penalty) control_delta = np.abs(current_actions[idx] - previous_actions[idx]) large_delta_mask = control_delta > 0.12 large_delta_penalty = 0.002 * float(np.sum(np.square(control_delta[large_delta_mask]))) if np.any(large_delta_mask) else 0.0 - # Combined Step Reward r_step = ( alive_reward + linear_speed_reward @@ -465,8 +398,7 @@ class JackBotEnv(gym.Env): rolls = [] pitches = [] for pb_id in self.pb_robots: - pos, orient = p.getBasePositionAndOrientation(pb_id, physicsClientId=self.sim_manager.physics_client) - roll, pitch, _ = p.getEulerFromQuaternion(orient) + pos, (roll, pitch, _) = self.sim_manager.get_robot_pose_and_rpy(pb_id) heights.append(pos[2]) rolls.append(math.degrees(roll)) pitches.append(math.degrees(pitch)) @@ -490,9 +422,7 @@ class JackBotEnv(gym.Env): best_idx = int(np.argmax(self.robot_rewards)) leader_pb_id = self.pb_robots[best_idx] - leader_pos, _ = p.getBasePositionAndOrientation( - leader_pb_id, physicsClientId=self.sim_manager.physics_client - ) + leader_pos, _ = self.sim_manager.get_robot_pose(leader_pb_id) self.leader_crown.update(leader_pos) def _update_robot_failures(self): @@ -501,10 +431,7 @@ class JackBotEnv(gym.Env): if self.failed_robots_mask[idx]: continue - position, orientation = p.getBasePositionAndOrientation( - pb_id, physicsClientId=self.sim_manager.physics_client - ) - roll, pitch, _ = p.getEulerFromQuaternion(orientation) + position, (roll, pitch, _) = self.sim_manager.get_robot_pose_and_rpy(pb_id) collapse_threshold = max(0.06, self.collapse_height_fraction * self.target_height) is_tilted = abs(roll) > self.tilt_failure_rad or abs(pitch) > self.tilt_failure_rad @@ -513,7 +440,7 @@ class JackBotEnv(gym.Env): if is_tilted or is_collapsed: self.failed_robots_mask[idx] = True if self.use_gui: - self._set_robot_color(pb_id, COLOR_FAILED) + self.sim_manager.set_robot_color(pb_id, COLOR_FAILED) def close(self): self.sim_manager.disconnect() \ No newline at end of file diff --git a/ml/train.py b/ml/train.py index 72a14ac..bbddf2e 100644 --- a/ml/train.py +++ b/ml/train.py @@ -178,13 +178,18 @@ def train( device = resolve_device(device) + policy_kwargs = dict( + log_std_init=-1.5, # Sets initial std ~ 0.22 instead of 1.0 + net_arch=dict(pi=[256, 256], vf=[256, 256]) +) + model = PPO( "MlpPolicy", env, verbose=1, seed=seed, - learning_rate=3.5e-4, # Cut LR in half (from 3e-4) to smooth out updates - n_steps=2048, # Larger rollout buffer per env for stable gradients + learning_rate=1.5e-4, # Cut LR in half (from 3e-4) to smooth out updates + n_steps=1024, # Larger rollout buffer per env for stable gradients batch_size=128, # Larger minibatches reduce noise n_epochs=10, # Number of epoch updates per rollout gamma=0.99, # Discount factor @@ -195,6 +200,7 @@ def train( vf_coef=0.5, max_grad_norm=0.5, device=device, + policy_kwargs=policy_kwargs, tensorboard_log=str(Path(__file__).resolve().parent / "tensorboard"), )