added Robot kinematic use for training

HUGE BUG -> SimManager physics broken (at least with Robot kinematics)

Robot.py

@@ -55,12 +55,16 @@ class PyBulletBackend:
         self.body_id = body_id
 
     def send_angles(self, rad_array: dt.RadArray) -> None:
-        if self.sim:
+        if not self.sim:
-            # If body_id is set, target that specific robot body
+            return
-            if self.body_id is not None and hasattr(self.sim, 'updatePosForBody'):
+
-                self.sim.updatePosForBody(self.body_id, rad_array)
+        if self.body_id is not None and hasattr(self.sim, 'updatePosForBody'):
-            else:
+            self.sim.updatePosForBody(self.body_id, rad_array)
-                self.sim.updatePos(rad_array)
+        elif hasattr(self.sim, 'updatePos'):
+            self.sim.updatePos(rad_array)
+        elif hasattr(self.sim, 'set_robot_joint_angles') and self.body_id is not None:
+            joint_indices = getattr(self.sim, 'joint_indices', list(range(18)))
+            self.sim.set_robot_joint_angles(self.body_id, joint_indices, rad_array.data.flatten())
 
     def step_simulation(self) -> None:
         if self.sim:
@@ -111,6 +115,7 @@ class Robot:
 
         # RL configuration
         self.action_scale = 0.1  # Joint delta step size (radians)
+        self.gait_phase = 0.0
 
         # Gait / motion variables
         self.leg_state = np.array(["step", "drag", "step", "drag", "step", "drag"])
@@ -146,6 +151,7 @@ class Robot:
         pose_deg = ri.init_deg if self.start_pose == "init_deg" else ri.init90_deg
         self.current_rad = pose_deg.to_rad()
         self.current_pos = kin.ikpyForward(self.current_rad)
+        self.gait_phase = 0.0
         self.set_joint_angles(self.current_rad)
         self.step_sim()
 
@@ -169,20 +175,129 @@ class Robot:
             self.transition_to(next_state_key)
         self.step_sim()
 
+    def step_kinematic_gait(self, vx: float, vy: float, omega: float) -> None:
+        """Procedural Tripod Gait IK solver for kinematic execution & evaluation."""
+        cmd_mag = math.hypot(vx, vy) + abs(omega)
+        if cmd_mag < 0.03:
+            target_rad = self.compute_ik(self.center_points)
+            self.set_joint_angles(target_rad)
+            return
+
+        # Advance gait step cycle phase
+        self.gait_phase = (self.gait_phase + 0.22) % (2.0 * math.pi)
+
+        stride_len = 0.045   # 4.5 cm maximum stride
+        step_height = 0.035  # 3.5 cm foot clearance height
+
+        center_data = (
+            self.center_points.data 
+            if hasattr(self.center_points, 'data') 
+            else np.array(self.center_points)
+        )
+
+        target_positions = []
+        for leg_id in range(6):
+            base_pos = np.array(center_data[leg_id], dtype=np.float32)
+            
+            # Tripod leg grouping phase offset (even vs odd leg IDs)
+            phase_offset = 0.0 if (leg_id % 2 == 0) else math.pi
+            leg_phase = (self.gait_phase + phase_offset) % (2.0 * math.pi)
+
+            # Directional motion unit vector calculation
+            lx, ly = base_pos[0], base_pos[1]
+            rot_dx = -omega * ly
+            rot_dy = omega * lx
+
+            dx_dir = vx + rot_dx
+            dy_dir = vy + rot_dy
+            dir_norm = math.hypot(dx_dir, dy_dir) + 1e-6
+
+            dx_unit = dx_dir / dir_norm
+            dy_unit = dy_dir / dir_norm
+
+            if leg_phase < math.pi:
+                # Swing Phase (Leg lifted & stepping forward)
+                progress = math.cos(leg_phase)
+                lift = math.sin(leg_phase) * step_height
+                dx = -progress * stride_len * dx_unit
+                dy = -progress * stride_len * dy_unit
+                dz = lift
+            else:
+                # Stance Phase (Leg grounded & propelling torso)
+                progress = math.cos(leg_phase - math.pi)
+                dx = progress * stride_len * dx_unit
+                dy = progress * stride_len * dy_unit
+                dz = 0.0
+
+            target_leg_pos = base_pos + np.array([dx, dy, dz], dtype=np.float32)
+            target_positions.append(target_leg_pos)
+
+        target_pos_array = dt.PosArray(np.array(target_positions))
+        target_rad = self.compute_ik(target_pos_array)
+        self.set_joint_angles(target_rad)
+
+    def step_with_command(
+        self, 
+        command: np.ndarray, 
+        action: Optional[np.ndarray] = None, 
+        mode: str = "direct"
+    ) -> None:
+        """
+        Unified motion execution method supporting Direct RL, Residual RL, and Pure Kinematics.
+
+        Args:
+            command: np.ndarray [vx, vy, vz, omega] from RL environment
+            action: np.ndarray [18,] RL action deltas from neural network ([-1, 1])
+            mode: "kinematics_only" | "residual" | "direct"
+        """
+        # 1. Map RL command [vx, vy, vz, omega] to Robot motion vector [vx, vy, omega]
+        cmd_vx, cmd_vy, _, cmd_omega = command
+        self.vector_dirmov = [float(cmd_vx), float(cmd_vy), float(cmd_omega)]
+
+        # 2. Automatically trigger state transitions based on command magnitude
+        cmd_magnitude = math.hypot(cmd_vx, cmd_vy) + abs(cmd_omega)
+        if cmd_magnitude > 0.05 and self.current_state_key == "idle":
+            # Transition to walking state if registered in state machine
+            target_state = "walk" if "walk" in STATE_REGISTRY else "move"
+            if target_state in STATE_REGISTRY:
+                self.transition_to(target_state)
+        elif cmd_magnitude <= 0.05 and self.current_state_key != "idle":
+            self.transition_to("idle")
+
+        # 3. Execute according to chosen mode
+        if mode == "kinematics_only":
+            self.step_kinematic_gait(cmd_vx, cmd_vy, cmd_omega)
+
+        elif mode == "residual":
+            if self.current_state_key != "idle" and self.current_state and self.current_state_key in STATE_REGISTRY:
+                self.current_state.execute(self)
+            else:
+                self.step_kinematic_gait(cmd_vx, cmd_vy, cmd_omega)
+
+            if action is not None:
+                action_flat = np.clip(np.asarray(action, dtype=np.float32), -1.0, 1.0) * self.action_scale
+                kin_flat = self.current_rad.data.flatten()
+                final_flat = np.clip(kin_flat + action_flat, -np.pi / 2, np.pi / 2)
+                self.set_joint_angles(dt.RadArray(data=final_flat.reshape(self.current_rad.data.shape)))
+
+        elif mode == "direct":
+            if action is not None:
+                self.apply_rl_action(action)
+
     # --- RL METHODS ---
 
     def apply_rl_action(self, action: np.ndarray) -> None:
-        """Applies continuous RL action deltas [-1, 1] to current joint angles."""
         action = np.asarray(action, dtype=np.float32)
         scaled_action = np.clip(action, -1.0, 1.0) * self.action_scale
 
-        current_flat = self.current_rad.data.flatten()
+        # Sync with actual PyBullet state if backend supports it
-        updated_flat = np.clip(
+        if isinstance(self.backend, PyBulletBackend) and self.backend.sim:
-            current_flat + scaled_action,
+            actual_angles = self.backend.sim.get_robot_joint_angles(self.backend.body_id)
-            -np.pi / 2,
+            current_flat = actual_angles.flatten()
-            np.pi / 2
+        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)
 

ml/SimManager.py

@@ -1,7 +1,7 @@
 """
-ml/SimManager.py - PyBullet Simulation & Multi-Body Manager
+ml/SimManager.py - PyBullet Simulation Manager (Single Robot Dedicated)
 """
-from typing import Dict, List, Tuple, Optional
+from typing import List, Tuple, Optional, Union
 import pybullet as p
 import pybullet_data
 import numpy as np
@@ -9,14 +9,16 @@ import DataTypes as dt
 
 
 class SimManager:
-    """Manages PyBullet simulation lifecycle and multi-robot physics."""
+    """Manages PyBullet simulation lifecycle for a single JackBot hexapod."""
 
     def __init__(self, use_gui: bool = True):
         self.use_gui = use_gui
         self.physics_client = None
-        self.robot_joints: Dict[int, List[int]] = {}
+        self.plane: Optional[int] = None
+        self.joint_indices: List[int] = []
+        self.foot_indices: List[int] = []
 
-    def connect(self):
+    def connect(self) -> None:
         """Connects to PyBullet and hides side GUI panels."""
         if self.physics_client is not None and p.isConnected(self.physics_client):
             return
@@ -28,62 +30,87 @@ class SimManager:
         p.setGravity(0, 0, -9.81, physicsClientId=self.physics_client)
 
         if self.use_gui:
-            # Disable PyBullet side panel and preview windows
             p.configureDebugVisualizer(p.COV_ENABLE_GUI, 0, physicsClientId=self.physics_client)
             p.configureDebugVisualizer(p.COV_ENABLE_SEGMENTATION_MARK_PREVIEW, 0, physicsClientId=self.physics_client)
             p.configureDebugVisualizer(p.COV_ENABLE_DEPTH_BUFFER_PREVIEW, 0, physicsClientId=self.physics_client)
             p.configureDebugVisualizer(p.COV_ENABLE_RGB_BUFFER_PREVIEW, 0, physicsClientId=self.physics_client)
 
     def load_scene(
-        self, urdf_path: str, robot_spacing: float, base_pos_fn
+        self, 
+        urdf_path: str, 
+        spacing: float = 0.0, 
+        position_func=None
     ) -> Tuple[int, List[int], List[List[int]]]:
-        """Loads the plane and a single hexapod body into the simulation scene."""
+        """Loads ground plane and the single robot URDF with high ground friction and force settings."""
-        plane_id = p.loadURDF("plane.urdf", physicsClientId=self.physics_client)
+        # 1. Load ground plane and set explicit friction
-        robots = []
+        self.plane = p.loadURDF("plane.urdf", physicsClientId=self.physics_client)
-        robot_joint_indices = []
-        self.robot_joints.clear()
 
-        base_pos = base_pos_fn(0, robot_spacing)
+        # 2. Determine starting position (matching simulation.py starting height of 0.20m)
-        robot = p.loadURDF(
+        spawn_pos = position_func(0) if position_func is not None else [0.0, 0.0, 0.20]
-            urdf_path,
-            basePosition=base_pos,
-            useFixedBase=False,
-            physicsClientId=self.physics_client
-        )
-        robots.append(robot)
 
-        joint_indices = [
+        # 3. Spawn single robot body
-            i for i in range(p.getNumJoints(robot, physicsClientId=self.physics_client))
+        self.robot_id = p.loadURDF(urdf_path, spawn_pos, physicsClientId=self.physics_client)
-            if p.getJointInfo(robot, i, physicsClientId=self.physics_client)[2] == p.JOINT_REVOLUTE
-        ]
-        robot_joint_indices.append(joint_indices)
-        self.robot_joints[robot] = joint_indices
 
-        return plane_id, robots, robot_joint_indices
+        # 4. Retrieve revolute joint indices
+        self.joint_indices = []
+        for j in range(p.getNumJoints(self.robot_id, physicsClientId=self.physics_client)):
+            info = p.getJointInfo(self.robot_id, j, physicsClientId=self.physics_client)
+            if info[2] == p.JOINT_REVOLUTE:
+                self.joint_indices.append(j)
 
-    def updatePosForBody(self, body_id: int, current_rad: dt.RadArray):
+        # Return format maintains 100% compatibility with JackBotEnv unpack sequence
-        """Sets joint motor position targets on individual robot bodies."""
+        return self.plane, [self.robot_id], [self.joint_indices]
-        if body_id not in self.robot_joints:
+
-            return
+    def _resolve_body_id(self, body_id: Optional[int] = None) -> int:
-            
+        """Internal helper to return the single active robot ID."""
-        joint_indices = self.robot_joints[body_id]
+        if body_id is not None:
-        radflat = current_rad.data.flatten()
+            return body_id
+        if self.robot_id is not None:
+            return self.robot_id
+        raise RuntimeError("No robot loaded in SimManager. Call load_scene() first.")
+
+    def set_robot_joint_angles(
+        self, 
+        target_angles: Union[np.ndarray, List[float]], 
+        joint_indices: Optional[List[int]] = None,
+        body_id: Optional[int] = None,
+    ) -> None:
+        """Applies position control with force=500 matching simulation.py logic."""
+        bid = self._resolve_body_id(body_id)
+        j_indices = joint_indices if joint_indices is not None else self.joint_indices
+        radflat = target_angles.flatten() if isinstance(target_angles, np.ndarray) else target_angles
         
-        for joint_index, target_angle in zip(joint_indices, radflat):
+        for joint_index, target_angle in zip(j_indices, radflat):
             p.setJointMotorControl2(
-                bodyIndex=body_id,
+                bodyIndex=bid,
                 jointIndex=joint_index,
                 controlMode=p.POSITION_CONTROL,
                 targetPosition=float(target_angle),
-                force=250,
+                force=500,
                 physicsClientId=self.physics_client
             )
 
-    def step(self):
+    def updatePosForBody(
+        self, 
+        body_id_or_rad: Union[int, dt.RadArray], 
+        current_rad: Optional[dt.RadArray] = None
+    ) -> None:
+        """Updates robot joint positions using RadArray input."""
+        if isinstance(body_id_or_rad, dt.RadArray):
+            rad_data = body_id_or_rad
+            bid = self.robot_id
+        else:
+            bid = self._resolve_body_id(body_id_or_rad)
+            rad_data = current_rad
+
+        if rad_data is not None:
+            self.set_robot_joint_angles(rad_data.data, body_id=bid)
+
+    def step(self) -> None:
         p.stepSimulation(physicsClientId=self.physics_client)
 
     def set_rendering(self, enabled: bool) -> None:
-        """Toggles PyBullet 3D rendering to speed up simulation."""
+        """Toggles PyBullet 3D rendering visualizer."""
         if self.physics_client is not None and p.isConnected(self.physics_client):
             p.configureDebugVisualizer(
                 p.COV_ENABLE_RENDERING, 
@@ -91,7 +118,7 @@ class SimManager:
                 physicsClientId=self.physics_client
             )
 
-    def disconnect(self):
+    def disconnect(self) -> None:
         if self.physics_client is not None and p.isConnected(self.physics_client):
             p.disconnect(self.physics_client)
             self.physics_client = None
@@ -116,114 +143,153 @@ class SimManager:
 
         return p.getContactPoints(**kwargs)
 
-    # --- ROBOT GETTERS AND SETTERS ---
+    # Add to ml/SimManager.py
+    def hard_reset_joint_angles(self, target_angles: np.ndarray, body_id: Optional[int] = None) -> None:
+        bid = self._resolve_body_id(body_id)
+        radflat = target_angles.flatten()
+        for joint_index, target_angle in zip(self.joint_indices, radflat):
+            p.resetJointState(
+                bodyUniqueId=bid,
+                jointIndex=joint_index,
+                targetValue=float(target_angle),
+                targetVelocity=0.0,
+                physicsClientId=self.physics_client
+            )
+
+    # --- SINGLE ROBOT GETTERS & SETTERS ---
 
     def reset_robot_base(
         self,
-        body_id: int,
+        body_id_or_pos: Union[int, List[float]],
-        position: List[float],
+        position_or_orn: Optional[List[float]] = None,
         orientation: Optional[List[float]] = None,
         linear_velocity: Optional[List[float]] = None,
         angular_velocity: Optional[List[float]] = None
     ) -> None:
-        """Resets a robot body's base position, orientation, and velocities."""
+        """Resets the robot base pose and clears linear/angular velocities."""
-        if orientation is None:
+        if isinstance(body_id_or_pos, int):
-            orientation = [0.0, 0.0, 0.0, 1.0]
+            bid = body_id_or_pos
-        if linear_velocity is None:
+            pos = position_or_orn if position_or_orn is not None else [0.0, 0.0, 0.20]
-            linear_velocity = [0.0, 0.0, 0.0]
+            orn = orientation if orientation is not None else [0.0, 0.0, 0.0, 1.0]
-        if angular_velocity is None:
+        else:
-            angular_velocity = [0.0, 0.0, 0.0]
+            bid = self.robot_id
+            pos = body_id_or_pos
+            orn = position_or_orn if position_or_orn is not None else [0.0, 0.0, 0.0, 1.0]
+
+        lin_v = linear_velocity if linear_velocity is not None else [0.0, 0.0, 0.0]
+        ang_v = angular_velocity if angular_velocity is not None else [0.0, 0.0, 0.0]
 
         p.resetBasePositionAndOrientation(
-            body_id, position, orientation, physicsClientId=self.physics_client
+            bid, pos, orn, physicsClientId=self.physics_client
         )
         p.resetBaseVelocity(
-            body_id, linearVelocity=linear_velocity, angularVelocity=angular_velocity,
+            bid, linearVelocity=lin_v, angularVelocity=ang_v,
             physicsClientId=self.physics_client
         )
 
-    def get_robot_pose(self, body_id: int) -> Tuple[List[float], List[float]]:
+    def get_robot_pose(self, body_id: Optional[int] = None) -> Tuple[List[float], List[float]]:
         """Returns base position (x, y, z) and orientation quaternion (x, y, z, w)."""
-        pos, orn = p.getBasePositionAndOrientation(body_id, physicsClientId=self.physics_client)
+        bid = self._resolve_body_id(body_id)
+        pos, orn = p.getBasePositionAndOrientation(bid, physicsClientId=self.physics_client)
         return list(pos), list(orn)
 
-    def get_robot_rpy(self, body_id: int) -> Tuple[float, float, float]:
+    def get_robot_rpy(self, body_id: Optional[int] = None) -> Tuple[float, float, float]:
-        """Returns roll, pitch, yaw angles in radians for the given robot body."""
+        """Returns roll, pitch, yaw angles in radians for the robot."""
-        _, orn = p.getBasePositionAndOrientation(body_id, physicsClientId=self.physics_client)
+        bid = self._resolve_body_id(body_id)
+        _, orn = p.getBasePositionAndOrientation(bid, physicsClientId=self.physics_client)
         roll, pitch, yaw = p.getEulerFromQuaternion(orn)
         return float(roll), float(pitch), float(yaw)
 
-    def get_foot_link_indices(self, body_id: int) -> list[int]:
+    def get_robot_pose_and_rpy(self, body_id: Optional[int] = None) -> Tuple[List[float], Tuple[float, float, float]]:
+        """Returns base position and (roll, pitch, yaw) tuple in radians."""
+        bid = self._resolve_body_id(body_id)
+        pos, orn = p.getBasePositionAndOrientation(bid, physicsClientId=self.physics_client)
+        roll, pitch, yaw = p.getEulerFromQuaternion(orn)
+        return list(pos), (float(roll), float(pitch), float(yaw))
+
+    def get_robot_velocity(self, body_id: Optional[int] = None) -> Tuple[List[float], List[float]]:
+        """Returns linear velocity (vx, vy, vz) and angular velocity (wx, wy, wz)."""
+        bid = self._resolve_body_id(body_id)
+        lin_v, ang_v = p.getBaseVelocity(bid, physicsClientId=self.physics_client)
+        return list(lin_v), list(ang_v)
+
+    def get_robot_joint_angles(
+        self, 
+        body_id: Optional[int] = None, 
+        joint_indices: Optional[List[int]] = None
+    ) -> np.ndarray:
+        """Returns joint angles as a 1D numpy float32 array."""
+        bid = self._resolve_body_id(body_id)
+        j_indices = joint_indices if joint_indices is not None else self.joint_indices
+        joint_states = p.getJointStates(bid, j_indices, physicsClientId=self.physics_client)
+        return np.array([state[0] for state in joint_states], dtype=np.float32)
+
+    def get_foot_link_indices(self, body_id: Optional[int] = None) -> List[int]:
         """Inspects URDF joint structure to extract link IDs for leg tips and tibias."""
+        bid = self._resolve_body_id(body_id)
         foot_indices = []
-        num_joints = p.getNumJoints(body_id, physicsClientId=self.physics_client)
+        num_joints = p.getNumJoints(bid, physicsClientId=self.physics_client)
         for j_idx in range(num_joints):
-            info = p.getJointInfo(body_id, j_idx, physicsClientId=self.physics_client)
+            info = p.getJointInfo(bid, j_idx, physicsClientId=self.physics_client)
             link_name = info[12].decode("utf-8")
             if "tip" in link_name or "tibia" in link_name:
                 foot_indices.append(j_idx)
         return foot_indices
 
-    def get_robot_pose_and_rpy(self, body_id: int) -> Tuple[List[float], Tuple[float, float, float]]:
+    def set_robot_color(
-        """Returns base position and (roll, pitch, yaw) tuple in radians."""
+        self, 
-        pos, orn = p.getBasePositionAndOrientation(body_id, physicsClientId=self.physics_client)
+        body_id_or_rgba: Union[int, List[float]], 
-        roll, pitch, yaw = p.getEulerFromQuaternion(orn)
+        rgba: Optional[List[float]] = None
-        return list(pos), (float(roll), float(pitch), float(yaw))
+    ) -> None:
+        """Changes visual color RGBA of base link and all joints."""
+        if isinstance(body_id_or_rgba, int):
+            bid = body_id_or_rgba
+            color = rgba
+        else:
+            bid = self.robot_id
+            color = body_id_or_rgba
 
-    def get_robot_velocity(self, body_id: int) -> Tuple[List[float], List[float]]:
+        if color is None:
-        """Returns linear velocity (vx, vy, vz) and angular velocity (wx, wy, wz)."""
+            color = [1.0, 1.0, 1.0, 1.0]
-        lin_v, ang_v = p.getBaseVelocity(body_id, physicsClientId=self.physics_client)
-        return list(lin_v), list(ang_v)
 
-    def get_robot_joint_angles(self, body_id: int, joint_indices: Optional[List[int]] = None) -> np.ndarray:
+        num_joints = p.getNumJoints(bid, physicsClientId=self.physics_client)
-        """Returns joint angles as a 1D numpy array float32 for specified or registered joint indices."""
+        p.changeVisualShape(bid, -1, rgbaColor=color, physicsClientId=self.physics_client)
-        if joint_indices is None:
-            joint_indices = self.robot_joints.get(body_id, list(range(18)))
-        joint_states = p.getJointStates(body_id, joint_indices, physicsClientId=self.physics_client)
-        return np.array([state[0] for state in joint_states], dtype=np.float32)
-
-    def set_robot_color(self, body_id: int, rgba: List[float]) -> None:
-        """Changes visual color RGBA of base link and all joints of the specified robot body."""
-        num_joints = p.getNumJoints(body_id, physicsClientId=self.physics_client)
-        p.changeVisualShape(body_id, -1, rgbaColor=rgba, physicsClientId=self.physics_client)
         for j in range(num_joints):
-            p.changeVisualShape(body_id, j, rgbaColor=rgba, physicsClientId=self.physics_client)
+            p.changeVisualShape(bid, j, rgbaColor=color, physicsClientId=self.physics_client)
 
-    def measure_robot_heights(self, robot_ids: List[int]) -> List[float]:
+    def measure_robot_height(self, body_id: Optional[int] = None) -> float:
-        """Gets current Z height for all specified robot body IDs."""
+        """Gets current Z height of the robot base."""
-        heights = []
+        bid = self._resolve_body_id(body_id)
-        for body_id in robot_ids:
+        pos, _ = p.getBasePositionAndOrientation(bid, physicsClientId=self.physics_client)
-            pos, _ = p.getBasePositionAndOrientation(body_id, physicsClientId=self.physics_client)
+        return pos[2]
-            heights.append(pos[2])
-        return heights
 
-    def settle_and_measure_height(self, robot_ids: List[int], steps: int = 200, fallback_height: float = 0.14) -> float:
+    def settle_and_measure_height(
-        """Steps simulation for designated steps so robot settles, then calculates target standing height."""
+        self, 
-        for _ in range(steps):
+        robot_ids_or_steps: Union[List[int], int] = 200, 
+        steps: int = 200, 
+        fallback_height: float = 0.122
+    ) -> float:
+        """Steps simulation so robot settles onto floor, then calculates target standing height."""
+        # Intelligently resolve whether the first argument was passed as robot_ids list or step count
+        actual_steps = robot_ids_or_steps if isinstance(robot_ids_or_steps, int) else steps
+
+        for _ in range(actual_steps):
             self.step()
-        heights = self.measure_robot_heights(robot_ids)
+        height = self.measure_robot_height()
-        mean_height = float(np.mean(heights)) if heights else fallback_height
+        return height if height > 0.0 else fallback_height
-        return mean_height if mean_height > 0.0 else fallback_height
 
     def apply_external_force(
         self,
-        body_id: int,
+        force: Union[List[float], np.ndarray],
-        force: list[float] | np.ndarray,
+        body_id: Optional[int] = None,
         link_index: int = -1,
-        position: list[float] | np.ndarray = (0.0, 0.0, 0.0),
+        position: Union[List[float], np.ndarray] = (0.0, 0.0, 0.0),
         frame: int = p.WORLD_FRAME,
-    ):
+    ) -> None:
-        """
+        """Applies external force vector to target link (defaults to base link)."""
-        Applies a 3D force vector (in Newtons) to a robot link.
+        bid = self._resolve_body_id(body_id)
-        
-        :param body_id: PyBullet body ID.
-        :param force: [fx, fy, fz] force vector in Newtons.
-        :param link_index: Target link index (-1 refers to the base/torso).
-        :param position: Offset [x, y, z] relative to link center where force is applied.
-        :param frame: p.WORLD_FRAME (global axes) or p.LINK_FRAME (robot's body axes).
-        """
         p.applyExternalForce(
-            objectUniqueId=body_id,
+            objectUniqueId=bid,
             linkIndex=link_index,
             forceObj=list(force),
             posObj=list(position),

ml/env.py

@@ -5,6 +5,7 @@ import time
 import math
 from enum import IntEnum
 from typing import Optional, Tuple, Dict, Any, List
+from collections import defaultdict
 
 import gymnasium as gym
 from gymnasium import spaces
@@ -37,8 +38,10 @@ class JackBotEnv(gym.Env):
         random_command: bool = True,
         max_episode_steps: int = 3000,
         urdf_path: str = cfg.urdf_path,
+        robot_mode: str = "direct",
     ):
         super().__init__()
+        self.robot_mode = robot_mode
         self.use_gui = use_gui
         self.random_command = random_command
         self.max_episode_steps = max_episode_steps
@@ -59,6 +62,10 @@ class JackBotEnv(gym.Env):
         self._curriculum_advanced = False
         self._first_reset = True
 
+        # Reward Component Tracking Initialization
+        self.last_reward_components: Dict[str, float] = {}
+        self.episode_reward_components_sum: Dict[str, float] = defaultdict(float)
+
         # Dynamic Command Resampling Timing (60 Hz control loop)
         self.control_freq = 60
         self.min_cmd_hold_steps = int(2.0 * self.control_freq)  # 120 steps (2s)
@@ -74,16 +81,14 @@ class JackBotEnv(gym.Env):
         self.plane, pb_robots, robot_joint_indices = self.sim_manager.load_scene(
             self.urdf_path, 0.0, self._robot_base_position
         )
-        self.pb_robot = pb_robots[0]
         self.joint_indices = robot_joint_indices[0]
 
-        # Instantiate Robot Python wrapper (start_pose is managed inside Robot.py)
+        # Instantiate Robot Python wrapper
         self.robot = Robot(
-            backend_type=PyBulletBackend(self.sim_manager, body_id=self.pb_robot),
+            backend_type=PyBulletBackend(self.sim_manager),
             urdf_path=self.urdf_path
         )
 
-        # Action (18 joint deltas) & Observation (18 angles + 4 command dims)
         action_dim = 18
         obs_dim = 18 + 4
 
@@ -109,7 +114,7 @@ class JackBotEnv(gym.Env):
             CurriculumPhase.FORWARD: {
                 "survival_steps": 300,
                 "min_avg_height_ratio": 0.88,
-                "max_avg_roll_pitch": 0.18,  # ~10 degrees average
+                "max_avg_roll_pitch": 0.18,
             },
             CurriculumPhase.TURN_AND_DIRECTION: {
                 "survival_steps": 500,
@@ -136,7 +141,7 @@ class JackBotEnv(gym.Env):
         self.last_time = time.time()
 
     def _robot_base_position(self, robot_id: int, spacing: float = 0.0) -> list[float]:
-        return [0.0, 0.0, 0.13]
+        return [0.0, 0.0, 0.2]
 
     def _find_foot_link_indices(self) -> list:
         return self.sim_manager.get_foot_link_indices(self.pb_robot)
@@ -187,12 +192,19 @@ class JackBotEnv(gym.Env):
         self.episode_roll_sum = 0.0
         self.episode_pitch_sum = 0.0
 
+        # Reset Component Tracking Dictionary
+        self.last_reward_components = {}
+        self.episode_reward_components_sum = defaultdict(float)
+
         spawn_pos = self._robot_base_position(0)
         spawn_orn = [0.0, 0.0, 0.0, 1.0]
 
         self.sim_manager.reset_robot_base(self.pb_robot, spawn_pos, spawn_orn)
         self.robot.reset_to_init()
 
+        init_angles = self.robot.current_rad.data.flatten()
+        self.sim_manager.hard_reset_joint_angles(init_angles, self.pb_robot)
+
         if self.use_gui:
             self.sim_manager.set_robot_color(self.pb_robot, [1.0, 1.0, 1.0, 1.0])
 
@@ -211,8 +223,9 @@ class JackBotEnv(gym.Env):
         pos, _ = self.sim_manager.get_robot_pose(self.pb_robot)
         self.start_position = [float(pos[0]), float(pos[1]), float(pos[2])]
 
+        # Fixed single-robot settlement call
         self.target_height = self.sim_manager.settle_and_measure_height(
-            [self.pb_robot], steps=200, fallback_height=0.122
+            steps=200, fallback_height=0.122
         )
 
         self.default_joint_angles = np.array(
@@ -226,8 +239,13 @@ class JackBotEnv(gym.Env):
 
         return self._get_obs(), {}
 
+    def get_reward_component_averages(self) -> Dict[str, float]:
+        """Calculates step-averaged scores for each sub-reward component."""
+        steps = max(1, self.step_count)
+        return {k: float(v / steps) for k, v in self.episode_reward_components_sum.items()}
+
     def get_current_robot_metrics(self) -> list:
-        """Returns metric summary for the callback."""
+        """Returns metric summary for callbacks."""
         if self.is_failed:
             return []
 
@@ -263,17 +281,16 @@ class JackBotEnv(gym.Env):
             random_interval = np.random.randint(self.min_cmd_hold_steps, self.max_cmd_hold_steps + 1)
             self.next_cmd_resample_step = self.step_count + random_interval
 
-        self.robot.apply_rl_action(action)
+        self.robot.step_with_command(command=self.command, action=action, mode=self.robot_mode)
 
-        if self.step_count % 60 == 0:
+        if self.robot_mode != "kinematics_only" and self.step_count % 60 == 0:
             random_force = np.random.uniform(-2.0, 2.0, size=2)
             self.sim_manager.apply_external_force(
                 body_id=self.pb_robot, 
                 force=[random_force[0], random_force[1], 0.0]
             )
 
-        render_freq = 10  # Only draw 1 in every 10 frames
+        render_freq = 10
-        
         if self.use_gui and self.step_count % render_freq != 0:
             self.sim_manager.set_rendering(False)
             
@@ -295,9 +312,15 @@ class JackBotEnv(gym.Env):
         terminated = self.is_failed
         truncated = self.step_count >= self.max_episode_steps
 
+        info = {
+            "reward_components": self.last_reward_components.copy()
+        }
+
         if self.step_count % 120 == 0 and self.use_gui:
             self._update_hud()
-        return obs, reward, terminated, truncated, {}
+
+        # Gymnasium standard 5-tuple return
+        return obs, reward, terminated, truncated, info
 
     def _update_distance_metrics(self):
         pos, _ = self.sim_manager.get_robot_pose(self.pb_robot)
@@ -312,22 +335,17 @@ class JackBotEnv(gym.Env):
             return False
 
         req = self.curriculum_stage_requirements[next_phase]
-        
-        # 1. Survival Check
         survival_ok = self.max_survival_steps >= req["survival_steps"]
 
-        # 2. Smooth Average Stability Checks (Prevents 1-frame spikes from failing curriculum)
         avg_roll = self.episode_roll_sum / max(1, self.step_count)
         avg_pitch = self.episode_pitch_sum / max(1, self.step_count)
         max_allowed_angle = req.get("max_avg_roll_pitch", 0.20)
         stability_ok = (avg_roll <= max_allowed_angle) and (avg_pitch <= max_allowed_angle)
 
-        # 3. Average Height Check
         avg_height = self.episode_height_sum / max(1, self.step_count)
         required_min_avg_height = self.target_height * req.get("min_avg_height_ratio", 0.85)
         height_ok = avg_height >= required_min_avg_height
 
-        # 4. Distance and Drift Checks
         pos, _ = self.sim_manager.get_robot_pose(self.pb_robot)
         start_x, start_y, _ = self.start_position
         dx = pos[0] - start_x
@@ -373,10 +391,8 @@ class JackBotEnv(gym.Env):
 
     def _compute_reward(self, action: np.ndarray, previous_action: np.ndarray) -> float:
         """
-        Calculates task rewards using normalized Exponential Kernels.
+        Calculates task rewards using normalized Exponential Kernels and tracks component terms.
-        Includes a deadband filter for jittering and zero-reward gating when stationary.
         """
-        # 1. Fetch Robot State
         pos, (roll, pitch, yaw) = self.sim_manager.get_robot_pose_and_rpy(self.pb_robot)
         linear_vel, angular_vel = self.sim_manager.get_robot_velocity(self.pb_robot)
         current_joints = np.array(
@@ -387,9 +403,8 @@ class JackBotEnv(gym.Env):
         cmd_vx, cmd_vy, _, cmd_yaw = self.command
         cmd_norm = math.hypot(cmd_vx, cmd_vy)
 
-        # 2. Velocity Deadband Filtering (Ignores jittering & micro-movements)
+        VEL_DEADBAND = 0.04
-        VEL_DEADBAND = 0.04   # 4 cm/s threshold
+        YAW_DEADBAND = 0.05
-        YAW_DEADBAND = 0.05   # 0.05 rad/s threshold
 
         raw_speed = math.hypot(linear_vel[0], linear_vel[1])
         if raw_speed < VEL_DEADBAND:
@@ -405,7 +420,6 @@ class JackBotEnv(gym.Env):
         else:
             filtered_yaw_rate = angular_vel[2]
 
-        # 3. Posture & Stability Sub-Rewards
         height_error = pos[2] - self.target_height
         r_height = math.exp(-150.0 * (height_error ** 2))
 
@@ -418,13 +432,19 @@ class JackBotEnv(gym.Env):
         action_delta = np.mean(np.square(action - previous_action))
         r_smoothness = math.exp(-0.1 * action_delta)
 
-        # 4. Mode Logic
+        r_lin_vel = 0.0
+        r_ang_vel = 0.0
+        stillness_penalty = 0.0
+        gated_zero = 0.0
+
         if cmd_norm < 0.05 and abs(cmd_yaw) < 0.05:
-            # STANDING MODE: Reward clean posture, height, and stability
+            # STANDING MODE
             w_height = 0.35
             w_stability = 0.35
             w_pose = 0.20
             w_smoothness = 0.10
+            w_lin_vel = 0.0
+            w_ang_vel = 0.0
 
             total_reward = (
                 (w_height * r_height)
@@ -436,42 +456,60 @@ class JackBotEnv(gym.Env):
             # WALKING / TURNING MODE
             is_moving = (filtered_speed > 0.0) or (abs(filtered_yaw_rate) > 0.0)
 
-            # HARD GATE: If commanded to move but standing still/jittering, reward is strictly 0.0
-            if not is_moving:
-                return 0.0
-
             target_vx = cmd_vx * self.max_robot_speed
             target_vy = cmd_vy * self.max_robot_speed
+            target_speed = math.hypot(target_vx, target_vy)
 
-            lin_vel_error = (filtered_vx - target_vx)**2 + (filtered_vy - target_vy)**2
+            if not is_moving:
-            r_lin_vel = math.exp(-25.0 * lin_vel_error)
+                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
+                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 * ang_vel_error)
 
-            # Stillness Check: Commanded to move, but staying virtually still
+                if target_speed > 0.08 and raw_speed < 0.03:
-            stillness_penalty = 0.0
+                    r_lin_vel = 0.0
-            if math.hypot(target_vx, target_vy) > 0.08 and math.hypot(linear_vel[0], linear_vel[1]) < 0.03:
+                    stillness_penalty = -0.25
-                r_lin_vel = 0.0            # Strip velocity credit completely
-                stillness_penalty = -0.25
 
-            w_lin_vel = 0.55
+                w_lin_vel = 0.55
-            w_ang_vel = 0.15
+                w_ang_vel = 0.15
-            w_height = 0.10
+                w_height = 0.10
-            w_stability = 0.12
+                w_stability = 0.12
-            w_smoothness = 0.08
+                w_pose = 0.0
+                w_smoothness = 0.08
 
-            total_reward = (
+                total_reward = (
-                (w_lin_vel * r_lin_vel)
+                    (w_lin_vel * r_lin_vel)
-                + (w_ang_vel * r_ang_vel)
+                    + (w_ang_vel * r_ang_vel)
-                + (w_height * r_height)
+                    + (w_height * r_height)
-                + (w_stability * r_stability)
+                    + (w_stability * r_stability)
-                + (w_smoothness * r_smoothness)
+                    + (w_smoothness * r_smoothness)
-                + stillness_penalty
+                    + stillness_penalty
-            )
+                )
 
-        # Scaled reward for policy stability
+        final_reward = float(total_reward / 10.0)
-        return float(total_reward / 10.0)
+
+        comp = {
+            "lin_vel": float((w_lin_vel * r_lin_vel) / 10.0),
+            "ang_vel": float((w_ang_vel * r_ang_vel) / 10.0),
+            "height": float((w_height * r_height) / 10.0),
+            "stability": float((w_stability * r_stability) / 10.0),
+            "pose": float((w_pose * r_pose) / 10.0),
+            "smoothness": float((w_smoothness * r_smoothness) / 10.0),
+            "stillness_penalty": float(stillness_penalty / 10.0),
+            "gated_zero": gated_zero,
+            "total_step_reward": final_reward,
+        }
+
+        self.last_reward_components = comp
+        for key, val in comp.items():
+            self.episode_reward_components_sum[key] += val
+
+        return final_reward
 
     def _update_hud(self):
         if not self.use_gui:
@@ -579,4 +617,34 @@ class CurriculumCallback(BaseCallback):
         except Exception:
             pass
 
+        return True
+
+
+class RewardLoggerCallback(BaseCallback):
+    """
+    Logs step-averaged individual reward components to TensorBoard during PPO training.
+    """
+
+    def __init__(self, verbose=0):
+        super().__init__(verbose)
+
+    def _on_step(self) -> bool:
+        return True
+
+    def _on_rollout_end(self) -> bool:
+        try:
+            vec_env = self.training_env
+            all_comp_averages = vec_env.env_method("get_reward_component_averages")
+
+            if not all_comp_averages:
+                return True
+
+            keys = all_comp_averages[0].keys()
+            for key in keys:
+                avg_val = np.mean([env_comp.get(key, 0.0) for env_comp in all_comp_averages])
+                self.logger.record(f"reward_components/{key}", float(avg_val))
+
+        except Exception:
+            pass
+
         return True

ml/run_eval_training.py

@@ -0,0 +1,76 @@
+"""
+ml/run_eval_training.py - Benchmark reward system across all curriculum phases.
+"""
+import time
+import numpy as np
+from ml.env import JackBotEnv, CurriculumPhase
+
+def evaluate_kinematics(episode_length: int = 1000):
+    # Use kinematics_only mode so inverse kinematics generates gait motion from commands
+    env = JackBotEnv(
+        use_gui=True, 
+        random_command=False, 
+        max_episode_steps=episode_length,
+        robot_mode="kinematics_only"
+    )
+
+    print("\n" + "=" * 70)
+    print(" RUNNING MULTI-PHASE REWARD BENCHMARK (KINEMATICS MODE)")
+    print("=" * 70 + "\n")
+
+    # Define test suite covering every curriculum stage
+    phase_configs = [
+        (CurriculumPhase.STAND_ONLY,          "STAND ONLY",          np.array([0.0, 0.0, 0.0, 0.0], dtype=np.float32)),
+        (CurriculumPhase.FORWARD,             "FORWARD GAIT",        np.array([0.1, 0.0, 0.0, 0.0], dtype=np.float32)),
+        (CurriculumPhase.TURN_AND_DIRECTION,  "FORWARD + YAW TURN",  np.array([0.3, 0.0, 0.0, 0.4], dtype=np.float32)),
+        (CurriculumPhase.OMNI_DIRECTION,      "STRIDE LATERAL",      np.array([0.2, 0.3, 0.0, 0.0], dtype=np.float32)),
+        (CurriculumPhase.FULL_COMMAND,        "FULL OMNI COMBINATION",np.array([0.3, 0.2, 0.0, 0.3], dtype=np.float32)),
+    ]
+
+    for phase_enum, label, cmd in phase_configs:
+        obs, _ = env.reset()
+        
+        # Force specific curriculum phase & target command
+        env.curriculum_phase = phase_enum
+        env.command = cmd.copy()
+        
+        done = False
+        total_reward = 0.0
+        step_count = 0
+
+        while not done:
+            # Action array is unused in kinematics_only mode
+            dummy_action = np.zeros(18, dtype=np.float32)
+            
+            obs, reward, terminated, truncated, _ = env.step(dummy_action)
+            total_reward += reward
+            step_count += 1
+            done = terminated or truncated
+
+            time.sleep(1.0 / 60.0)
+
+        # Retrieve detailed component averages
+        comp_averages = env.get_reward_component_averages()
+
+        print(f"\n--- Episode Stage: [{phase_enum.name}] ({label}) ---")
+        print(f"Command Applied: vx={cmd[0]:.2f}, vy={cmd[1]:.2f}, vz={cmd[2]:.2f}, yaw={cmd[3]:.2f}")
+        print(f"Total Episode Reward: {total_reward:.4f}")
+        print("Component Step Averages:")
+        for name, value in comp_averages.items():
+            print(f"  • {name:<20}: {value:+.5f}")
+
+        metrics = env.get_current_robot_metrics()
+        dist = metrics[0]["distance_from_start"] if metrics else 0.0
+        speed = metrics[0]["speed"] if metrics else 0.0
+        avg_reward = total_reward / max(1, step_count)
+
+        print(f"  ├─ Average Reward / Step: {avg_reward:.4f}")
+        print(f"  ├─ Distance Travelled: {dist:.2f} m")
+        print(f"  ├─ Actual Avg Speed: {speed:.2f} m/s")
+        print(f"  └─ Steps Survived: {step_count} / {episode_length}")
+        print("-" * 70)
+
+    env.close()
+
+if __name__ == "__main__":
+    evaluate_kinematics()

ml/run_train.py

@@ -77,7 +77,7 @@ def main():
     model = PPO(
         policy="MlpPolicy",
         env=vec_env,
-        learning_rate=3e-4,
+        learning_rate=1e-4,
         n_steps=256,
         batch_size=256,
         n_epochs=10,