Complete Restructered Robot Code
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
This commit is contained in:
@@ -1,71 +1,33 @@
|
||||
import math
|
||||
import random
|
||||
import numpy as np
|
||||
import pybullet as p
|
||||
# ml/env.py
|
||||
import gymnasium as gym
|
||||
from gymnasium import spaces
|
||||
import numpy as np
|
||||
import math
|
||||
|
||||
import config as cfg
|
||||
import robot_init as ri
|
||||
from .sim_manager import SimManager
|
||||
|
||||
from Robot import Robot, PyBulletBackend
|
||||
from ml.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 __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
|
||||
@@ -175,44 +137,24 @@ class JackBotEnv(gym.Env):
|
||||
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,
|
||||
)
|
||||
action_per_robot = action.reshape(len(self.robots), 18)
|
||||
|
||||
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,
|
||||
)
|
||||
# Apply RL actions independently to each Robot object instance
|
||||
for robot, act in zip(self.robots, action_per_robot):
|
||||
robot.apply_rl_action(act)
|
||||
|
||||
for _ in range(self.frame_skip):
|
||||
p.stepSimulation()
|
||||
# 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
|
||||
])
|
||||
|
||||
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
|
||||
done = False
|
||||
return obs, reward, done, False, {}
|
||||
|
||||
def _update_gui_hud(self, reward: float):
|
||||
"""Passes current state metrics to the SimManager HUD renderer."""
|
||||
|
||||
Reference in New Issue
Block a user