JackM323
b537677277
Robot into its own Class instead of lose Global Variables that cause circular imports StateClass usage instead of the old RobotState.py New Input Class for Controller and randome intputs
204 lines
8.1 KiB
Python
204 lines
8.1 KiB
Python
# ml/env.py
|
|
import gymnasium as gym
|
|
from gymnasium import spaces
|
|
import numpy as np
|
|
import math
|
|
|
|
from Robot import Robot, PyBulletBackend
|
|
from ml.sim_manager import SimManager
|
|
|
|
class JackBotEnv(gym.Env):
|
|
def __init__(self, use_gui: bool = False, num_robots: int = 1):
|
|
self.sim_manager = SimManager(use_gui=use_gui)
|
|
self.sim_manager.connect()
|
|
|
|
# Load simulation bodies
|
|
self.plane, self.pb_robots, self.joint_indices = self.sim_manager.load_scene(
|
|
cfg.urdf_path, num_robots=num_robots
|
|
)
|
|
|
|
# Instantiate dedicated Robot Python object for EACH spawned robot
|
|
self.robots = [
|
|
Robot(backend=PyBulletBackend(self.sim_manager, body_id=pb_id))
|
|
for pb_id in self.pb_robots
|
|
]
|
|
|
|
# Action & Observation Spaces
|
|
action_dim = num_robots * 18
|
|
obs_dim = num_robots * (18 + 4)
|
|
self.action_space = spaces.Box(-1.0, 1.0, shape=(action_dim,), dtype=np.float32)
|
|
self.observation_space = spaces.Box(-np.inf, np.inf, shape=(obs_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_per_robot = action.reshape(len(self.robots), 18)
|
|
|
|
# Apply RL actions independently to each Robot object instance
|
|
for robot, act in zip(self.robots, action_per_robot):
|
|
robot.apply_rl_action(act)
|
|
|
|
# Step PyBullet physics engine once
|
|
self.sim_manager.step()
|
|
|
|
# Gather observations across all robot objects
|
|
obs = np.concatenate([
|
|
robot.get_observation(command=np.zeros(4))
|
|
for robot in self.robots
|
|
])
|
|
|
|
reward = self._compute_reward()
|
|
done = False
|
|
return obs, reward, done, False, {}
|
|
|
|
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() |