diff --git a/Robot.py b/Robot.py index 841b460..43807a2 100644 --- a/Robot.py +++ b/Robot.py @@ -1,12 +1,9 @@ """ -Robot.py - Unified Robot Class for JackBot -Handles state, kinematics, backends (Hardware/Simulation), and motion execution. +Robot.py - Central Robot Control, Kinematics, State Machine & Hardware Abstraction """ - -from typing import Protocol, Optional, Union +from typing import Protocol, Optional, Union, Tuple, List import numpy as np import math -import pybullet as p from states import STATE_REGISTRY from states.State import State @@ -15,33 +12,51 @@ import kinematics as kin import robot_init as ri from config import cfg, BackendType -# Import communications and simulation modules from simulation import Simulation from EspCommunication import ESP32Communication from ArduinoCommunication import ArduinoCommunication class RobotBackend(Protocol): - """Abstraction layer for hardware vs simulation output.""" - def send_angles(self, rad_array: dt.RadArray) -> None: - ... - def step_simulation(self) -> None: - ... - def cleanup(self) -> None: - ... + """Protocol defining hardware abstraction for both Simulation and Hardware backends.""" + def send_angles(self, rad_array: dt.RadArray) -> None: ... + def step_simulation(self) -> None: ... + def hard_reset_joints(self, target_angles: np.ndarray) -> None: ... + def reset_base(self, position: Optional[List[float]] = None, orientation: Optional[List[float]] = None) -> None: ... + def get_joint_angles(self) -> np.ndarray: ... + def get_base_pose_and_rpy(self) -> Tuple[List[float], Tuple[float, float, float]]: ... + def get_base_velocity(self) -> Tuple[List[float], List[float]]: ... + def cleanup(self) -> None: ... class HardwareBackend: - """Backend for physical ESP32 or Arduino robot.""" + """Backend for physical ESP32 or Arduino microcontrollers.""" def __init__(self, comm_channel): self.comm_channel = comm_channel + self.internal_angles = np.zeros(18, dtype=np.float32) def send_angles(self, rad_array: dt.RadArray) -> None: + self.internal_angles = rad_array.data.flatten().copy() if self.comm_channel: self.comm_channel.send_motion(rad_array) def step_simulation(self) -> None: - pass # Physical hardware steps in real-time + pass + + def hard_reset_joints(self, target_angles: np.ndarray) -> None: + self.internal_angles = target_angles.flatten().copy() + + def reset_base(self, position: Optional[List[float]] = None, orientation: Optional[List[float]] = None) -> None: + pass + + def get_joint_angles(self) -> np.ndarray: + return self.internal_angles.copy() + + def get_base_pose_and_rpy(self) -> Tuple[List[float], Tuple[float, float, float]]: + return [0.0, 0.0, 0.122], (0.0, 0.0, 0.0) + + def get_base_velocity(self) -> Tuple[List[float], List[float]]: + return [0.0, 0.0, 0.0], [0.0, 0.0, 0.0] def cleanup(self) -> None: if self.comm_channel and hasattr(self.comm_channel, 'close'): @@ -49,52 +64,62 @@ class HardwareBackend: class PyBulletBackend: - """Backend for PyBullet simulation execution.""" - def __init__(self, sim_instance, body_id: Optional[int] = None): + """Backend mapping Robot operations directly to PyBullet simulation engine.""" + def __init__(self, sim_instance: Simulation): self.sim = sim_instance - self.body_id = body_id def send_angles(self, rad_array: dt.RadArray) -> None: - 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()) + if self.sim: + self.sim.set_robot_joint_angles(rad_array) def step_simulation(self) -> None: if self.sim: self.sim.step() + def hard_reset_joints(self, target_angles: np.ndarray) -> None: + if self.sim: + self.sim.hard_reset_joint_angles(target_angles) + + def reset_base(self, position: Optional[List[float]] = None, orientation: Optional[List[float]] = None) -> None: + if self.sim: + self.sim.reset_robot_base(pos=position, orn=orientation) + + def get_joint_angles(self) -> np.ndarray: + return self.sim.get_robot_joint_angles() if self.sim else np.zeros(18, dtype=np.float32) + + def get_base_pose_and_rpy(self) -> Tuple[List[float], Tuple[float, float, float]]: + return self.sim.get_robot_pose_and_rpy() if self.sim else ([0, 0, 0], (0, 0, 0)) + + def get_base_velocity(self) -> Tuple[List[float], List[float]]: + return self.sim.get_robot_velocity() if self.sim else ([0, 0, 0], [0, 0, 0]) + def cleanup(self) -> None: - if self.sim and hasattr(self.sim, 'disconnect'): + if self.sim: self.sim.disconnect() class Robot: """ - Encapsulates a single JackBot hexapod instance. - Maintains joint states, leg positions, kinematics, and backend control. + Unified JackBot Class. + Coordinates joint memory, IK solvers, procedural tripods, and backend communication. """ def __init__( self, backend_type: Union[BackendType, RobotBackend] = BackendType.SIMULATION, start_pose: str = "init_deg", - urdf_path: str = cfg.urdf_path + urdf_path: str = cfg.urdf_path, + mode: str = "kinematics" # "kinematics", "residual", or "direct" ): self.urdf_path = urdf_path self.start_pose = start_pose + self.mode = mode - # --- BACKEND FACTORY CREATION --- + # --- BACKEND INSTANTIATION --- if isinstance(backend_type, BackendType): if backend_type == BackendType.SIMULATION: - # Launch PyBullet 3D Simulation GUI - sim_instance = Simulation(urdf_path=self.urdf_path) + sim_instance = Simulation(urdf_path=self.urdf_path, use_gui=True) + sim_instance.load_scene() self.backend: RobotBackend = PyBulletBackend(sim_instance) elif backend_type == BackendType.ESP32: comm = ESP32Communication(ip=cfg.esp32_ip, port=cfg.esp32_port) @@ -105,7 +130,7 @@ class Robot: else: self.backend = backend_type - # Kinematics and position initialization + # Kinematics initialization pose_deg = ri.init_deg if start_pose == "init_deg" else ri.init90_deg self.current_rad: dt.RadArray = pose_deg.to_rad() self.current_pos: dt.PosArray = kin.ikpyForward(self.current_rad) @@ -113,32 +138,20 @@ class Robot: ri.get_center_points() if hasattr(ri, "get_center_points") else ri.center_points ) - # RL configuration - self.action_scale = 0.1 # Joint delta step size (radians) + # RL configuration & Gait state variables + self.action_scale = 0.1 # Radian step scale for RL deltas self.gait_phase = 0.0 - - # Gait / motion variables - self.leg_state = np.array(["step", "drag", "step", "drag", "step", "drag"]) - self.robot_state = "idle" self.vector_dirmov = [0.0, 0.0, 0.0] # [vx, vy, omega] + self.robot_state = "idle" + self.leg_state = np.array(["step", "drag", "step", "drag", "step", "drag"]) - # State machine initialization + # State Machine Initialization self.current_state_key: str = "idle" self.current_state: State = STATE_REGISTRY["idle"] self.current_state.enter(self) - def change_state(self, new_state: State) -> None: - if self.current_state: - self.current_state.exit(self) - self.current_state = new_state - self.current_state.enter(self) - - def update(self) -> None: - if self.current_state: - self.current_state.execute(self) - self.step_sim() - def set_joint_angles(self, target_rad: dt.RadArray) -> None: + """Updates internal Python memory state and sends angles to active backend.""" self.current_rad = target_rad if self.backend: self.backend.send_angles(target_rad) @@ -148,12 +161,18 @@ class Robot: self.backend.step_simulation() def reset_to_init(self) -> None: + """Resets kinematics state and forces instant joint alignment in backend.""" 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() + self.robot_state = "idle" + self.leg_state = np.array(["step", "drag", "step", "drag", "step", "drag"]) + + init_flat = self.current_rad.data.flatten() + if self.backend: + self.backend.hard_reset_joints(init_flat) + self.backend.send_angles(self.current_rad) def compute_ik(self, target_pos: dt.PosArray) -> dt.RadArray: return kin.ikpyInverse(target_pos, initial_rad=self.current_rad) @@ -169,41 +188,51 @@ class Robot: self.current_state = STATE_REGISTRY[next_state_key] self.current_state.enter(self) - def tick(self) -> None: - next_state_key = self.current_state.execute(self) - if next_state_key: - self.transition_to(next_state_key) + def tick(self, action: Optional[np.ndarray] = None) -> None: + """ + Unified control loop tick. + Processes commands through direct RL, residual RL, or State Machine kinematics. + """ + vx, vy, omega = self.vector_dirmov + + if self.mode == "direct": + if action is not None: + self.apply_rl_action(action) + + elif self.mode == "residual": + self.step_kinematic_gait(vx, vy, omega) + 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) + self.step_sim() def step_kinematic_gait(self, vx: float, vy: float, omega: float) -> None: - """Procedural Tripod Gait IK solver for kinematic execution & evaluation.""" + """Procedural Tripod Gait solver.""" 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 + stride_len = 0.045 + step_height = 0.035 center_data = ( - self.center_points.data - if hasattr(self.center_points, 'data') - else np.array(self.center_points) + 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 @@ -216,116 +245,49 @@ class Robot: dy_unit = dy_dir / dir_norm if leg_phase < math.pi: - # Swing Phase (Leg lifted & stepping forward) + # Swing phase (leg lifted in air, moving 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 = lift + else: + # Stance phase (leg on ground, pushing body forward) + 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_positions.append(base_pos + np.array([dx, dy, dz], dtype=np.float32)) 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 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)) + self.set_joint_angles(new_rad) + + def apply_rl_action_delta(self, action: np.ndarray) -> None: + """Applies action deltas on top of joint state for Residual RL.""" action = np.asarray(action, dtype=np.float32) scaled_action = np.clip(action, -1.0, 1.0) * self.action_scale - - # 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() - + current_flat = self.backend.get_joint_angles().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) def get_observation(self, command: Optional[np.ndarray] = None) -> np.ndarray: - """ - Returns observation vector [18 joint angles] + [optional 4 command dimensions]. - Queries SimManager helper if PyBulletBackend is used; falls back to internal state otherwise. - """ - 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 - 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_angles = self.current_rad.data.flatten().astype(np.float32) - else: - joint_angles = self.current_rad.data.flatten().astype(np.float32) - + """Extracts joint positions directly from active backend.""" + joint_angles = self.backend.get_joint_angles().flatten().astype(np.float32) if command is not None: cmd = np.asarray(command, dtype=np.float32).flatten() return np.concatenate([joint_angles, cmd]).astype(np.float32) - - return joint_angles.astype(np.float32) + return joint_angles def cleanup(self) -> None: - if self.backend and hasattr(self.backend, 'cleanup'): + if self.backend: self.backend.cleanup() \ No newline at end of file diff --git a/ml/SimManager.py b/ml/SimManager.py deleted file mode 100644 index 84217c6..0000000 --- a/ml/SimManager.py +++ /dev/null @@ -1,298 +0,0 @@ -""" -ml/SimManager.py - PyBullet Simulation Manager (Single Robot Dedicated) -""" -from typing import List, Tuple, Optional, Union -import pybullet as p -import pybullet_data -import numpy as np -import DataTypes as dt - - -class SimManager: - """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.plane: Optional[int] = None - self.joint_indices: List[int] = [] - self.foot_indices: List[int] = [] - - 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 - - flags = p.GUI if self.use_gui else p.DIRECT - self.physics_client = p.connect(flags) - - p.setAdditionalSearchPath(pybullet_data.getDataPath()) - p.setGravity(0, 0, -9.81, physicsClientId=self.physics_client) - - if self.use_gui: - 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, - spacing: float = 0.0, - position_func=None - ) -> Tuple[int, List[int], List[List[int]]]: - """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) - - # 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] - - # 3. Spawn single robot body - self.robot_id = p.loadURDF(urdf_path, spawn_pos, physicsClientId=self.physics_client) - - # 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) - - # 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(j_indices, radflat): - p.setJointMotorControl2( - bodyIndex=bid, - jointIndex=joint_index, - controlMode=p.POSITION_CONTROL, - targetPosition=float(target_angle), - force=500, - physicsClientId=self.physics_client - ) - - 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 visualizer.""" - if self.physics_client is not None and p.isConnected(self.physics_client): - p.configureDebugVisualizer( - p.COV_ENABLE_RENDERING, - 1 if enabled else 0, - physicsClientId=self.physics_client - ) - - 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 - - def get_contact_points( - self, - bodyA: int = -1, - bodyB: int = -1, - linkIndexA: int = -1, - linkIndexB: int = -1, - ): - """Wrapper around pybullet.getContactPoints bound to this simulation client.""" - kwargs = {"physicsClientId": self.physics_client} - if bodyA != -1: - kwargs["bodyA"] = bodyA - if bodyB != -1: - kwargs["bodyB"] = bodyB - if linkIndexA != -1: - kwargs["linkIndexA"] = linkIndexA - if linkIndexB != -1: - kwargs["linkIndexB"] = linkIndexB - - return p.getContactPoints(**kwargs) - - # 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_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 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( - bid, pos, orn, physicsClientId=self.physics_client - ) - p.resetBaseVelocity( - bid, linearVelocity=lin_v, angularVelocity=ang_v, - physicsClientId=self.physics_client - ) - - 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).""" - 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: 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_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(bid, physicsClientId=self.physics_client) - for j_idx in range(num_joints): - 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 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 - - if color is None: - color = [1.0, 1.0, 1.0, 1.0] - - 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(bid, j, rgbaColor=color, physicsClientId=self.physics_client) - - 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_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() - height = self.measure_robot_height() - return height if height > 0.0 else fallback_height - - def apply_external_force( - self, - force: Union[List[float], np.ndarray], - body_id: Optional[int] = None, - link_index: int = -1, - position: Union[List[float], np.ndarray] = (0.0, 0.0, 0.0), - frame: int = p.WORLD_FRAME, - ) -> None: - """Applies external force vector to target link (defaults to base link).""" - bid = self._resolve_body_id(body_id) - p.applyExternalForce( - objectUniqueId=bid, - linkIndex=link_index, - forceObj=list(force), - posObj=list(position), - flags=frame, - physicsClientId=self.physics_client, - ) \ No newline at end of file diff --git a/ml/__init__.py b/ml/__init__.py index 8f6caba..ebf19aa 100644 --- a/ml/__init__.py +++ b/ml/__init__.py @@ -1,4 +1,4 @@ -from .env import JackBotEnv, CurriculumCallback +from .env import JackBotEnv from .model import ActorCritic -__all__ = ["JackBotEnv", "CurriculumCallback", "ActorCritic"] +__all__ = ["JackBotEnv", "ActorCritic"] diff --git a/ml/callbacks.py b/ml/callbacks.py new file mode 100644 index 0000000..4bb9ea2 --- /dev/null +++ b/ml/callbacks.py @@ -0,0 +1,44 @@ +""" +ml/callbacks.py - Stable-Baselines3 Custom Callbacks for Logging & Curriculum Advancement +""" +import os +import numpy as np +from stable_baselines3.common.callbacks import BaseCallback + + +class RewardLoggerCallback(BaseCallback): + """Logs individual reward component averages to TensorBoard.""" + + def __init__(self, verbose: int = 0): + super().__init__(verbose) + + def _on_step(self) -> bool: + # Pull component averages from the environment vector + for env_idx, env in enumerate(self.training_env.envs): + if hasattr(env, "get_reward_component_averages"): + averages = env.get_reward_component_averages() + for key, val in averages.items(): + self.logger.record(f"reward_components/{key}", val) + return True + + +class CurriculumCallback(BaseCallback): + """Monitors evaluation performance and automatically manages curriculum progression.""" + + def __init__(self, eval_freq: int = 10000, verbose: int = 1): + super().__init__(verbose) + self.eval_freq = eval_freq + + def _on_step(self) -> bool: + if self.n_calls % self.eval_freq == 0: + for env in self.training_env.envs: + if hasattr(env, "get_current_robot_metrics"): + metrics = env.get_current_robot_metrics() + if metrics: + phase = metrics[0].get("phase_name", "UNKNOWN") + dist = metrics[0].get("distance_from_start", 0.0) + self.logger.record("curriculum/phase_idx", phase) + self.logger.record("curriculum/max_distance", dist) + if self.verbose > 0: + print(f"[CurriculumCallback] Step {self.num_timesteps}: Current Phase = {phase}, Max Dist = {dist:.2f}m") + return True \ No newline at end of file diff --git a/ml/env.py b/ml/env.py index f995074..8688067 100644 --- a/ml/env.py +++ b/ml/env.py @@ -10,14 +10,12 @@ from collections import defaultdict import gymnasium as gym from gymnasium import spaces import numpy as np -from stable_baselines3.common.callbacks import BaseCallback from config import cfg +from simulation import Simulation from Robot import Robot, PyBulletBackend -from ml.SimManager import SimManager from ml.MetricsOverlay import MetricsHUD -# Color Palette RGBA for Terminated/Failed Robots COLOR_FAILED = [0.3, 0.3, 0.3, 0.6] @@ -30,7 +28,7 @@ class CurriculumPhase(IntEnum): class JackBotEnv(gym.Env): - """Gymnasium environment wrapping a single JackBot hexapod.""" + """Gymnasium environment wrapping JackBot hexapod simulation.""" def __init__( self, @@ -38,7 +36,7 @@ class JackBotEnv(gym.Env): random_command: bool = True, max_episode_steps: int = 3000, urdf_path: str = cfg.urdf_path, - robot_mode: str = "direct", + robot_mode: str = "direct", # "direct", "residual", or "kinematics" ): super().__init__() self.robot_mode = robot_mode @@ -52,7 +50,6 @@ class JackBotEnv(gym.Env): self.total_steps = 0 self.cumulative_reward = 0.0 self.robot_reward = 0.0 - self.consecutive_still_steps = 0 self.is_failed = False self.episode_count = 0 @@ -62,94 +59,55 @@ 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) - self.max_cmd_hold_steps = int(6.0 * self.control_freq) # 360 steps (6s) + self.min_cmd_hold_steps = int(2.0 * self.control_freq) + self.max_cmd_hold_steps = int(6.0 * self.control_freq) self.next_cmd_resample_step = 0 - self.initial_stand_steps = 120 # Mandatory 2s standing window at episode reset + self.initial_stand_steps = 120 - # Initialize Simulation Manager - self.sim_manager = SimManager(use_gui=self.use_gui) - self.sim_manager.connect() + # --- INITIALIZE PYBULLET SIMULATION ENGINE --- + self.sim = Simulation(urdf_path=self.urdf_path, use_gui=self.use_gui) + self.plane_id, self.pb_robot, self.revolute_joints = self.sim.load_scene() - # Connect physics world & load single robot - self.plane, pb_robots, robot_joint_indices = self.sim_manager.load_scene( - self.urdf_path, 0.0, self._robot_base_position - ) - self.joint_indices = robot_joint_indices[0] - - # Instantiate Robot Python wrapper - self.robot = Robot( - backend_type=PyBulletBackend(self.sim_manager), - urdf_path=self.urdf_path - ) + # --- INITIALIZE ROBOT WITH BACKEND --- + self.backend = PyBulletBackend(self.sim) + self.robot = Robot(backend_type=self.backend, urdf_path=self.urdf_path, mode=self.robot_mode) action_dim = 18 - obs_dim = 18 + 4 + obs_dim = 18 + 3 # 18 Joint Angles + 3 Command Inputs [vx, vy, omega] 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.command = np.zeros(4, dtype=np.float32) + self.command = np.zeros(3, dtype=np.float32) # [vx, vy, omega] self.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 - self.foot_link_indices = self._find_foot_link_indices() self.start_position = [0.0, 0.0, 0.0] self.max_distance_from_start = 0.0 self.max_survival_steps = 0 self.default_joint_angles = np.zeros(18, dtype=np.float32) - - # Curriculum Initialization via Enum + self.curriculum_phase = CurriculumPhase.STAND_ONLY - self.curriculum_stage_requirements = { - CurriculumPhase.FORWARD: { - "survival_steps": 300, - "min_avg_height_ratio": 0.88, - "max_avg_roll_pitch": 0.18, - }, - CurriculumPhase.TURN_AND_DIRECTION: { - "survival_steps": 500, - "min_forward_distance": 2.5, - "max_lateral_drift": 0.8, - "min_avg_height_ratio": 0.85, - "stability_roll_pitch": 0.25, - }, - CurriculumPhase.OMNI_DIRECTION: { - "survival_steps": 600, - "min_distance": 5.0, - "min_avg_height_ratio": 0.85, - "stability_roll_pitch": 0.25, - }, - CurriculumPhase.FULL_COMMAND: { - "survival_steps": 750, - "min_distance": 8.0, - "min_avg_height_ratio": 0.85, - "stability_roll_pitch": 0.20, - }, + CurriculumPhase.FORWARD: {"survival_steps": 300, "min_avg_height_ratio": 0.88, "max_avg_roll_pitch": 0.18}, + CurriculumPhase.TURN_AND_DIRECTION: {"survival_steps": 500, "min_forward_distance": 2.5, "max_lateral_drift": 0.8, "min_avg_height_ratio": 0.85, "stability_roll_pitch": 0.25}, + CurriculumPhase.OMNI_DIRECTION: {"survival_steps": 600, "min_distance": 5.0, "min_avg_height_ratio": 0.85, "stability_roll_pitch": 0.25}, + CurriculumPhase.FULL_COMMAND: {"survival_steps": 750, "min_distance": 8.0, "min_avg_height_ratio": 0.85, "stability_roll_pitch": 0.20}, } - self.hud = MetricsHUD(physics_client_id=self.sim_manager.physics_client) + self.hud = MetricsHUD(physics_client_id=self.sim.physics_client) 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.2] - - def _find_foot_link_indices(self) -> list: - return self.sim_manager.get_foot_link_indices(self.pb_robot) - def sample_command(self) -> np.ndarray: + """Samples a command vector [vx, vy, omega] based on active curriculum phase.""" phase = self.curriculum_phase - - stand_probabilities = { + stand_probs = { CurriculumPhase.STAND_ONLY: 1.0, CurriculumPhase.FORWARD: 0.25, CurriculumPhase.TURN_AND_DIRECTION: 0.20, @@ -157,28 +115,19 @@ class JackBotEnv(gym.Env): CurriculumPhase.FULL_COMMAND: 0.15, } - if np.random.random() < stand_probabilities.get(phase, 0.15): - return np.zeros(4, dtype=np.float32) + if np.random.random() < stand_probs.get(phase, 0.15): + return np.zeros(3, dtype=np.float32) if phase == CurriculumPhase.FORWARD: - vx = np.random.uniform(0.15, 0.50) - vy, vz, omega = 0.0, 0.0, 0.0 + vx, vy, omega = np.random.uniform(0.15, 0.50), 0.0, 0.0 elif phase == CurriculumPhase.TURN_AND_DIRECTION: - vx = np.random.uniform(-0.8, 0.8) - vy, vz = 0.0, 0.0 - omega = np.random.uniform(-0.8, 0.8) + vx, vy, omega = np.random.uniform(-0.8, 0.8), 0.0, np.random.uniform(-0.8, 0.8) elif phase == CurriculumPhase.OMNI_DIRECTION: - vx = np.random.uniform(-0.8, 0.8) - vy = np.random.uniform(-0.5, 0.5) - vz = 0.0 - omega = np.random.uniform(-0.8, 0.8) + vx, vy, omega = np.random.uniform(-0.8, 0.8), np.random.uniform(-0.5, 0.5), np.random.uniform(-0.8, 0.8) else: - vx = np.random.uniform(-1.0, 1.0) - vy = np.random.uniform(-1.0, 1.0) - vz = 0.0 - omega = np.random.uniform(-1.0, 1.0) + vx, vy, omega = np.random.uniform(-1.0, 1.0), np.random.uniform(-1.0, 1.0), np.random.uniform(-1.0, 1.0) - return np.array([vx, vy, vz, omega], dtype=np.float32) + return np.array([vx, vy, omega], dtype=np.float32) def reset(self, seed: Optional[int] = None, options: Optional[Dict[str, Any]] = None): super().reset(seed=seed) @@ -191,24 +140,21 @@ class JackBotEnv(gym.Env): self.episode_height_sum = 0.0 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_pos = [0.0, 0.0, 0.20] spawn_orn = [0.0, 0.0, 0.0, 1.0] - self.sim_manager.reset_robot_base(self.pb_robot, spawn_pos, spawn_orn) + # 1. Reset base pose and velocities + self.sim.reset_robot_base(spawn_pos, spawn_orn) + + # 2. Reset internal kinematics & hard reset joints in PyBullet 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]) + self.sim.set_robot_color([1.0, 1.0, 1.0, 1.0]) - self.consecutive_still_steps = 0 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 @@ -217,21 +163,15 @@ class JackBotEnv(gym.Env): self.curriculum_phase = CurriculumPhase.STAND_ONLY self._first_reset = False - self.command = np.zeros(4, dtype=np.float32) + self.command = np.zeros(3, dtype=np.float32) self.next_cmd_resample_step = self.initial_stand_steps - pos, _ = self.sim_manager.get_robot_pose(self.pb_robot) - self.start_position = [float(pos[0]), float(pos[1]), float(pos[2])] + pos, _ = self.sim.get_robot_pose() + self.start_position = list(pos) - # Fixed single-robot settlement call - self.target_height = self.sim_manager.settle_and_measure_height( - steps=200, fallback_height=0.122 - ) - - self.default_joint_angles = np.array( - self.sim_manager.get_robot_joint_angles(self.pb_robot, self.joint_indices), - dtype=np.float32 - ) + # Drop settlement + self.target_height = self.sim.settle_and_measure_height(steps=200, fallback_height=0.122) + self.default_joint_angles = self.sim.get_robot_joint_angles() if self.use_gui: self.hud.reset() @@ -239,36 +179,8 @@ 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 callbacks.""" - if self.is_failed: - return [] - - pos, _ = self.sim_manager.get_robot_pose(self.pb_robot) - linear_vel, angular_vel = self.sim_manager.get_robot_velocity(self.pb_robot) - start_x, start_y, _ = self.start_position - 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])) - - return [{ - "reward": float(self.robot_reward), - "distance_from_start": dist, - "speed": speed, - "yaw_rate": yaw_rate, - "alive": True, - "survival_steps": int(self.step_count), - "phase_name": self.curriculum_phase.name, - }] - def _get_obs(self) -> np.ndarray: - joint_angles = self.sim_manager.get_robot_joint_angles(self.pb_robot, self.joint_indices) - return np.concatenate([joint_angles, self.command]).astype(np.float32) + return self.robot.get_observation(command=self.command) def step(self, action: np.ndarray) -> Tuple[np.ndarray, float, bool, bool, Dict[str, Any]]: self.step_count += 1 @@ -276,28 +188,25 @@ class JackBotEnv(gym.Env): previous_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): 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 - self.robot.step_with_command(command=self.command, action=action, mode=self.robot_mode) + # Extract [vx, vy, omega] + cmd_vx, cmd_vy, cmd_omega = self.command - if self.robot_mode != "kinematics_only" and self.step_count % 60 == 0: + # 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) + + if self.robot_mode != "kinematics" 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 - if self.use_gui and self.step_count % render_freq != 0: - self.sim_manager.set_rendering(False) - - self.sim_manager.step() - - if self.use_gui and self.step_count % render_freq == 0: - self.sim_manager.set_rendering(True) + self.sim.apply_external_force(force=[random_force[0], random_force[1], 0.0]) self._update_robot_failure() self._update_distance_metrics() @@ -305,25 +214,21 @@ class JackBotEnv(gym.Env): obs = self._get_obs() reward = self._compute_reward(action, previous_action) - + self.cumulative_reward += reward self.robot_reward += reward terminated = self.is_failed truncated = self.step_count >= self.max_episode_steps - - info = { - "reward_components": self.last_reward_components.copy() - } + info = {"reward_components": self.last_reward_components.copy()} if self.step_count % 120 == 0 and self.use_gui: self._update_hud() - # 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) + pos, _ = self.sim.get_robot_pose() self.episode_height_sum += float(pos[2]) start_x, start_y, _ = self.start_position dist = float(np.linalg.norm(np.array([pos[0] - start_x, pos[1] - start_y], dtype=np.float32))) @@ -336,7 +241,6 @@ class JackBotEnv(gym.Env): req = self.curriculum_stage_requirements[next_phase] survival_ok = self.max_survival_steps >= req["survival_steps"] - 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) @@ -346,10 +250,8 @@ class JackBotEnv(gym.Env): required_min_avg_height = self.target_height * req.get("min_avg_height_ratio", 0.85) height_ok = avg_height >= required_min_avg_height - pos, _ = self.sim_manager.get_robot_pose(self.pb_robot) - start_x, start_y, _ = self.start_position - dx = pos[0] - start_x - dy = pos[1] - start_y + pos, _ = self.sim.get_robot_pose() + dx, dy = pos[0] - self.start_position[0], pos[1] - self.start_position[1] dist_2d = math.hypot(dx, dy) distance_ok = True @@ -358,12 +260,9 @@ class JackBotEnv(gym.Env): elif "min_distance" in req: distance_ok = dist_2d >= req["min_distance"] - drift_ok = True - if "max_lateral_drift" in req: - drift_ok = abs(dy) <= req["max_lateral_drift"] - + drift_ok = abs(dy) <= req["max_lateral_drift"] if "max_lateral_drift" in req else True return survival_ok and height_ok and stability_ok and distance_ok and drift_ok - + def _update_curriculum(self): self._curriculum_advanced = False forced_forward = (self.curriculum_phase < CurriculumPhase.FORWARD) and (self.step_count >= 2500) @@ -371,174 +270,117 @@ class JackBotEnv(gym.Env): if self.curriculum_phase < CurriculumPhase.FORWARD and (self._phase_progress_ready(CurriculumPhase.FORWARD) or forced_forward): self.curriculum_phase = CurriculumPhase.FORWARD self._curriculum_advanced = True - reason = "FORCED (2500 steps)" if forced_forward else "MET" - print(f"[Curriculum] Phase {self.curriculum_phase.name} unlocked [{reason}] at step {self.step_count}") - elif self.curriculum_phase < CurriculumPhase.TURN_AND_DIRECTION and self._phase_progress_ready(CurriculumPhase.TURN_AND_DIRECTION): self.curriculum_phase = CurriculumPhase.TURN_AND_DIRECTION self._curriculum_advanced = True - print(f"[Curriculum] Phase {self.curriculum_phase.name} unlocked at total step {self.total_steps}") - elif self.curriculum_phase < CurriculumPhase.OMNI_DIRECTION and self._phase_progress_ready(CurriculumPhase.OMNI_DIRECTION): self.curriculum_phase = CurriculumPhase.OMNI_DIRECTION self._curriculum_advanced = True - print(f"[Curriculum] Phase {self.curriculum_phase.name} unlocked at total step {self.total_steps}") - elif self.curriculum_phase < CurriculumPhase.FULL_COMMAND and self._phase_progress_ready(CurriculumPhase.FULL_COMMAND): self.curriculum_phase = CurriculumPhase.FULL_COMMAND self._curriculum_advanced = True - print(f"[Curriculum] Phase {self.curriculum_phase.name} unlocked at total step {self.total_steps}") def _compute_reward(self, action: np.ndarray, previous_action: np.ndarray) -> float: - """ - Calculates task rewards using normalized Exponential Kernels and tracks component terms. - """ - 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( - self.sim_manager.get_robot_joint_angles(self.pb_robot, self.joint_indices), - dtype=np.float32 - ) + pos, (roll, pitch, yaw) = self.sim.get_robot_pose_and_rpy() + linear_vel, angular_vel = self.sim.get_robot_velocity() + current_joints = self.sim.get_robot_joint_angles() - cmd_vx, cmd_vy, _, cmd_yaw = self.command + cmd_vx, cmd_vy, cmd_yaw = self.command cmd_norm = math.hypot(cmd_vx, cmd_vy) - VEL_DEADBAND = 0.04 - YAW_DEADBAND = 0.05 - raw_speed = math.hypot(linear_vel[0], linear_vel[1]) - if raw_speed < VEL_DEADBAND: - filtered_vx, filtered_vy = 0.0, 0.0 - filtered_speed = 0.0 - else: - filtered_vx, filtered_vy = linear_vel[0], linear_vel[1] - filtered_speed = raw_speed + 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]) - if raw_yaw_rate < YAW_DEADBAND: - filtered_yaw_rate = 0.0 - else: - filtered_yaw_rate = angular_vel[2] + filtered_yaw_rate = angular_vel[2] if raw_yaw_rate >= 0.05 else 0.0 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)) + r_pose = math.exp(-2.0 * np.mean(np.square(current_joints - self.default_joint_angles))) + r_smoothness = math.exp(-0.1 * np.mean(np.square(action - previous_action))) - orientation_error = roll**2 + pitch**2 - r_stability = math.exp(-25.0 * orientation_error) - - joint_error = np.mean(np.square(current_joints - self.default_joint_angles)) - r_pose = math.exp(-2.0 * joint_error) - - action_delta = np.mean(np.square(action - previous_action)) - r_smoothness = math.exp(-0.1 * action_delta) - - r_lin_vel = 0.0 - r_ang_vel = 0.0 - stillness_penalty = 0.0 - gated_zero = 0.0 + r_lin_vel, r_ang_vel, stillness_penalty = 0.0, 0.0, 0.0 if cmd_norm < 0.05 and abs(cmd_yaw) < 0.05: - # 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) - + (w_stability * r_stability) - + (w_pose * r_pose) - + (w_smoothness * r_smoothness) - ) + 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) else: - # WALKING / TURNING MODE is_moving = (filtered_speed > 0.0) or (abs(filtered_yaw_rate) > 0.0) - - target_vx = cmd_vx * self.max_robot_speed - target_vy = cmd_vy * self.max_robot_speed + target_vx, target_vy = cmd_vx * self.max_robot_speed, cmd_vy * self.max_robot_speed target_speed = math.hypot(target_vx, target_vy) 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) + r_ang_vel = math.exp(-15.0 * ((filtered_yaw_rate - cmd_yaw)**2)) 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_pose = 0.0 - w_smoothness = 0.08 - + w_lin_vel, w_ang_vel, w_height, w_stability, w_smoothness = 0.55, 0.15, 0.10, 0.12, 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 + (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 ) 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 = { + "height": float(r_height), + "stability": float(r_stability), + "pose": float(r_pose), + "smoothness": float(r_smoothness), + "lin_vel": float(r_lin_vel), + "ang_vel": float(r_ang_vel), + "total": final_reward, } - - self.last_reward_components = comp - for key, val in comp.items(): - self.episode_reward_components_sum[key] += val + 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()} + + def get_current_robot_metrics(self) -> List[Dict[str, Any]]: + linear_vel, angular_vel = self.sim.get_robot_velocity() + speed = float(math.hypot(linear_vel[0], linear_vel[1])) + yaw_rate = float(abs(angular_vel[2])) + + metrics = { + "alive": not self.is_failed, + "phase_name": self.curriculum_phase.name, + "reward": float(self.cumulative_reward), + "speed": speed, + "yaw_rate": yaw_rate, + "distance_from_start": float(self.max_distance_from_start), + "survival_steps": int(self.max_survival_steps), + } + return [metrics] + def _update_hud(self): if not self.use_gui: return - now = time.time() fps = 1.0 / max(now - self.last_time, 1e-5) self.last_time = now - - pos, (roll, pitch, _) = self.sim_manager.get_robot_pose_and_rpy(self.pb_robot) - + pos, (roll, pitch, _) = self.sim.get_robot_pose_and_rpy() self.hud.update( - episode=self.episode_count, - step=self.total_steps, - robot_rewards=[self.robot_reward], - cmd_vel=self.command, - fps=fps, - height=pos[2], - roll_pitch=(math.degrees(roll), math.degrees(pitch)) + episode=self.episode_count, step=self.total_steps, robot_rewards=[self.robot_reward], + cmd_vel=self.command, fps=fps, height=pos[2], roll_pitch=(math.degrees(roll), math.degrees(pitch)) ) def _update_robot_failure(self): - if self.is_failed: + if self.is_failed or self.step_count < 15: return - if self.step_count < 15: - return - - position, (roll, pitch, _) = self.sim_manager.get_robot_pose_and_rpy(self.pb_robot) + position, (roll, pitch, _) = self.sim.get_robot_pose_and_rpy() collapse_threshold = max(0.04, self.collapse_height_fraction * self.target_height) is_tilted = abs(roll) > self.tilt_failure_rad or abs(pitch) > self.tilt_failure_rad is_collapsed = position[2] < collapse_threshold @@ -546,105 +388,7 @@ class JackBotEnv(gym.Env): if is_tilted or is_collapsed: self.is_failed = True if self.use_gui: - self.sim_manager.set_robot_color(self.pb_robot, COLOR_FAILED) + self.sim.set_robot_color(COLOR_FAILED) def close(self): - self.sim_manager.disconnect() - - -class CurriculumCallback(BaseCallback): - """Logs curriculum phase breakdown and best performance metrics to TensorBoard.""" - - def __init__(self, verbose=0): - super().__init__(verbose) - self.best_speed = 0.0 - self.best_yaw_rate = 0.0 - self.best_distance = 0.0 - self.best_survival_steps = 0.0 - self.best_reward = -float('inf') - - def _on_step(self) -> bool: - return True - - def _on_rollout_end(self) -> bool: - try: - vec_env = self.training_env - alive_metrics = vec_env.env_method("get_current_robot_metrics") - - self.best_speed = 0.0 - self.best_yaw_rate = 0.0 - self.best_distance = 0.0 - self.best_survival_steps = 0.0 - self.best_reward = -float('inf') - - phase_counts = { - "stand_only": 0, - "forward": 0, - "turn_and_direction": 0, - "omni_direction": 0, - "full_command": 0, - } - - for worker_res in alive_metrics: - for metrics in worker_res: - if not metrics.get("alive", False): - continue - - phase_key = metrics.get("phase_name", "STAND_ONLY").lower() - if phase_key in phase_counts: - phase_counts[phase_key] += 1 - - if metrics["reward"] > self.best_reward: - self.best_reward = float(metrics["reward"]) - if metrics["speed"] > self.best_speed: - self.best_speed = float(metrics["speed"]) - if metrics["yaw_rate"] > self.best_yaw_rate: - self.best_yaw_rate = float(metrics["yaw_rate"]) - if metrics["distance_from_start"] > self.best_distance: - self.best_distance = float(metrics["distance_from_start"]) - if metrics["survival_steps"] > self.best_survival_steps: - self.best_survival_steps = float(metrics["survival_steps"]) - - for phase_name, count in phase_counts.items(): - self.logger.record(f"phase/{phase_name}", count) - - self.logger.record("custom/best_reward", float(self.best_reward) if np.isfinite(self.best_reward) else 0.0) - self.logger.record("custom/best_survival_steps", float(self.best_survival_steps)) - self.logger.record("custom/best_distance_from_start_m", float(self.best_distance)) - self.logger.record("custom/best_speed_mps", float(self.best_speed)) - self.logger.record("custom/best_yaw_rate_rads", float(self.best_yaw_rate)) - - 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 + self.sim.disconnect() \ No newline at end of file diff --git a/ml/run_eval_training.py b/ml/run_eval_training.py index b3d3e6d..f1afc53 100644 --- a/ml/run_eval_training.py +++ b/ml/run_eval_training.py @@ -6,31 +6,27 @@ 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" + robot_mode="kinematics" ) 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)), + (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)), + (CurriculumPhase.FULL_COMMAND, "FULL OMNI COMBINATION",np.array([0.3, 0.2, 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() @@ -39,9 +35,7 @@ def evaluate_kinematics(episode_length: int = 1000): 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 @@ -49,11 +43,10 @@ def evaluate_kinematics(episode_length: int = 1000): 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"Command Applied: vx={cmd[0]:.2f}, vy={cmd[1]:.2f}, yaw={cmd[2]:.2f}") print(f"Total Episode Reward: {total_reward:.4f}") print("Component Step Averages:") for name, value in comp_averages.items(): diff --git a/ml/run_train.py b/ml/run_train.py index 86acf8d..a7c3b4c 100644 --- a/ml/run_train.py +++ b/ml/run_train.py @@ -15,7 +15,8 @@ from stable_baselines3.common.vec_env import SubprocVecEnv, DummyVecEnv from stable_baselines3.common.callbacks import CheckpointCallback, EvalCallback sys.path.append(str(Path(__file__).resolve().parent.parent)) -from ml.env import JackBotEnv, CurriculumCallback +from ml.env import JackBotEnv +from ml.callbacks import CurriculumCallback # Silence SB3's UserWarning about SubprocVecEnv vs DummyVecEnv warnings.filterwarnings("ignore", category=UserWarning, module="stable_baselines3") diff --git a/simulation.py b/simulation.py index b8a7c2a..5735f7a 100644 --- a/simulation.py +++ b/simulation.py @@ -1,89 +1,200 @@ """ -simulation.py - PyBullet Simulation Interface & Standalone Runner +simulation.py - PyBullet Simulation Interface & Physics Engine +Consolidates scene management, physics queries, motor control, and rendering. """ import time import math +from typing import List, Tuple, Optional, Union import numpy as np import pybullet as p -import pybullet_data # Added missing import +import pybullet_data from config import cfg import DataTypes as dt class Simulation: - def __init__(self, urdf_path: str = cfg.urdf_path): + """Single PyBullet simulation manager providing getters/setters for JackBot.""" + + def __init__(self, urdf_path: str = cfg.urdf_path, use_gui: bool = True): self.urdf_path = urdf_path - - # Connect to PyBullet GUI - self.physicsClient = p.connect(p.GUI) + self.use_gui = use_gui + self.physics_client: Optional[int] = None + self.plane_id: Optional[int] = None + self.robot_id: Optional[int] = None + self.revolute_joints: List[int] = [] + + self.connect() + + def connect(self) -> None: + """Establishes connection to PyBullet GUI or DIRECT mode.""" + if self.physics_client is not None and p.isConnected(self.physics_client): + return + + flags = p.GUI if self.use_gui else p.DIRECT + self.physics_client = p.connect(flags) + p.setAdditionalSearchPath(pybullet_data.getDataPath()) - p.setGravity(0, 0, -9.81) - - # Load plane and robot URDF - self.planeId = p.loadURDF("plane.urdf") - self.robot = p.loadURDF(self.urdf_path, [0, 0, 0.2]) + p.setGravity(0, 0, -9.81, physicsClientId=self.physics_client) - # Discover revolute joint indices dynamically - self.revolute_joints = [] - for j in range(p.getNumJoints(self.robot)): - joint_info = p.getJointInfo(self.robot, j) - if joint_info[2] == p.JOINT_REVOLUTE: - self.revolute_joints.append(j) - - self.set_all_joints_to_90() - p.resetDebugVisualizerCamera( - cameraDistance=1.0, - cameraYaw=50, - cameraPitch=-35, - cameraTargetPosition=[0, 0, 0], - ) - - def set_all_joints_to_90(self): - for joint_index in self.revolute_joints: - p.resetJointState(self.robot, joint_index, math.radians(90)) - - def updatePos(self, current_rad: dt.RadArray): - radflat = current_rad.data.flatten() - for joint_index, target_angle in zip(self.revolute_joints, radflat): - p.setJointMotorControl2( - bodyIndex=self.robot, - jointIndex=joint_index, - controlMode=p.POSITION_CONTROL, - targetPosition=target_angle, - force=500, + if self.use_gui: + 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) + p.resetDebugVisualizerCamera( + cameraDistance=1.0, cameraYaw=50, cameraPitch=-35, cameraTargetPosition=[0, 0, 0], + physicsClientId=self.physics_client ) - def step(self): - p.stepSimulation() + def load_scene(self, spawn_pos: Optional[List[float]] = None) -> Tuple[int, int, List[int]]: + """Loads plane and robot URDF, discovering revolute joint indices dynamically.""" + if spawn_pos is None: + spawn_pos = [0.0, 0.0, 0.20] - def disconnect(self): - if p.isConnected(self.physicsClient): - p.disconnect(self.physicsClient) + self.plane_id = p.loadURDF("plane.urdf", physicsClientId=self.physics_client) + self.robot_id = p.loadURDF(self.urdf_path, spawn_pos, physicsClientId=self.physics_client) - def close(self): - """Cleanup wrapper for Robot backend compatibility.""" + self.revolute_joints = [] + 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.revolute_joints.append(j) + + return self.plane_id, self.robot_id, self.revolute_joints + + # --- ACTUATION (SETTERS) --- + + def set_robot_joint_angles( + self, target_angles: Union[np.ndarray, List[float], dt.RadArray] + ) -> None: + """Applies motor torque to pull joints toward target position angles.""" + if isinstance(target_angles, dt.RadArray): + radflat = target_angles.data.flatten() + elif isinstance(target_angles, np.ndarray): + radflat = target_angles.flatten() + else: + radflat = target_angles + + for joint_index, target_angle in zip(self.revolute_joints, radflat): + p.setJointMotorControl2( + bodyIndex=self.robot_id, + jointIndex=joint_index, + controlMode=p.POSITION_CONTROL, + targetPosition=float(target_angle), + force=500, + 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() + for joint_index, target_angle in zip(self.revolute_joints, radflat): + p.resetJointState( + bodyUniqueId=self.robot_id, + jointIndex=joint_index, + targetValue=float(target_angle), + targetVelocity=0.0, + physicsClientId=self.physics_client + ) + + def reset_robot_base( + self, + pos: Optional[List[float]] = None, + orn: Optional[List[float]] = None, + linear_velocity: Optional[List[float]] = None, + angular_velocity: Optional[List[float]] = None + ) -> None: + """Resets root torso position, orientation quaternion, and clears base velocities.""" + if pos is None: + pos = [0.0, 0.0, 0.20] + if orn is None: + orn = [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(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]]: + pos, orn = p.getBasePositionAndOrientation(self.robot_id, physicsClientId=self.physics_client) + return list(pos), list(orn) + + def get_robot_rpy(self) -> Tuple[float, float, float]: + _, orn = p.getBasePositionAndOrientation(self.robot_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) -> Tuple[List[float], Tuple[float, float, float]]: + pos, orn = p.getBasePositionAndOrientation(self.robot_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) -> Tuple[List[float], List[float]]: + lin_v, ang_v = p.getBaseVelocity(self.robot_id, physicsClientId=self.physics_client) + return list(lin_v), list(ang_v) + + def get_robot_joint_angles(self) -> np.ndarray: + 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) + + # --- SIMULATION LIFECYCLE CONTROLS --- + + def step(self) -> None: + """Advances physics simulation by 1 time step.""" + p.stepSimulation(physicsClientId=self.physics_client) + + def set_robot_color(self, rgba: List[float]) -> None: + num_joints = p.getNumJoints(self.robot_id, physicsClientId=self.physics_client) + p.changeVisualShape(self.robot_id, -1, rgbaColor=rgba, physicsClientId=self.physics_client) + 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: + for _ in range(steps): + self.step() + pos, _ = self.get_robot_pose() + return pos[2] if pos[2] > 0.0 else fallback_height + + def apply_external_force( + self, force: Union[List[float], np.ndarray], link_index: int = -1, position: Tuple[float, float, float] = (0.0, 0.0, 0.0) + ) -> None: + p.applyExternalForce( + objectUniqueId=self.robot_id, + linkIndex=link_index, + forceObj=list(force), + posObj=list(position), + flags=p.WORLD_FRAME, + physicsClientId=self.physics_client, + ) + + 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 + + def close(self) -> None: self.disconnect() if __name__ == "__main__": from Robot import Robot, PyBulletBackend - # 1. Initialize PyBullet simulation environment - sim_instance = Simulation() - backend = PyBulletBackend(sim_instance) + sim_instance = Simulation(use_gui=True) + sim_instance.load_scene() - # 2. Instantiate Robot with simulation backend - robot = Robot(backend_type=backend) + backend = PyBulletBackend(sim_instance) + robot = Robot(backend_type=backend, mode="kinematics") robot.reset_to_init() - # 3. Command forward movement [vx, vy, omega] - robot.vector_dirmov = [1.0, 0.0, 0.0] + # Forward velocity command + robot.vector_dirmov = [0.3, 0.0, 0.0] - # 4. Main test execution loop try: while True: - # Executes state machine logic robot.tick() time.sleep(1.0 / 60.0) except KeyboardInterrupt: