Text Metrics for Testing in the Training

ml/env.py

@@ -1,17 +1,17 @@
 import math
-import os
 import random
 import numpy as np
 import pybullet as p
-import pybullet_data
 import gymnasium as gym
 from gymnasium import spaces
+
 import config as cfg
 import robot_init as ri
+from .sim_manager import SimManager
 
 
 class JackBotEnv(gym.Env):
-    """Gymnasium environment for joint-command learning with the JackBot URDF."""
+    """Gymnasium environment for joint-command learning using SimManager for GUI and simulation control."""
 
     metadata = {"render_modes": ["human", "rgb_array"]}
 
@@ -26,20 +26,6 @@ class JackBotEnv(gym.Env):
         robot_spacing: float = 0.5,
         start_pose: str = "init_deg",
     ):
-        """
-        Initialize the JackBot environment.
-
-        Args:
-            urdf_path: Path to the robot's URDF file. Defaults to config value.
-            use_gui: If True, uses p.GUI (shows window), else p.DIRECT (headless).
-            frame_skip: Number of physics simulation steps per agent step. 
-                Higher values mean the agent takes actions less frequently (e.g., 4 steps at 240Hz = 60Hz control).
-            random_command: If True, assigns a random velocity command (vx, vy, omega) to each robot every reset.
-            max_episode_steps: Maximum number of steps (interactions) before the episode is truncated.
-            num_robots: Number of robots to spawn for parallel training in the same scene.
-            robot_spacing: Distance in meters between robots when spawning multiple.
-            start_pose: The initial joint configuration name (e.g., 'init_deg' or 'init90_deg').
-        """
         self.urdf_path = urdf_path or cfg.urdf_path
         self.use_gui = use_gui
         self.frame_skip = frame_skip
@@ -49,7 +35,10 @@ class JackBotEnv(gym.Env):
         self.robot_spacing = robot_spacing
         self.start_pose = start_pose
         self.action_scale = math.radians(8.0)
-        self.physics_client = None
+
+        # Delegate physics simulation & GUI management
+        self.sim_manager = SimManager(use_gui=self.use_gui)
+
         self.robots = []
         self.plane = None
         self.robot_joint_indices = []
@@ -58,6 +47,8 @@ class JackBotEnv(gym.Env):
         self.initial_angles = None
         self.commands = None
         self.step_count = 0
+        self.episode_count = 0
+        self.cumulative_reward = 0.0
         self.last_action = None
         self._first_reset = True
         self._min_steps_before_done = 150
@@ -69,55 +60,37 @@ class JackBotEnv(gym.Env):
         action_dim = total_joints
 
         self.observation_space = spaces.Box(
-            low=-np.inf,
-            high=np.inf,
-            shape=(observation_dim,),
-            dtype=np.float32,
+            low=-np.inf, high=np.inf, shape=(observation_dim,), dtype=np.float32
         )
         self.action_space = spaces.Box(
-            low=-1.0,
-            high=1.0,
-            shape=(action_dim,),
-            dtype=np.float32,
+            low=-1.0, high=1.0, shape=(action_dim,), dtype=np.float32
         )
 
+    def _robot_base_position(self, robot_id: int, num_robots: int = 1, spacing: float = 0.5) -> list[float]:
+            """Calculates grid coordinates for spawning multiple robots in PyBullet."""
+            cols = int(math.sqrt(num_robots - 1)) + 1
+            row = robot_id // cols
+            col = robot_id % cols
+            x = (col - (cols - 1) / 2.0) * spacing
+            y = (row - (cols - 1) / 2.0) * spacing
+            return [x, y, 0.1]
+
     def _connect_sim(self):
-        if self.physics_client is not None:
-            p.disconnect(self.physics_client)
-
-        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)
-
-        self.plane = p.loadURDF("plane.urdf")
-        self.robots = []
-        self.robot_joint_indices = []
-
-        for robot_id in range(self.num_robots):
-            base_pos = self._robot_base_position(robot_id)
-            robot = p.loadURDF(self.urdf_path, basePosition=base_pos, useFixedBase=False)
-            self.robots.append(robot)
-            joint_indices = [
-                i
-                for i in range(p.getNumJoints(robot))
-                if p.getJointInfo(robot, i)[2] == p.JOINT_REVOLUTE
-            ]
-            self.robot_joint_indices.append(joint_indices)
+        self.sim_manager.connect()
+        self.plane, self.robots, self.robot_joint_indices = self.sim_manager.load_scene(
+            self.urdf_path, self.num_robots, self.robot_spacing, self._robot_base_position
+        )
 
         joint_count = len(self.robot_joint_indices[0])
         if any(len(indices) != joint_count for indices in self.robot_joint_indices):
             raise ValueError("All robots must have the same number of revolute joints.")
 
-        lower_limits = []
-        upper_limits = []
+        lower_limits, upper_limits = [], []
         for joint_index in self.robot_joint_indices[0]:
             info = p.getJointInfo(self.robots[0], joint_index)
-            lower = info[8]
-            upper = info[9]
+            lower, upper = info[8], info[9]
             if lower >= upper:
-                lower = -math.pi
-                upper = math.pi
+                lower, upper = -math.pi, math.pi
             lower_limits.append(lower)
             upper_limits.append(upper)
 
@@ -128,53 +101,27 @@ class JackBotEnv(gym.Env):
         pose = ri.init_deg if self.start_pose == "init_deg" else ri.init90_deg
         initial_pose = pose.to_rad().data.flatten()
         self.initial_angles = np.tile(initial_pose, (self.num_robots, 1)).astype(np.float32)
-        self.initial_angles = np.clip(
-            self.initial_angles,
-            self.joint_lower,
-            self.joint_upper,
-        )
+        self.initial_angles = np.clip(self.initial_angles, self.joint_lower, self.joint_upper)
 
         self.last_action = np.zeros(self.num_robots * joint_count, dtype=np.float32)
         self.commands = np.zeros((self.num_robots, 4), dtype=np.float32)
 
     def reset(self, command: np.ndarray | None = None, seed: int | None = None, options: dict | None = None):
-        if self._first_reset:
-            self.step_count = 0
-            self.last_action = np.zeros(self.num_robots * len(self.robot_joint_indices[0]), dtype=np.float32)
-
-            if seed is not None:
-                self.seed(seed)
-
-            if command is not None:
-                self.commands = np.array(command, dtype=np.float32).reshape(self.num_robots, 4)
-            elif self.random_command:
-                self.commands = np.stack([self.sample_command() for _ in range(self.num_robots)])
-            else:
-                self.commands = np.zeros((self.num_robots, 4), dtype=np.float32)
-
-            self._reset_pose()
-            self._first_reset = False
-            return self._get_obs(), {}
-
-        p.resetSimulation()
-        p.setGravity(0, 0, -9.81)
-        p.setAdditionalSearchPath(pybullet_data.getDataPath())
-        self.plane = p.loadURDF("plane.urdf")
-        self.robots = []
-        self.robot_joint_indices = []
-
-        for robot_id in range(self.num_robots):
-            base_pos = self._robot_base_position(robot_id)
-            robot = p.loadURDF(self.urdf_path, basePosition=base_pos, useFixedBase=False)
-            self.robots.append(robot)
-            joint_indices = [
-                i
-                for i in range(p.getNumJoints(robot))
-                if p.getJointInfo(robot, i)[2] == p.JOINT_REVOLUTE
-            ]
-            self.robot_joint_indices.append(joint_indices)
-
+        self.episode_count += 1
         self.step_count = 0
+        self.cumulative_reward = 0.0
+
+        if not self._first_reset:
+            # Reset simulation bodies without reconnecting PyBullet
+            p.resetSimulation()
+            p.setGravity(0, 0, -9.81)
+            self.sim_manager.reset_hud()
+            self.plane, self.robots, self.robot_joint_indices = self.sim_manager.load_scene(
+                self.urdf_path, self.num_robots, self.robot_spacing, self._robot_base_position
+            )
+        else:
+            self._first_reset = False
+
         self.last_action = np.zeros(self.num_robots * len(self.robot_joint_indices[0]), dtype=np.float32)
 
         if seed is not None:
@@ -188,10 +135,11 @@ class JackBotEnv(gym.Env):
             self.commands = np.zeros((self.num_robots, 4), dtype=np.float32)
 
         self._reset_pose()
+        self._update_gui_hud(reward=0.0)
+
         return self._get_obs(), {}
 
     def seed(self, seed: int | None = None):
-        """Set RNG seed for reproducibility (Gym compatibility)."""
         if seed is None:
             seed = int(np.random.randint(0, 2**31 - 1))
         random.seed(seed)
@@ -252,21 +200,35 @@ class JackBotEnv(gym.Env):
         self.step_count += 1
         observation = self._get_obs()
         reward = self._compute_reward()
+        self.cumulative_reward += reward
+
         done = self._is_done()
         if self.step_count <= self._min_steps_before_done:
             done = False
-        info = {"step": self.step_count}
 
-        # Gymnasium-style return: (obs, reward, terminated, truncated, info)
-        if done:
-            terminated = True
-            truncated = False
-        else:
-            terminated = False
-            truncated = False
+        self._update_gui_hud(reward=reward)
+
+        terminated = done
+        truncated = False
+        info = {"step": self.step_count, "episode_reward": self.cumulative_reward}
 
         return observation, float(reward), terminated, truncated, info
 
+    def _update_gui_hud(self, reward: float):
+        """Passes current state metrics to the SimManager HUD renderer."""
+        linear_vel, _ = p.getBaseVelocity(self.robots[0])
+        cmd = self.commands[0]
+
+        stats = {
+            "Episode": self.episode_count,
+            "Step": f"{self.step_count} / {self.max_episode_steps}",
+            "Step Reward": reward,
+            "Total Reward": self.cumulative_reward,
+            "Target Cmd (vx,vy,w)": [cmd[0], cmd[1], cmd[3]],
+            "Actual Vel (vx,vy)": [linear_vel[0], linear_vel[1]],
+        }
+        self.sim_manager.update_hud(stats)
+
     def _compute_reward(self) -> float:
         rewards = []
         for robot_id, robot in enumerate(self.robots):
@@ -294,24 +256,7 @@ class JackBotEnv(gym.Env):
             if position[2] < 0.02:
                 return True
 
-        if self.step_count >= self.max_episode_steps:
-            return True
-        return False
-
-    def render(self, mode="human"):
-        if self.use_gui:
-            return None
-        raise NotImplementedError("Render is only supported with use_gui=True")
-
-    def _robot_base_position(self, robot_id: int) -> list[float]:
-        cols = int(math.sqrt(self.num_robots - 1)) + 1
-        row = robot_id // cols
-        col = robot_id % cols
-        x = (col - (cols - 1) / 2.0) * self.robot_spacing
-        y = (row - (cols - 1) / 2.0) * self.robot_spacing
-        return [x, y, 0.1]
+        return self.step_count >= self.max_episode_steps
 
     def close(self):
-        if self.physics_client is not None:
-            p.disconnect(self.physics_client)
-            self.physics_client = None
+        self.sim_manager.disconnect()