simulation physics fixed

fuck them physics

Robot.py

@@ -1,12 +1,9 @@
 """
-Robot.py - Unified Robot Class for JackBot
+Robot.py - Central Robot Control, Kinematics, State Machine & Hardware Abstraction
-Handles state, kinematics, backends (Hardware/Simulation), and motion execution.
 """
-
+from typing import Protocol, Optional, Union, Tuple, List
-from typing import Protocol, Optional, Union
 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."""
+    """Protocol defining hardware abstraction for both Simulation and Hardware backends."""
-    def send_angles(self, rad_array: dt.RadArray) -> None:
+    def send_angles(self, rad_array: dt.RadArray) -> None: ...
-        ...
+    def step_simulation(self) -> 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 cleanup(self) -> 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."""
+    """Backend mapping Robot operations directly to PyBullet simulation engine."""
-    def __init__(self, sim_instance, body_id: Optional[int] = None):
+    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:
+        if self.sim:
-            return
+            self.sim.set_robot_joint_angles(rad_array)
-
-        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:
             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.
+    Unified JackBot Class.
-    Maintains joint states, leg positions, kinematics, and backend control.
+    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, use_gui=True)
-                sim_instance = Simulation(urdf_path=self.urdf_path)
+                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
+        # RL configuration & Gait state variables
-        self.action_scale = 0.1  # Joint delta step size (radians)
+        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.robot_state = "idle"
-        self.step_sim()
+        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:
+    def tick(self, action: Optional[np.ndarray] = None) -> None:
-        next_state_key = self.current_state.execute(self)
+        """
-        if next_state_key:
+        Unified control loop tick.
-            self.transition_to(next_state_key)
+        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
-        stride_len = 0.045   # 4.5 cm maximum stride
+        step_height = 0.035
-        step_height = 0.035  # 3.5 cm foot clearance height
 
         center_data = (
-            self.center_points.data 
+            self.center_points.data if hasattr(self.center_points, 'data') else np.array(self.center_points)
-            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(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 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
-
+        current_flat = self.backend.get_joint_angles().flatten()
-        # 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)
 
     def get_observation(self, command: Optional[np.ndarray] = None) -> np.ndarray:
-        """
+        """Extracts joint positions directly from active backend."""
-        Returns observation vector [18 joint angles] + [optional 4 command dimensions].
+        joint_angles = self.backend.get_joint_angles().flatten().astype(np.float32)
-        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)
-
         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
-        return joint_angles.astype(np.float32)
 
     def cleanup(self) -> None:
-        if self.backend and hasattr(self.backend, 'cleanup'):
+        if self.backend:
             self.backend.cleanup()

ml/SimManager.py

@@ -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,
-        )

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"]

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

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.min_cmd_hold_steps = int(2.0 * self.control_freq)
-        self.max_cmd_hold_steps = int(6.0 * self.control_freq)  # 360 steps (6s)
+        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
+        # --- INITIALIZE PYBULLET SIMULATION ENGINE ---
-        self.sim_manager = SimManager(use_gui=self.use_gui)
+        self.sim = Simulation(urdf_path=self.urdf_path, use_gui=self.use_gui)
-        self.sim_manager.connect()
+        self.plane_id, self.pb_robot, self.revolute_joints = self.sim.load_scene()
 
-        # Connect physics world & load single robot
+        # --- INITIALIZE ROBOT WITH BACKEND ---
-        self.plane, pb_robots, robot_joint_indices = self.sim_manager.load_scene(
+        self.backend = PyBulletBackend(self.sim)
-            self.urdf_path, 0.0, self._robot_base_position
+        self.robot = Robot(backend_type=self.backend, urdf_path=self.urdf_path, mode=self.robot_mode)
-        )
-        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
-        )
 
         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: {
+            CurriculumPhase.FORWARD: {"survival_steps": 300, "min_avg_height_ratio": 0.88, "max_avg_roll_pitch": 0.18},
-                "survival_steps": 300,
+            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},
-                "min_avg_height_ratio": 0.88,
+            CurriculumPhase.OMNI_DIRECTION: {"survival_steps": 600, "min_distance": 5.0, "min_avg_height_ratio": 0.85, "stability_roll_pitch": 0.25},
-                "max_avg_roll_pitch": 0.18,
+            CurriculumPhase.FULL_COMMAND: {"survival_steps": 750, "min_distance": 8.0, "min_avg_height_ratio": 0.85, "stability_roll_pitch": 0.20},
-            },
-            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_probs = {
-        stand_probabilities = {
             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):
+        if np.random.random() < stand_probs.get(phase, 0.15):
-            return np.zeros(4, dtype=np.float32)
+            return np.zeros(3, dtype=np.float32)
 
         if phase == CurriculumPhase.FORWARD:
-            vx = np.random.uniform(0.15, 0.50)
+            vx, vy, omega = np.random.uniform(0.15, 0.50), 0.0, 0.0
-            vy, vz, omega = 0.0, 0.0, 0.0
         elif phase == CurriculumPhase.TURN_AND_DIRECTION:
-            vx = 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)
-            vy, vz = 0.0, 0.0
-            omega = np.random.uniform(-0.8, 0.8)
         elif phase == CurriculumPhase.OMNI_DIRECTION:
-            vx = 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)
-            vy = np.random.uniform(-0.5, 0.5)
-            vz = 0.0
-            omega = np.random.uniform(-0.8, 0.8)
         else:
-            vx = 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)
-            vy = np.random.uniform(-1.0, 1.0)
-            vz = 0.0
-            omega = 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)
+        pos, _ = self.sim.get_robot_pose()
-        self.start_position = [float(pos[0]), float(pos[1]), float(pos[2])]
+        self.start_position = list(pos)
 
-        # Fixed single-robot settlement call
+        # Drop settlement
-        self.target_height = self.sim_manager.settle_and_measure_height(
+        self.target_height = self.sim.settle_and_measure_height(steps=200, fallback_height=0.122)
-            steps=200, fallback_height=0.122
+        self.default_joint_angles = self.sim.get_robot_joint_angles()
-        )
-
-        self.default_joint_angles = np.array(
-            self.sim_manager.get_robot_joint_angles(self.pb_robot, self.joint_indices), 
-            dtype=np.float32
-        )
 
         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 self.robot.get_observation(command=self.command)
-        return np.concatenate([joint_angles, self.command]).astype(np.float32)
 
     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(
+            self.sim.apply_external_force(force=[random_force[0], random_force[1], 0.0])
-                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._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)
+        pos, _ = self.sim.get_robot_pose()
-        start_x, start_y, _ = self.start_position
+        dx, dy = pos[0] - self.start_position[0], pos[1] - self.start_position[1]
-        dx = pos[0] - start_x
-        dy = pos[1] - start_y
         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
+        drift_ok = abs(dy) <= req["max_lateral_drift"] if "max_lateral_drift" in req else True
-        if "max_lateral_drift" in req:
-            drift_ok = abs(dy) <= req["max_lateral_drift"]
-
         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:
-        """
+        pos, (roll, pitch, yaw) = self.sim.get_robot_pose_and_rpy()
-        Calculates task rewards using normalized Exponential Kernels and tracks component terms.
+        linear_vel, angular_vel = self.sim.get_robot_velocity()
-        """
+        current_joints = self.sim.get_robot_joint_angles()
-        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
-        )
 
-        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 = (linear_vel[0], linear_vel[1]) if raw_speed >= 0.04 else (0.0, 0.0)
-            filtered_vx, filtered_vy = 0.0, 0.0
+        filtered_speed = raw_speed if raw_speed >= 0.04 else 0.0
-            filtered_speed = 0.0
-        else:
-            filtered_vx, filtered_vy = linear_vel[0], linear_vel[1]
-            filtered_speed = raw_speed
 
         raw_yaw_rate = abs(angular_vel[2])
-        if raw_yaw_rate < YAW_DEADBAND:
+        filtered_yaw_rate = angular_vel[2] if raw_yaw_rate >= 0.05 else 0.0
-            filtered_yaw_rate = 0.0
-        else:
-            filtered_yaw_rate = angular_vel[2]
 
         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_lin_vel, r_ang_vel, stillness_penalty = 0.0, 0.0, 0.0
-        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
 
         if cmd_norm < 0.05 and abs(cmd_yaw) < 0.05:
-            # STANDING MODE
+            w_height, w_stability, w_pose, w_smoothness = 0.35, 0.35, 0.20, 0.10
-            w_height = 0.35
+            total_reward = (w_height * r_height) + (w_stability * r_stability) + (w_pose * r_pose) + (w_smoothness * r_smoothness)
-            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)
-            )
         else:
-            # WALKING / TURNING MODE
             is_moving = (filtered_speed > 0.0) or (abs(filtered_yaw_rate) > 0.0)
-
+            target_vx, target_vy = cmd_vx * self.max_robot_speed, cmd_vy * self.max_robot_speed
-            target_vx = cmd_vx * self.max_robot_speed
-            target_vy = 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)
-
+                r_ang_vel = math.exp(-15.0 * ((filtered_yaw_rate - cmd_yaw)**2))
-                ang_vel_error = (filtered_yaw_rate - cmd_yaw)**2
-                r_ang_vel = math.exp(-15.0 * ang_vel_error)
 
                 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_lin_vel, w_ang_vel, w_height, w_stability, w_smoothness = 0.55, 0.15, 0.10, 0.12, 0.08
-                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_lin_vel * r_lin_vel) + (w_ang_vel * r_ang_vel) + (w_height * r_height)
-                    + (w_ang_vel * r_ang_vel)
+                    + (w_stability * r_stability) + (w_smoothness * r_smoothness) + stillness_penalty
-                    + (w_height * r_height)
-                    + (w_stability * r_stability)
-                    + (w_smoothness * r_smoothness)
-                    + stillness_penalty
                 )
 
         final_reward = float(total_reward / 10.0)
 
-        comp = {
+        self.last_reward_components = {
-            "lin_vel": float((w_lin_vel * r_lin_vel) / 10.0),
+            "height": float(r_height),
-            "ang_vel": float((w_ang_vel * r_ang_vel) / 10.0),
+            "stability": float(r_stability),
-            "height": float((w_height * r_height) / 10.0),
+            "pose": float(r_pose),
-            "stability": float((w_stability * r_stability) / 10.0),
+            "smoothness": float(r_smoothness),
-            "pose": float((w_pose * r_pose) / 10.0),
+            "lin_vel": float(r_lin_vel),
-            "smoothness": float((w_smoothness * r_smoothness) / 10.0),
+            "ang_vel": float(r_ang_vel),
-            "stillness_penalty": float(stillness_penalty / 10.0),
+            "total": final_reward,
-            "gated_zero": gated_zero,
-            "total_step_reward": final_reward,
         }
-
+        for k, v in self.last_reward_components.items():
-        self.last_reward_components = comp
+            self.episode_reward_components_sum[k] += v
-        for key, val in comp.items():
-            self.episode_reward_components_sum[key] += val
 
         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.get_robot_pose_and_rpy()
-        pos, (roll, pitch, _) = self.sim_manager.get_robot_pose_and_rpy(self.pb_robot)
-        
         self.hud.update(
-            episode=self.episode_count,
+            episode=self.episode_count, step=self.total_steps, robot_rewards=[self.robot_reward],
-            step=self.total_steps,
+            cmd_vel=self.command, fps=fps, height=pos[2], roll_pitch=(math.degrees(roll), math.degrees(pitch))
-            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:
+        position, (roll, pitch, _) = self.sim.get_robot_pose_and_rpy()
-            return
-
-        position, (roll, pitch, _) = self.sim_manager.get_robot_pose_and_rpy(self.pb_robot)
         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()
+        self.sim.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

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([1.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.5, 0.0, 0.4], 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.5, 0.5, 0.0], 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.3], 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()
         
@@ -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():

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")

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
-        
+        self.use_gui = use_gui
-        # Connect to PyBullet GUI
+        self.physics_client: Optional[int] = None
-        self.physicsClient = p.connect(p.GUI)
+        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)
+        p.setGravity(0, 0, -9.81, physicsClientId=self.physics_client)
-        
-        # Load plane and robot URDF
-        self.planeId = p.loadURDF("plane.urdf")
-        self.robot = p.loadURDF(self.urdf_path, [0, 0, 0.2])
 
-        # Discover revolute joint indices dynamically
+        if self.use_gui:
-        self.revolute_joints = []
+            p.configureDebugVisualizer(p.COV_ENABLE_GUI, 0, physicsClientId=self.physics_client)
-        for j in range(p.getNumJoints(self.robot)):
+            p.configureDebugVisualizer(p.COV_ENABLE_SEGMENTATION_MARK_PREVIEW, 0, physicsClientId=self.physics_client)
-            joint_info = p.getJointInfo(self.robot, j)
+            p.configureDebugVisualizer(p.COV_ENABLE_DEPTH_BUFFER_PREVIEW, 0, physicsClientId=self.physics_client)
-            if joint_info[2] == p.JOINT_REVOLUTE:
+            p.configureDebugVisualizer(p.COV_ENABLE_RGB_BUFFER_PREVIEW, 0, physicsClientId=self.physics_client)
-                self.revolute_joints.append(j)
+            p.resetDebugVisualizerCamera(
-
+                cameraDistance=1.0, cameraYaw=50, cameraPitch=-35, cameraTargetPosition=[0, 0, 0],
-        self.set_all_joints_to_90()
+                physicsClientId=self.physics_client
-        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,
             )
 
-    def step(self):
+    def load_scene(self, spawn_pos: Optional[List[float]] = None) -> Tuple[int, int, List[int]]:
-        p.stepSimulation()
+        """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):
+        self.plane_id = p.loadURDF("plane.urdf", physicsClientId=self.physics_client)
-        if p.isConnected(self.physicsClient):
+        self.robot_id = p.loadURDF(self.urdf_path, spawn_pos, physicsClientId=self.physics_client)
-            p.disconnect(self.physicsClient)
 
-    def close(self):
+        self.revolute_joints = []
-        """Cleanup wrapper for Robot backend compatibility."""
+        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(use_gui=True)
-    sim_instance = Simulation()
+    sim_instance.load_scene()
-    backend = PyBulletBackend(sim_instance)
 
-    # 2. Instantiate Robot with simulation backend
+    backend = PyBulletBackend(sim_instance)
-    robot = Robot(backend_type=backend)
+    robot = Robot(backend_type=backend, mode="kinematics")
     robot.reset_to_init()
 
-    # 3. Command forward movement [vx, vy, omega]
+    # Forward velocity command
-    robot.vector_dirmov = [1.0, 0.0, 0.0]
+    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: