diff --git a/Robot.py b/Robot.py index d8fd537..841b460 100644 --- a/Robot.py +++ b/Robot.py @@ -55,12 +55,16 @@ class PyBulletBackend: self.body_id = body_id def send_angles(self, rad_array: dt.RadArray) -> None: - if self.sim: - # If body_id is set, target that specific robot body - if self.body_id is not None and hasattr(self.sim, 'updatePosForBody'): - self.sim.updatePosForBody(self.body_id, rad_array) - else: - self.sim.updatePos(rad_array) + if not self.sim: + return + + if self.body_id is not None and hasattr(self.sim, 'updatePosForBody'): + self.sim.updatePosForBody(self.body_id, rad_array) + elif hasattr(self.sim, 'updatePos'): + self.sim.updatePos(rad_array) + elif hasattr(self.sim, 'set_robot_joint_angles') and self.body_id is not None: + joint_indices = getattr(self.sim, 'joint_indices', list(range(18))) + self.sim.set_robot_joint_angles(self.body_id, joint_indices, rad_array.data.flatten()) def step_simulation(self) -> None: if self.sim: @@ -111,6 +115,7 @@ class Robot: # RL configuration self.action_scale = 0.1 # Joint delta step size (radians) + self.gait_phase = 0.0 # Gait / motion variables self.leg_state = np.array(["step", "drag", "step", "drag", "step", "drag"]) @@ -146,6 +151,7 @@ class Robot: pose_deg = ri.init_deg if self.start_pose == "init_deg" else ri.init90_deg self.current_rad = pose_deg.to_rad() self.current_pos = kin.ikpyForward(self.current_rad) + self.gait_phase = 0.0 self.set_joint_angles(self.current_rad) self.step_sim() @@ -169,20 +175,129 @@ class Robot: self.transition_to(next_state_key) self.step_sim() + def step_kinematic_gait(self, vx: float, vy: float, omega: float) -> None: + """Procedural Tripod Gait IK solver for kinematic execution & evaluation.""" + cmd_mag = math.hypot(vx, vy) + abs(omega) + if cmd_mag < 0.03: + target_rad = self.compute_ik(self.center_points) + self.set_joint_angles(target_rad) + return + + # Advance gait step cycle phase + self.gait_phase = (self.gait_phase + 0.22) % (2.0 * math.pi) + + stride_len = 0.045 # 4.5 cm maximum stride + step_height = 0.035 # 3.5 cm foot clearance height + + center_data = ( + self.center_points.data + if hasattr(self.center_points, 'data') + else np.array(self.center_points) + ) + + target_positions = [] + for leg_id in range(6): + base_pos = np.array(center_data[leg_id], dtype=np.float32) + + # Tripod leg grouping phase offset (even vs odd leg IDs) + phase_offset = 0.0 if (leg_id % 2 == 0) else math.pi + leg_phase = (self.gait_phase + phase_offset) % (2.0 * math.pi) + + # Directional motion unit vector calculation + lx, ly = base_pos[0], base_pos[1] + rot_dx = -omega * ly + rot_dy = omega * lx + + dx_dir = vx + rot_dx + dy_dir = vy + rot_dy + dir_norm = math.hypot(dx_dir, dy_dir) + 1e-6 + + dx_unit = dx_dir / dir_norm + dy_unit = dy_dir / dir_norm + + if leg_phase < math.pi: + # Swing Phase (Leg lifted & stepping forward) + progress = math.cos(leg_phase) + lift = math.sin(leg_phase) * step_height + dx = -progress * stride_len * dx_unit + dy = -progress * stride_len * dy_unit + dz = lift + else: + # Stance Phase (Leg grounded & propelling torso) + progress = math.cos(leg_phase - math.pi) + dx = progress * stride_len * dx_unit + dy = progress * stride_len * dy_unit + dz = 0.0 + + target_leg_pos = base_pos + np.array([dx, dy, dz], dtype=np.float32) + target_positions.append(target_leg_pos) + + target_pos_array = dt.PosArray(np.array(target_positions)) + target_rad = self.compute_ik(target_pos_array) + self.set_joint_angles(target_rad) + + def step_with_command( + self, + command: np.ndarray, + action: Optional[np.ndarray] = None, + mode: str = "direct" + ) -> None: + """ + Unified motion execution method supporting Direct RL, Residual RL, and Pure Kinematics. + + Args: + command: np.ndarray [vx, vy, vz, omega] from RL environment + action: np.ndarray [18,] RL action deltas from neural network ([-1, 1]) + mode: "kinematics_only" | "residual" | "direct" + """ + # 1. Map RL command [vx, vy, vz, omega] to Robot motion vector [vx, vy, omega] + cmd_vx, cmd_vy, _, cmd_omega = command + self.vector_dirmov = [float(cmd_vx), float(cmd_vy), float(cmd_omega)] + + # 2. Automatically trigger state transitions based on command magnitude + cmd_magnitude = math.hypot(cmd_vx, cmd_vy) + abs(cmd_omega) + if cmd_magnitude > 0.05 and self.current_state_key == "idle": + # Transition to walking state if registered in state machine + target_state = "walk" if "walk" in STATE_REGISTRY else "move" + if target_state in STATE_REGISTRY: + self.transition_to(target_state) + elif cmd_magnitude <= 0.05 and self.current_state_key != "idle": + self.transition_to("idle") + + # 3. Execute according to chosen mode + if mode == "kinematics_only": + self.step_kinematic_gait(cmd_vx, cmd_vy, cmd_omega) + + elif mode == "residual": + if self.current_state_key != "idle" and self.current_state and self.current_state_key in STATE_REGISTRY: + self.current_state.execute(self) + else: + self.step_kinematic_gait(cmd_vx, cmd_vy, cmd_omega) + + if action is not None: + action_flat = np.clip(np.asarray(action, dtype=np.float32), -1.0, 1.0) * self.action_scale + kin_flat = self.current_rad.data.flatten() + final_flat = np.clip(kin_flat + action_flat, -np.pi / 2, np.pi / 2) + self.set_joint_angles(dt.RadArray(data=final_flat.reshape(self.current_rad.data.shape))) + + elif mode == "direct": + if action is not None: + self.apply_rl_action(action) + # --- RL METHODS --- def apply_rl_action(self, action: np.ndarray) -> None: - """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 - current_flat = self.current_rad.data.flatten() - updated_flat = np.clip( - current_flat + scaled_action, - -np.pi / 2, - np.pi / 2 - ) + # Sync with actual PyBullet state if backend supports it + if isinstance(self.backend, PyBulletBackend) and self.backend.sim: + actual_angles = self.backend.sim.get_robot_joint_angles(self.backend.body_id) + current_flat = actual_angles.flatten() + else: + current_flat = self.current_rad.data.flatten() + updated_flat = np.clip(current_flat + scaled_action, -np.pi / 2, np.pi / 2) new_rad = dt.RadArray(data=updated_flat.reshape(self.current_rad.data.shape)) self.set_joint_angles(new_rad) diff --git a/ml/SimManager.py b/ml/SimManager.py index fbd2ed9..84217c6 100644 --- a/ml/SimManager.py +++ b/ml/SimManager.py @@ -1,7 +1,7 @@ """ -ml/SimManager.py - PyBullet Simulation & Multi-Body Manager +ml/SimManager.py - PyBullet Simulation Manager (Single Robot Dedicated) """ -from typing import Dict, List, Tuple, Optional +from typing import List, Tuple, Optional, Union import pybullet as p import pybullet_data import numpy as np @@ -9,14 +9,16 @@ import DataTypes as dt class SimManager: - """Manages PyBullet simulation lifecycle and multi-robot physics.""" + """Manages PyBullet simulation lifecycle for a single JackBot hexapod.""" def __init__(self, use_gui: bool = True): self.use_gui = use_gui self.physics_client = None - self.robot_joints: Dict[int, List[int]] = {} + self.plane: Optional[int] = None + self.joint_indices: List[int] = [] + self.foot_indices: List[int] = [] - def connect(self): + def connect(self) -> None: """Connects to PyBullet and hides side GUI panels.""" if self.physics_client is not None and p.isConnected(self.physics_client): return @@ -28,62 +30,87 @@ class SimManager: p.setGravity(0, 0, -9.81, physicsClientId=self.physics_client) if self.use_gui: - # Disable PyBullet side panel and preview windows p.configureDebugVisualizer(p.COV_ENABLE_GUI, 0, physicsClientId=self.physics_client) p.configureDebugVisualizer(p.COV_ENABLE_SEGMENTATION_MARK_PREVIEW, 0, physicsClientId=self.physics_client) p.configureDebugVisualizer(p.COV_ENABLE_DEPTH_BUFFER_PREVIEW, 0, physicsClientId=self.physics_client) p.configureDebugVisualizer(p.COV_ENABLE_RGB_BUFFER_PREVIEW, 0, physicsClientId=self.physics_client) def load_scene( - self, urdf_path: str, robot_spacing: float, base_pos_fn + self, + urdf_path: str, + spacing: float = 0.0, + position_func=None ) -> Tuple[int, List[int], List[List[int]]]: - """Loads the plane and a single hexapod body into the simulation scene.""" - plane_id = p.loadURDF("plane.urdf", physicsClientId=self.physics_client) - robots = [] - robot_joint_indices = [] - self.robot_joints.clear() + """Loads ground plane and the single robot URDF with high ground friction and force settings.""" + # 1. Load ground plane and set explicit friction + self.plane = p.loadURDF("plane.urdf", physicsClientId=self.physics_client) - base_pos = base_pos_fn(0, robot_spacing) - robot = p.loadURDF( - urdf_path, - basePosition=base_pos, - useFixedBase=False, - physicsClientId=self.physics_client - ) - robots.append(robot) + # 2. Determine starting position (matching simulation.py starting height of 0.20m) + spawn_pos = position_func(0) if position_func is not None else [0.0, 0.0, 0.20] - joint_indices = [ - i for i in range(p.getNumJoints(robot, physicsClientId=self.physics_client)) - if p.getJointInfo(robot, i, physicsClientId=self.physics_client)[2] == p.JOINT_REVOLUTE - ] - robot_joint_indices.append(joint_indices) - self.robot_joints[robot] = joint_indices + # 3. Spawn single robot body + self.robot_id = p.loadURDF(urdf_path, spawn_pos, physicsClientId=self.physics_client) - return plane_id, robots, robot_joint_indices + # 4. Retrieve revolute joint indices + self.joint_indices = [] + for j in range(p.getNumJoints(self.robot_id, physicsClientId=self.physics_client)): + info = p.getJointInfo(self.robot_id, j, physicsClientId=self.physics_client) + if info[2] == p.JOINT_REVOLUTE: + self.joint_indices.append(j) - def updatePosForBody(self, body_id: int, current_rad: dt.RadArray): - """Sets joint motor position targets on individual robot bodies.""" - if body_id not in self.robot_joints: - return - - joint_indices = self.robot_joints[body_id] - radflat = current_rad.data.flatten() + # Return format maintains 100% compatibility with JackBotEnv unpack sequence + return self.plane, [self.robot_id], [self.joint_indices] + + def _resolve_body_id(self, body_id: Optional[int] = None) -> int: + """Internal helper to return the single active robot ID.""" + if body_id is not None: + return body_id + if self.robot_id is not None: + return self.robot_id + raise RuntimeError("No robot loaded in SimManager. Call load_scene() first.") + + def set_robot_joint_angles( + self, + target_angles: Union[np.ndarray, List[float]], + joint_indices: Optional[List[int]] = None, + body_id: Optional[int] = None, + ) -> None: + """Applies position control with force=500 matching simulation.py logic.""" + bid = self._resolve_body_id(body_id) + j_indices = joint_indices if joint_indices is not None else self.joint_indices + radflat = target_angles.flatten() if isinstance(target_angles, np.ndarray) else target_angles - for joint_index, target_angle in zip(joint_indices, radflat): + for joint_index, target_angle in zip(j_indices, radflat): p.setJointMotorControl2( - bodyIndex=body_id, + bodyIndex=bid, jointIndex=joint_index, controlMode=p.POSITION_CONTROL, targetPosition=float(target_angle), - force=250, + force=500, physicsClientId=self.physics_client ) - def step(self): + def updatePosForBody( + self, + body_id_or_rad: Union[int, dt.RadArray], + current_rad: Optional[dt.RadArray] = None + ) -> None: + """Updates robot joint positions using RadArray input.""" + if isinstance(body_id_or_rad, dt.RadArray): + rad_data = body_id_or_rad + bid = self.robot_id + else: + bid = self._resolve_body_id(body_id_or_rad) + rad_data = current_rad + + if rad_data is not None: + self.set_robot_joint_angles(rad_data.data, body_id=bid) + + def step(self) -> None: p.stepSimulation(physicsClientId=self.physics_client) def set_rendering(self, enabled: bool) -> None: - """Toggles PyBullet 3D rendering to speed up simulation.""" + """Toggles PyBullet 3D rendering visualizer.""" if self.physics_client is not None and p.isConnected(self.physics_client): p.configureDebugVisualizer( p.COV_ENABLE_RENDERING, @@ -91,7 +118,7 @@ class SimManager: physicsClientId=self.physics_client ) - def disconnect(self): + def disconnect(self) -> None: if self.physics_client is not None and p.isConnected(self.physics_client): p.disconnect(self.physics_client) self.physics_client = None @@ -116,114 +143,153 @@ class SimManager: return p.getContactPoints(**kwargs) - # --- ROBOT GETTERS AND SETTERS --- + # Add to ml/SimManager.py + def hard_reset_joint_angles(self, target_angles: np.ndarray, body_id: Optional[int] = None) -> None: + bid = self._resolve_body_id(body_id) + radflat = target_angles.flatten() + for joint_index, target_angle in zip(self.joint_indices, radflat): + p.resetJointState( + bodyUniqueId=bid, + jointIndex=joint_index, + targetValue=float(target_angle), + targetVelocity=0.0, + physicsClientId=self.physics_client + ) + + # --- SINGLE ROBOT GETTERS & SETTERS --- def reset_robot_base( self, - body_id: int, - position: List[float], + body_id_or_pos: Union[int, List[float]], + position_or_orn: Optional[List[float]] = None, 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] + """Resets the robot base pose and clears linear/angular velocities.""" + if isinstance(body_id_or_pos, int): + bid = body_id_or_pos + pos = position_or_orn if position_or_orn is not None else [0.0, 0.0, 0.20] + orn = orientation if orientation is not None else [0.0, 0.0, 0.0, 1.0] + else: + bid = self.robot_id + pos = body_id_or_pos + orn = position_or_orn if position_or_orn is not None else [0.0, 0.0, 0.0, 1.0] + + lin_v = linear_velocity if linear_velocity is not None else [0.0, 0.0, 0.0] + ang_v = angular_velocity if angular_velocity is not None else [0.0, 0.0, 0.0] p.resetBasePositionAndOrientation( - body_id, position, orientation, physicsClientId=self.physics_client + bid, pos, orn, physicsClientId=self.physics_client ) p.resetBaseVelocity( - body_id, linearVelocity=linear_velocity, angularVelocity=angular_velocity, + bid, linearVelocity=lin_v, angularVelocity=ang_v, physicsClientId=self.physics_client ) - def get_robot_pose(self, body_id: int) -> Tuple[List[float], List[float]]: + def get_robot_pose(self, body_id: Optional[int] = None) -> 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) + bid = self._resolve_body_id(body_id) + pos, orn = p.getBasePositionAndOrientation(bid, 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) + def get_robot_rpy(self, body_id: Optional[int] = None) -> Tuple[float, float, float]: + """Returns roll, pitch, yaw angles in radians for the robot.""" + bid = self._resolve_body_id(body_id) + _, orn = p.getBasePositionAndOrientation(bid, physicsClientId=self.physics_client) roll, pitch, yaw = p.getEulerFromQuaternion(orn) return float(roll), float(pitch), float(yaw) - def get_foot_link_indices(self, body_id: int) -> list[int]: + def get_robot_pose_and_rpy(self, body_id: Optional[int] = None) -> Tuple[List[float], Tuple[float, float, float]]: + """Returns base position and (roll, pitch, yaw) tuple in radians.""" + bid = self._resolve_body_id(body_id) + pos, orn = p.getBasePositionAndOrientation(bid, 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: Optional[int] = None) -> Tuple[List[float], List[float]]: + """Returns linear velocity (vx, vy, vz) and angular velocity (wx, wy, wz).""" + bid = self._resolve_body_id(body_id) + lin_v, ang_v = p.getBaseVelocity(bid, physicsClientId=self.physics_client) + return list(lin_v), list(ang_v) + + def get_robot_joint_angles( + self, + body_id: Optional[int] = None, + joint_indices: Optional[List[int]] = None + ) -> np.ndarray: + """Returns joint angles as a 1D numpy float32 array.""" + bid = self._resolve_body_id(body_id) + j_indices = joint_indices if joint_indices is not None else self.joint_indices + joint_states = p.getJointStates(bid, j_indices, physicsClientId=self.physics_client) + return np.array([state[0] for state in joint_states], dtype=np.float32) + + def get_foot_link_indices(self, body_id: Optional[int] = None) -> List[int]: """Inspects URDF joint structure to extract link IDs for leg tips and tibias.""" + bid = self._resolve_body_id(body_id) foot_indices = [] - num_joints = p.getNumJoints(body_id, physicsClientId=self.physics_client) + num_joints = p.getNumJoints(bid, physicsClientId=self.physics_client) for j_idx in range(num_joints): - info = p.getJointInfo(body_id, j_idx, physicsClientId=self.physics_client) + info = p.getJointInfo(bid, j_idx, physicsClientId=self.physics_client) link_name = info[12].decode("utf-8") if "tip" in link_name or "tibia" in link_name: foot_indices.append(j_idx) return foot_indices - 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 set_robot_color( + self, + body_id_or_rgba: Union[int, List[float]], + rgba: Optional[List[float]] = None + ) -> None: + """Changes visual color RGBA of base link and all joints.""" + if isinstance(body_id_or_rgba, int): + bid = body_id_or_rgba + color = rgba + else: + bid = self.robot_id + color = body_id_or_rgba - 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) + if color is None: + color = [1.0, 1.0, 1.0, 1.0] - 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) + num_joints = p.getNumJoints(bid, physicsClientId=self.physics_client) + p.changeVisualShape(bid, -1, rgbaColor=color, physicsClientId=self.physics_client) for j in range(num_joints): - p.changeVisualShape(body_id, j, rgbaColor=rgba, physicsClientId=self.physics_client) + p.changeVisualShape(bid, j, rgbaColor=color, 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 measure_robot_height(self, body_id: Optional[int] = None) -> float: + """Gets current Z height of the robot base.""" + bid = self._resolve_body_id(body_id) + pos, _ = p.getBasePositionAndOrientation(bid, physicsClientId=self.physics_client) + return pos[2] - 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): + def settle_and_measure_height( + self, + robot_ids_or_steps: Union[List[int], int] = 200, + steps: int = 200, + fallback_height: float = 0.122 + ) -> float: + """Steps simulation so robot settles onto floor, then calculates target standing height.""" + # Intelligently resolve whether the first argument was passed as robot_ids list or step count + actual_steps = robot_ids_or_steps if isinstance(robot_ids_or_steps, int) else steps + + for _ in range(actual_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 + height = self.measure_robot_height() + return height if height > 0.0 else fallback_height def apply_external_force( self, - body_id: int, - force: list[float] | np.ndarray, + force: Union[List[float], np.ndarray], + body_id: Optional[int] = None, link_index: int = -1, - position: list[float] | np.ndarray = (0.0, 0.0, 0.0), + position: Union[List[float], np.ndarray] = (0.0, 0.0, 0.0), frame: int = p.WORLD_FRAME, - ): - """ - Applies a 3D force vector (in Newtons) to a robot link. - - :param body_id: PyBullet body ID. - :param force: [fx, fy, fz] force vector in Newtons. - :param link_index: Target link index (-1 refers to the base/torso). - :param position: Offset [x, y, z] relative to link center where force is applied. - :param frame: p.WORLD_FRAME (global axes) or p.LINK_FRAME (robot's body axes). - """ + ) -> None: + """Applies external force vector to target link (defaults to base link).""" + bid = self._resolve_body_id(body_id) p.applyExternalForce( - objectUniqueId=body_id, + objectUniqueId=bid, linkIndex=link_index, forceObj=list(force), posObj=list(position), diff --git a/ml/env.py b/ml/env.py index 73c5d3b..f995074 100644 --- a/ml/env.py +++ b/ml/env.py @@ -5,6 +5,7 @@ import time import math from enum import IntEnum from typing import Optional, Tuple, Dict, Any, List +from collections import defaultdict import gymnasium as gym from gymnasium import spaces @@ -37,8 +38,10 @@ class JackBotEnv(gym.Env): random_command: bool = True, max_episode_steps: int = 3000, urdf_path: str = cfg.urdf_path, + robot_mode: str = "direct", ): super().__init__() + self.robot_mode = robot_mode self.use_gui = use_gui self.random_command = random_command self.max_episode_steps = max_episode_steps @@ -59,6 +62,10 @@ class JackBotEnv(gym.Env): self._curriculum_advanced = False self._first_reset = True + # Reward Component Tracking Initialization + self.last_reward_components: Dict[str, float] = {} + self.episode_reward_components_sum: Dict[str, float] = defaultdict(float) + # Dynamic Command Resampling Timing (60 Hz control loop) self.control_freq = 60 self.min_cmd_hold_steps = int(2.0 * self.control_freq) # 120 steps (2s) @@ -74,16 +81,14 @@ class JackBotEnv(gym.Env): self.plane, pb_robots, robot_joint_indices = self.sim_manager.load_scene( self.urdf_path, 0.0, self._robot_base_position ) - self.pb_robot = pb_robots[0] self.joint_indices = robot_joint_indices[0] - # Instantiate Robot Python wrapper (start_pose is managed inside Robot.py) + # Instantiate Robot Python wrapper self.robot = Robot( - backend_type=PyBulletBackend(self.sim_manager, body_id=self.pb_robot), + backend_type=PyBulletBackend(self.sim_manager), urdf_path=self.urdf_path ) - # Action (18 joint deltas) & Observation (18 angles + 4 command dims) action_dim = 18 obs_dim = 18 + 4 @@ -109,7 +114,7 @@ class JackBotEnv(gym.Env): CurriculumPhase.FORWARD: { "survival_steps": 300, "min_avg_height_ratio": 0.88, - "max_avg_roll_pitch": 0.18, # ~10 degrees average + "max_avg_roll_pitch": 0.18, }, CurriculumPhase.TURN_AND_DIRECTION: { "survival_steps": 500, @@ -136,7 +141,7 @@ class JackBotEnv(gym.Env): self.last_time = time.time() def _robot_base_position(self, robot_id: int, spacing: float = 0.0) -> list[float]: - return [0.0, 0.0, 0.13] + return [0.0, 0.0, 0.2] def _find_foot_link_indices(self) -> list: return self.sim_manager.get_foot_link_indices(self.pb_robot) @@ -187,12 +192,19 @@ class JackBotEnv(gym.Env): self.episode_roll_sum = 0.0 self.episode_pitch_sum = 0.0 + # Reset Component Tracking Dictionary + self.last_reward_components = {} + self.episode_reward_components_sum = defaultdict(float) + spawn_pos = self._robot_base_position(0) spawn_orn = [0.0, 0.0, 0.0, 1.0] self.sim_manager.reset_robot_base(self.pb_robot, spawn_pos, spawn_orn) self.robot.reset_to_init() + init_angles = self.robot.current_rad.data.flatten() + self.sim_manager.hard_reset_joint_angles(init_angles, self.pb_robot) + if self.use_gui: self.sim_manager.set_robot_color(self.pb_robot, [1.0, 1.0, 1.0, 1.0]) @@ -211,8 +223,9 @@ class JackBotEnv(gym.Env): pos, _ = self.sim_manager.get_robot_pose(self.pb_robot) self.start_position = [float(pos[0]), float(pos[1]), float(pos[2])] + # Fixed single-robot settlement call self.target_height = self.sim_manager.settle_and_measure_height( - [self.pb_robot], steps=200, fallback_height=0.122 + steps=200, fallback_height=0.122 ) self.default_joint_angles = np.array( @@ -226,8 +239,13 @@ class JackBotEnv(gym.Env): return self._get_obs(), {} + def get_reward_component_averages(self) -> Dict[str, float]: + """Calculates step-averaged scores for each sub-reward component.""" + steps = max(1, self.step_count) + return {k: float(v / steps) for k, v in self.episode_reward_components_sum.items()} + def get_current_robot_metrics(self) -> list: - """Returns metric summary for the callback.""" + """Returns metric summary for callbacks.""" if self.is_failed: return [] @@ -263,17 +281,16 @@ class JackBotEnv(gym.Env): 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 - self.robot.apply_rl_action(action) + self.robot.step_with_command(command=self.command, action=action, mode=self.robot_mode) - if self.step_count % 60 == 0: + if self.robot_mode != "kinematics_only" and self.step_count % 60 == 0: random_force = np.random.uniform(-2.0, 2.0, size=2) self.sim_manager.apply_external_force( body_id=self.pb_robot, force=[random_force[0], random_force[1], 0.0] ) - render_freq = 10 # Only draw 1 in every 10 frames - + render_freq = 10 if self.use_gui and self.step_count % render_freq != 0: self.sim_manager.set_rendering(False) @@ -295,9 +312,15 @@ class JackBotEnv(gym.Env): terminated = self.is_failed truncated = self.step_count >= self.max_episode_steps + info = { + "reward_components": self.last_reward_components.copy() + } + if self.step_count % 120 == 0 and self.use_gui: self._update_hud() - return obs, reward, terminated, truncated, {} + + # Gymnasium standard 5-tuple return + return obs, reward, terminated, truncated, info def _update_distance_metrics(self): pos, _ = self.sim_manager.get_robot_pose(self.pb_robot) @@ -312,22 +335,17 @@ class JackBotEnv(gym.Env): return False req = self.curriculum_stage_requirements[next_phase] - - # 1. Survival Check survival_ok = self.max_survival_steps >= req["survival_steps"] - # 2. Smooth Average Stability Checks (Prevents 1-frame spikes from failing curriculum) avg_roll = self.episode_roll_sum / max(1, self.step_count) avg_pitch = self.episode_pitch_sum / max(1, self.step_count) max_allowed_angle = req.get("max_avg_roll_pitch", 0.20) stability_ok = (avg_roll <= max_allowed_angle) and (avg_pitch <= max_allowed_angle) - # 3. Average Height Check avg_height = self.episode_height_sum / max(1, self.step_count) required_min_avg_height = self.target_height * req.get("min_avg_height_ratio", 0.85) height_ok = avg_height >= required_min_avg_height - # 4. Distance and Drift Checks pos, _ = self.sim_manager.get_robot_pose(self.pb_robot) start_x, start_y, _ = self.start_position dx = pos[0] - start_x @@ -373,10 +391,8 @@ class JackBotEnv(gym.Env): def _compute_reward(self, action: np.ndarray, previous_action: np.ndarray) -> float: """ - Calculates task rewards using normalized Exponential Kernels. - Includes a deadband filter for jittering and zero-reward gating when stationary. + Calculates task rewards using normalized Exponential Kernels and tracks component terms. """ - # 1. Fetch Robot State pos, (roll, pitch, yaw) = self.sim_manager.get_robot_pose_and_rpy(self.pb_robot) linear_vel, angular_vel = self.sim_manager.get_robot_velocity(self.pb_robot) current_joints = np.array( @@ -387,9 +403,8 @@ class JackBotEnv(gym.Env): cmd_vx, cmd_vy, _, cmd_yaw = self.command cmd_norm = math.hypot(cmd_vx, cmd_vy) - # 2. Velocity Deadband Filtering (Ignores jittering & micro-movements) - VEL_DEADBAND = 0.04 # 4 cm/s threshold - YAW_DEADBAND = 0.05 # 0.05 rad/s threshold + VEL_DEADBAND = 0.04 + YAW_DEADBAND = 0.05 raw_speed = math.hypot(linear_vel[0], linear_vel[1]) if raw_speed < VEL_DEADBAND: @@ -405,7 +420,6 @@ class JackBotEnv(gym.Env): else: filtered_yaw_rate = angular_vel[2] - # 3. Posture & Stability Sub-Rewards height_error = pos[2] - self.target_height r_height = math.exp(-150.0 * (height_error ** 2)) @@ -418,13 +432,19 @@ class JackBotEnv(gym.Env): action_delta = np.mean(np.square(action - previous_action)) r_smoothness = math.exp(-0.1 * action_delta) - # 4. Mode Logic + r_lin_vel = 0.0 + r_ang_vel = 0.0 + stillness_penalty = 0.0 + gated_zero = 0.0 + if cmd_norm < 0.05 and abs(cmd_yaw) < 0.05: - # STANDING MODE: Reward clean posture, height, and stability + # STANDING MODE w_height = 0.35 w_stability = 0.35 w_pose = 0.20 w_smoothness = 0.10 + w_lin_vel = 0.0 + w_ang_vel = 0.0 total_reward = ( (w_height * r_height) @@ -436,42 +456,60 @@ class JackBotEnv(gym.Env): # WALKING / TURNING MODE is_moving = (filtered_speed > 0.0) or (abs(filtered_yaw_rate) > 0.0) - # HARD GATE: If commanded to move but standing still/jittering, reward is strictly 0.0 - if not is_moving: - return 0.0 - target_vx = cmd_vx * self.max_robot_speed target_vy = cmd_vy * self.max_robot_speed + target_speed = math.hypot(target_vx, target_vy) - lin_vel_error = (filtered_vx - target_vx)**2 + (filtered_vy - target_vy)**2 - r_lin_vel = math.exp(-25.0 * lin_vel_error) + if not is_moving: + gated_zero = 1.0 + total_reward = 0.0 + w_lin_vel, w_ang_vel, w_height, w_stability, w_pose, w_smoothness = 0, 0, 0, 0, 0, 0 + else: + lin_vel_error = (filtered_vx - target_vx)**2 + (filtered_vy - target_vy)**2 + r_lin_vel = math.exp(-25.0 * lin_vel_error) - ang_vel_error = (filtered_yaw_rate - cmd_yaw)**2 - r_ang_vel = math.exp(-15.0 * ang_vel_error) + ang_vel_error = (filtered_yaw_rate - cmd_yaw)**2 + r_ang_vel = math.exp(-15.0 * ang_vel_error) - # Stillness Check: Commanded to move, but staying virtually still - stillness_penalty = 0.0 - if math.hypot(target_vx, target_vy) > 0.08 and math.hypot(linear_vel[0], linear_vel[1]) < 0.03: - r_lin_vel = 0.0 # Strip velocity credit completely - stillness_penalty = -0.25 + if target_speed > 0.08 and raw_speed < 0.03: + r_lin_vel = 0.0 + stillness_penalty = -0.25 - w_lin_vel = 0.55 - w_ang_vel = 0.15 - w_height = 0.10 - w_stability = 0.12 - w_smoothness = 0.08 + w_lin_vel = 0.55 + w_ang_vel = 0.15 + w_height = 0.10 + w_stability = 0.12 + w_pose = 0.0 + w_smoothness = 0.08 - total_reward = ( - (w_lin_vel * r_lin_vel) - + (w_ang_vel * r_ang_vel) - + (w_height * r_height) - + (w_stability * r_stability) - + (w_smoothness * r_smoothness) - + stillness_penalty - ) + total_reward = ( + (w_lin_vel * r_lin_vel) + + (w_ang_vel * r_ang_vel) + + (w_height * r_height) + + (w_stability * r_stability) + + (w_smoothness * r_smoothness) + + stillness_penalty + ) - # Scaled reward for policy stability - return float(total_reward / 10.0) + final_reward = float(total_reward / 10.0) + + comp = { + "lin_vel": float((w_lin_vel * r_lin_vel) / 10.0), + "ang_vel": float((w_ang_vel * r_ang_vel) / 10.0), + "height": float((w_height * r_height) / 10.0), + "stability": float((w_stability * r_stability) / 10.0), + "pose": float((w_pose * r_pose) / 10.0), + "smoothness": float((w_smoothness * r_smoothness) / 10.0), + "stillness_penalty": float(stillness_penalty / 10.0), + "gated_zero": gated_zero, + "total_step_reward": final_reward, + } + + self.last_reward_components = comp + for key, val in comp.items(): + self.episode_reward_components_sum[key] += val + + return final_reward def _update_hud(self): if not self.use_gui: @@ -579,4 +617,34 @@ class CurriculumCallback(BaseCallback): except Exception: pass + return True + + +class RewardLoggerCallback(BaseCallback): + """ + Logs step-averaged individual reward components to TensorBoard during PPO training. + """ + + def __init__(self, verbose=0): + super().__init__(verbose) + + def _on_step(self) -> bool: + return True + + def _on_rollout_end(self) -> bool: + try: + vec_env = self.training_env + all_comp_averages = vec_env.env_method("get_reward_component_averages") + + if not all_comp_averages: + return True + + keys = all_comp_averages[0].keys() + for key in keys: + avg_val = np.mean([env_comp.get(key, 0.0) for env_comp in all_comp_averages]) + self.logger.record(f"reward_components/{key}", float(avg_val)) + + except Exception: + pass + return True \ No newline at end of file diff --git a/ml/run_eval_training.py b/ml/run_eval_training.py new file mode 100644 index 0000000..b3d3e6d --- /dev/null +++ b/ml/run_eval_training.py @@ -0,0 +1,76 @@ +""" +ml/run_eval_training.py - Benchmark reward system across all curriculum phases. +""" +import time +import numpy as np +from ml.env import JackBotEnv, CurriculumPhase + +def evaluate_kinematics(episode_length: int = 1000): + # Use kinematics_only mode so inverse kinematics generates gait motion from commands + env = JackBotEnv( + use_gui=True, + random_command=False, + max_episode_steps=episode_length, + robot_mode="kinematics_only" + ) + + print("\n" + "=" * 70) + print(" RUNNING MULTI-PHASE REWARD BENCHMARK (KINEMATICS MODE)") + print("=" * 70 + "\n") + + # Define test suite covering every curriculum stage + phase_configs = [ + (CurriculumPhase.STAND_ONLY, "STAND ONLY", np.array([0.0, 0.0, 0.0, 0.0], dtype=np.float32)), + (CurriculumPhase.FORWARD, "FORWARD GAIT", np.array([0.1, 0.0, 0.0, 0.0], dtype=np.float32)), + (CurriculumPhase.TURN_AND_DIRECTION, "FORWARD + YAW TURN", np.array([0.3, 0.0, 0.0, 0.4], dtype=np.float32)), + (CurriculumPhase.OMNI_DIRECTION, "STRIDE LATERAL", np.array([0.2, 0.3, 0.0, 0.0], dtype=np.float32)), + (CurriculumPhase.FULL_COMMAND, "FULL OMNI COMBINATION",np.array([0.3, 0.2, 0.0, 0.3], dtype=np.float32)), + ] + + for phase_enum, label, cmd in phase_configs: + obs, _ = env.reset() + + # Force specific curriculum phase & target command + env.curriculum_phase = phase_enum + env.command = cmd.copy() + + done = False + total_reward = 0.0 + step_count = 0 + + while not done: + # Action array is unused in kinematics_only mode + dummy_action = np.zeros(18, dtype=np.float32) + + obs, reward, terminated, truncated, _ = env.step(dummy_action) + total_reward += reward + step_count += 1 + done = terminated or truncated + + time.sleep(1.0 / 60.0) + + # Retrieve detailed component averages + comp_averages = env.get_reward_component_averages() + + print(f"\n--- Episode Stage: [{phase_enum.name}] ({label}) ---") + print(f"Command Applied: vx={cmd[0]:.2f}, vy={cmd[1]:.2f}, vz={cmd[2]:.2f}, yaw={cmd[3]:.2f}") + print(f"Total Episode Reward: {total_reward:.4f}") + print("Component Step Averages:") + for name, value in comp_averages.items(): + print(f" • {name:<20}: {value:+.5f}") + + metrics = env.get_current_robot_metrics() + dist = metrics[0]["distance_from_start"] if metrics else 0.0 + speed = metrics[0]["speed"] if metrics else 0.0 + avg_reward = total_reward / max(1, step_count) + + print(f" ├─ Average Reward / Step: {avg_reward:.4f}") + print(f" ├─ Distance Travelled: {dist:.2f} m") + print(f" ├─ Actual Avg Speed: {speed:.2f} m/s") + print(f" └─ Steps Survived: {step_count} / {episode_length}") + print("-" * 70) + + env.close() + +if __name__ == "__main__": + evaluate_kinematics() \ No newline at end of file diff --git a/ml/run_train.py b/ml/run_train.py index 66e5155..86acf8d 100644 --- a/ml/run_train.py +++ b/ml/run_train.py @@ -77,7 +77,7 @@ def main(): model = PPO( policy="MlpPolicy", env=vec_env, - learning_rate=3e-4, + learning_rate=1e-4, n_steps=256, batch_size=256, n_epochs=10,