import math import random import numpy as np import pybullet as p 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 using SimManager for GUI and simulation control.""" metadata = {"render_modes": ["human", "rgb_array"]} def __init__( self, urdf_path: str | None = None, use_gui: bool = True, frame_skip: int = 4, random_command: bool = True, max_episode_steps: int = 2000, num_robots: int = 1, robot_spacing: float = 0.5, start_pose: str = "init_deg", ): self.urdf_path = urdf_path or cfg.urdf_path self.use_gui = use_gui self.frame_skip = frame_skip self.random_command = random_command self.max_episode_steps = max_episode_steps self.num_robots = max(1, num_robots) self.robot_spacing = robot_spacing self.start_pose = start_pose self.action_scale = math.radians(8.0) # Delegate physics simulation & GUI management self.sim_manager = SimManager(use_gui=self.use_gui) self.robots = [] self.plane = None self.robot_joint_indices = [] self.joint_lower = None self.joint_upper = None 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 self._connect_sim() total_joints = self.num_robots * len(self.robot_joint_indices[0]) observation_dim = total_joints + self.num_robots * 4 action_dim = total_joints self.observation_space = spaces.Box( 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 ) 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): 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 = [], [] for joint_index in self.robot_joint_indices[0]: info = p.getJointInfo(self.robots[0], joint_index) lower, upper = info[8], info[9] if lower >= upper: lower, upper = -math.pi, math.pi lower_limits.append(lower) upper_limits.append(upper) limits = np.array([lower_limits, upper_limits], dtype=np.float32) self.joint_lower = np.tile(limits[0], (self.num_robots, 1)) self.joint_upper = np.tile(limits[1], (self.num_robots, 1)) 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.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): 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: 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._update_gui_hud(reward=0.0) return self._get_obs(), {} def seed(self, seed: int | None = None): if seed is None: seed = int(np.random.randint(0, 2**31 - 1)) random.seed(seed) np.random.seed(seed) self._seed = seed return [seed] def _reset_pose(self): for robot, joint_indices, init_angles in zip(self.robots, self.robot_joint_indices, self.initial_angles): for joint_index, target_angle in zip(joint_indices, init_angles): p.resetJointState(robot, joint_index, target_angle) p.setJointMotorControl2( bodyIndex=robot, jointIndex=joint_index, controlMode=p.POSITION_CONTROL, targetPosition=target_angle, force=250, ) for _ in range(400): p.stepSimulation() self.joint_angles = self.initial_angles.copy() def sample_command(self) -> np.ndarray: vx = np.random.uniform(-1.0, 1.0) vy = np.random.uniform(-0.5, 0.5) vz = 0.0 omega = np.random.uniform(-1.0, 1.0) return np.array([vx, vy, vz, omega], dtype=np.float32) def _get_obs(self) -> np.ndarray: return np.concatenate([self.joint_angles.flatten(), self.commands.flatten()]).astype(np.float32) def step(self, action: np.ndarray): action = np.clip(action, self.action_space.low, self.action_space.high).astype(np.float32) self.last_action = action action_matrix = action.reshape(self.num_robots, -1) self.joint_angles = np.clip( self.joint_angles + action_matrix * self.action_scale, self.joint_lower, self.joint_upper, ) for robot, joint_indices, angles in zip(self.robots, self.robot_joint_indices, self.joint_angles): for joint_index, target_angle in zip(joint_indices, angles): p.setJointMotorControl2( bodyIndex=robot, jointIndex=joint_index, controlMode=p.POSITION_CONTROL, targetPosition=target_angle, force=250, ) for _ in range(self.frame_skip): p.stepSimulation() 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 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): linear_vel, angular_vel = p.getBaseVelocity(robot) _, orientation = p.getBasePositionAndOrientation(robot) roll, pitch, _ = p.getEulerFromQuaternion(orientation) command = self.commands[robot_id] forward_reward = command[0] * linear_vel[0] + command[1] * linear_vel[1] rotation_reward = command[3] * angular_vel[2] stability_penalty = abs(roll) + abs(pitch) action_penalty = float(np.sum(np.square(self.last_action.reshape(self.num_robots, -1)[robot_id]))) * 0.01 rewards.append(0.1 + forward_reward + rotation_reward - 0.2 * stability_penalty - action_penalty) return float(np.sum(rewards)) def _is_done(self) -> bool: for robot in self.robots: _, orientation = p.getBasePositionAndOrientation(robot) roll, pitch, _ = p.getEulerFromQuaternion(orientation) if abs(roll) > 0.7 or abs(pitch) > 0.7: return True position, _ = p.getBasePositionAndOrientation(robot) if position[2] < 0.02: return True return self.step_count >= self.max_episode_steps def close(self): self.sim_manager.disconnect()