JackM323
c5ca79a354
Training environment to make a walk model for the hexapod generated code that will be checked
318 lines
12 KiB
Python
318 lines
12 KiB
Python
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
|
|
|
|
|
|
class JackBotEnv(gym.Env):
|
|
"""Gymnasium environment for joint-command learning with the JackBot URDF."""
|
|
|
|
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",
|
|
):
|
|
"""
|
|
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
|
|
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)
|
|
self.physics_client = None
|
|
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.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 _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)
|
|
|
|
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 = info[8]
|
|
upper = info[9]
|
|
if lower >= upper:
|
|
lower = -math.pi
|
|
upper = 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):
|
|
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.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()
|
|
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)
|
|
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()
|
|
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
|
|
|
|
return observation, float(reward), terminated, truncated, info
|
|
|
|
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
|
|
|
|
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]
|
|
|
|
def close(self):
|
|
if self.physics_client is not None:
|
|
p.disconnect(self.physics_client)
|
|
self.physics_client = None
|