Training uses new Robot.py

new Metrics and SimManager for live training viewing
Robot.py usage added to machinelearning

Robot.py

@@ -6,6 +6,7 @@ Handles state, kinematics, backends (Hardware/Simulation), and motion execution.
 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
@@ -49,20 +50,25 @@ class HardwareBackend:
 
 class PyBulletBackend:
     """Backend for PyBullet simulation execution."""
-    def __init__(self, sim_instance):
+    def __init__(self, sim_instance, body_id: Optional[int] = None):
         self.sim = sim_instance
+        self.body_id = body_id
 
     def send_angles(self, rad_array: dt.RadArray) -> None:
         if self.sim:
-            self.sim.updatePos(rad_array)
+            # If body_id is set, target that specific robot body
+            if self.body_id is not None and hasattr(self.sim, 'updatePosForBody'):
+                self.sim.updatePosForBody(self.body_id, rad_array)
+            else:
+                self.sim.updatePos(rad_array)
 
     def step_simulation(self) -> None:
         if self.sim:
             self.sim.step()
 
     def cleanup(self) -> None:
-        if self.sim and hasattr(self.sim, 'close'):
-            self.sim.close()
+        if self.sim and hasattr(self.sim, 'disconnect'):
+            self.sim.disconnect()
 
 
 class Robot:
@@ -102,6 +108,9 @@ class Robot:
             ri.get_center_points() if hasattr(ri, "get_center_points") else ri.center_points
         )
 
+        # RL configuration
+        self.action_scale = 0.1  # Joint delta step size (radians)
+
         # Gait / motion variables
         self.leg_state = np.array(["step", "drag", "step", "drag", "step", "drag"])
         self.robot_state = "idle"
@@ -158,6 +167,51 @@ class Robot:
             self.transition_to(next_state_key)
         self.step_sim()
 
+    # --- RL METHODS ---
+
+    def apply_rl_action(self, action: np.ndarray) -> None:
+        """
+        Applies continuous RL action deltas [-1, 1] to current joint angles.
+        """
+        action = np.asarray(action, dtype=np.float32)
+        scaled_action = np.clip(action, -1.0, 1.0) * self.action_scale
+
+        current_flat = self.current_rad.data.flatten()
+        updated_flat = np.clip(
+            current_flat + scaled_action,
+            -np.pi / 2,
+            np.pi / 2
+        )
+
+        new_rad = dt.RadArray(data=updated_flat.reshape(self.current_rad.data.shape))
+        self.set_joint_angles(new_rad)
+
+    def get_observation(self, command: Optional[np.ndarray] = None) -> np.ndarray:
+        """
+        Returns observation vector [18 joint angles] + [optional 4 command dimensions].
+        Queries PyBullet if backend is PyBulletBackend; otherwise falls back to internal state.
+        """
+        if isinstance(self.backend, PyBulletBackend) and self.backend.sim and hasattr(self.backend.sim, 'physics_client'):
+            physics_client = self.backend.sim.physics_client
+            body_id = self.backend.body_id if self.backend.body_id is not None else 0
+            
+            # Retrieve joint mapping from SimManager/Simulation if available
+            if hasattr(self.backend.sim, 'robot_joints') and body_id in self.backend.sim.robot_joints:
+                joint_indices = self.backend.sim.robot_joints[body_id]
+            else:
+                joint_indices = list(range(18))
+
+            joint_states = p.getJointStates(body_id, joint_indices, physicsClientId=physics_client)
+            joint_angles = np.array([state[0] for state in joint_states], dtype=np.float32)
+        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.astype(np.float32)
+
     def cleanup(self) -> None:
         if self.backend and hasattr(self.backend, 'cleanup'):
             self.backend.cleanup()