env cleanup

rewards adjustment
pybullet logic contained in SimManager

ml/SimManager.py

@@ -1,9 +1,10 @@
 """
 ml/SimManager.py - PyBullet Simulation & Multi-Body Manager
 """
-from typing import Dict, List, Tuple
+from typing import Dict, List, Tuple, Optional
 import pybullet as p
 import pybullet_data
+import numpy as np
 import DataTypes as dt
 
 
@@ -84,4 +85,82 @@ class SimManager:
     def disconnect(self):
         if self.physics_client is not None and p.isConnected(self.physics_client):
             p.disconnect(self.physics_client)
-            self.physics_client = None
+            self.physics_client = None
+
+    # --- ROBOT GETTERS AND SETTERS ---
+
+    def reset_robot_base(
+        self,
+        body_id: int,
+        position: List[float],
+        orientation: Optional[List[float]] = None,
+        linear_velocity: Optional[List[float]] = None,
+        angular_velocity: Optional[List[float]] = None
+    ) -> None:
+        """Resets a robot body's base position, orientation, and velocities."""
+        if orientation is None:
+            orientation = [0.0, 0.0, 0.0, 1.0]
+        if linear_velocity is None:
+            linear_velocity = [0.0, 0.0, 0.0]
+        if angular_velocity is None:
+            angular_velocity = [0.0, 0.0, 0.0]
+
+        p.resetBasePositionAndOrientation(
+            body_id, position, orientation, physicsClientId=self.physics_client
+        )
+        p.resetBaseVelocity(
+            body_id, linearVelocity=linear_velocity, angularVelocity=angular_velocity,
+            physicsClientId=self.physics_client
+        )
+
+    def get_robot_pose(self, body_id: int) -> 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)
+        return list(pos), list(orn)
+
+    def get_robot_rpy(self, body_id: int) -> Tuple[float, float, float]:
+        """Returns roll, pitch, yaw angles in radians for the given robot body."""
+        _, orn = p.getBasePositionAndOrientation(body_id, physicsClientId=self.physics_client)
+        roll, pitch, yaw = p.getEulerFromQuaternion(orn)
+        return float(roll), float(pitch), float(yaw)
+
+    def get_robot_pose_and_rpy(self, body_id: int) -> Tuple[List[float], Tuple[float, float, float]]:
+        """Returns base position and (roll, pitch, yaw) tuple in radians."""
+        pos, orn = p.getBasePositionAndOrientation(body_id, physicsClientId=self.physics_client)
+        roll, pitch, yaw = p.getEulerFromQuaternion(orn)
+        return list(pos), (float(roll), float(pitch), float(yaw))
+
+    def get_robot_velocity(self, body_id: int) -> Tuple[List[float], List[float]]:
+        """Returns linear velocity (vx, vy, vz) and angular velocity (wx, wy, wz)."""
+        lin_v, ang_v = p.getBaseVelocity(body_id, physicsClientId=self.physics_client)
+        return list(lin_v), list(ang_v)
+
+    def get_robot_joint_angles(self, body_id: int, joint_indices: Optional[List[int]] = None) -> np.ndarray:
+        """Returns joint angles as a 1D numpy array float32 for specified or registered joint indices."""
+        if joint_indices is None:
+            joint_indices = self.robot_joints.get(body_id, list(range(18)))
+        joint_states = p.getJointStates(body_id, joint_indices, physicsClientId=self.physics_client)
+        return np.array([state[0] for state in joint_states], dtype=np.float32)
+
+    def set_robot_color(self, body_id: int, rgba: List[float]) -> None:
+        """Changes visual color RGBA of base link and all joints of the specified robot body."""
+        num_joints = p.getNumJoints(body_id, physicsClientId=self.physics_client)
+        p.changeVisualShape(body_id, -1, rgbaColor=rgba, physicsClientId=self.physics_client)
+        for j in range(num_joints):
+            p.changeVisualShape(body_id, j, rgbaColor=rgba, physicsClientId=self.physics_client)
+
+    def measure_robot_heights(self, robot_ids: List[int]) -> List[float]:
+        """Gets current Z height for all specified robot body IDs."""
+        heights = []
+        for body_id in robot_ids:
+            pos, _ = p.getBasePositionAndOrientation(body_id, physicsClientId=self.physics_client)
+            heights.append(pos[2])
+        return heights
+
+    def settle_and_measure_height(self, robot_ids: List[int], steps: int = 200, fallback_height: float = 0.14) -> float:
+        """Steps simulation for designated steps so robot settles, then calculates target standing height."""
+        for _ in range(steps):
+            self.step()
+        heights = self.measure_robot_heights(robot_ids)
+        mean_height = float(np.mean(heights)) if heights else fallback_height
+        return mean_height if mean_height > 0.0 else fallback_height