From 5448335b114a9d5d99ab9fe248612f9e358d7a98 Mon Sep 17 00:00:00 2001 From: JackM323 Date: Thu, 30 Jul 2026 17:35:12 +0200 Subject: [PATCH] Text Metrics for Testing in the Training --- ml/env.py | 191 +++++++++++++++++----------------------------- ml/sim_manager.py | 2 +- 2 files changed, 69 insertions(+), 124 deletions(-) diff --git a/ml/env.py b/ml/env.py index f4f3830..404bddc 100644 --- a/ml/env.py +++ b/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() \ No newline at end of file diff --git a/ml/sim_manager.py b/ml/sim_manager.py index fc20567..f056647 100644 --- a/ml/sim_manager.py +++ b/ml/sim_manager.py @@ -30,7 +30,7 @@ class SimManager: p.configureDebugVisualizer(p.COV_ENABLE_GUI, 0) p.configureDebugVisualizer(p.COV_ENABLE_SEGMENTATION_MARK_PREVIEW, 0) p.configureDebugVisualizer(p.COV_ENABLE_DEPTH_BUFFER_PREVIEW, 0) - p.configureDebugVisualizer(p.COV_ENABLE_RGB_OUTPUT_PREVIEW, 0) + p.configureDebugVisualizer(p.COV_ENABLE_RGB_BUFFER_PREVIEW, 0) def load_scene( self, urdf_path: str, num_robots: int, robot_spacing: float, base_pos_fn