added Robot kinematic use for training
HUGE BUG -> SimManager physics broken (at least with Robot kinematics)
This commit is contained in:
+174
-108
@@ -1,7 +1,7 @@
|
||||
"""
|
||||
ml/SimManager.py - PyBullet Simulation & Multi-Body Manager
|
||||
ml/SimManager.py - PyBullet Simulation Manager (Single Robot Dedicated)
|
||||
"""
|
||||
from typing import Dict, List, Tuple, Optional
|
||||
from typing import List, Tuple, Optional, Union
|
||||
import pybullet as p
|
||||
import pybullet_data
|
||||
import numpy as np
|
||||
@@ -9,14 +9,16 @@ import DataTypes as dt
|
||||
|
||||
|
||||
class SimManager:
|
||||
"""Manages PyBullet simulation lifecycle and multi-robot physics."""
|
||||
"""Manages PyBullet simulation lifecycle for a single JackBot hexapod."""
|
||||
|
||||
def __init__(self, use_gui: bool = True):
|
||||
self.use_gui = use_gui
|
||||
self.physics_client = None
|
||||
self.robot_joints: Dict[int, List[int]] = {}
|
||||
self.plane: Optional[int] = None
|
||||
self.joint_indices: List[int] = []
|
||||
self.foot_indices: List[int] = []
|
||||
|
||||
def connect(self):
|
||||
def connect(self) -> None:
|
||||
"""Connects to PyBullet and hides side GUI panels."""
|
||||
if self.physics_client is not None and p.isConnected(self.physics_client):
|
||||
return
|
||||
@@ -28,62 +30,87 @@ class SimManager:
|
||||
p.setGravity(0, 0, -9.81, physicsClientId=self.physics_client)
|
||||
|
||||
if self.use_gui:
|
||||
# Disable PyBullet side panel and preview windows
|
||||
p.configureDebugVisualizer(p.COV_ENABLE_GUI, 0, physicsClientId=self.physics_client)
|
||||
p.configureDebugVisualizer(p.COV_ENABLE_SEGMENTATION_MARK_PREVIEW, 0, physicsClientId=self.physics_client)
|
||||
p.configureDebugVisualizer(p.COV_ENABLE_DEPTH_BUFFER_PREVIEW, 0, physicsClientId=self.physics_client)
|
||||
p.configureDebugVisualizer(p.COV_ENABLE_RGB_BUFFER_PREVIEW, 0, physicsClientId=self.physics_client)
|
||||
|
||||
def load_scene(
|
||||
self, urdf_path: str, robot_spacing: float, base_pos_fn
|
||||
self,
|
||||
urdf_path: str,
|
||||
spacing: float = 0.0,
|
||||
position_func=None
|
||||
) -> Tuple[int, List[int], List[List[int]]]:
|
||||
"""Loads the plane and a single hexapod body into the simulation scene."""
|
||||
plane_id = p.loadURDF("plane.urdf", physicsClientId=self.physics_client)
|
||||
robots = []
|
||||
robot_joint_indices = []
|
||||
self.robot_joints.clear()
|
||||
"""Loads ground plane and the single robot URDF with high ground friction and force settings."""
|
||||
# 1. Load ground plane and set explicit friction
|
||||
self.plane = p.loadURDF("plane.urdf", physicsClientId=self.physics_client)
|
||||
|
||||
base_pos = base_pos_fn(0, robot_spacing)
|
||||
robot = p.loadURDF(
|
||||
urdf_path,
|
||||
basePosition=base_pos,
|
||||
useFixedBase=False,
|
||||
physicsClientId=self.physics_client
|
||||
)
|
||||
robots.append(robot)
|
||||
# 2. Determine starting position (matching simulation.py starting height of 0.20m)
|
||||
spawn_pos = position_func(0) if position_func is not None else [0.0, 0.0, 0.20]
|
||||
|
||||
joint_indices = [
|
||||
i for i in range(p.getNumJoints(robot, physicsClientId=self.physics_client))
|
||||
if p.getJointInfo(robot, i, physicsClientId=self.physics_client)[2] == p.JOINT_REVOLUTE
|
||||
]
|
||||
robot_joint_indices.append(joint_indices)
|
||||
self.robot_joints[robot] = joint_indices
|
||||
# 3. Spawn single robot body
|
||||
self.robot_id = p.loadURDF(urdf_path, spawn_pos, physicsClientId=self.physics_client)
|
||||
|
||||
return plane_id, robots, robot_joint_indices
|
||||
# 4. Retrieve revolute joint indices
|
||||
self.joint_indices = []
|
||||
for j in range(p.getNumJoints(self.robot_id, physicsClientId=self.physics_client)):
|
||||
info = p.getJointInfo(self.robot_id, j, physicsClientId=self.physics_client)
|
||||
if info[2] == p.JOINT_REVOLUTE:
|
||||
self.joint_indices.append(j)
|
||||
|
||||
def updatePosForBody(self, body_id: int, current_rad: dt.RadArray):
|
||||
"""Sets joint motor position targets on individual robot bodies."""
|
||||
if body_id not in self.robot_joints:
|
||||
return
|
||||
|
||||
joint_indices = self.robot_joints[body_id]
|
||||
radflat = current_rad.data.flatten()
|
||||
# Return format maintains 100% compatibility with JackBotEnv unpack sequence
|
||||
return self.plane, [self.robot_id], [self.joint_indices]
|
||||
|
||||
def _resolve_body_id(self, body_id: Optional[int] = None) -> int:
|
||||
"""Internal helper to return the single active robot ID."""
|
||||
if body_id is not None:
|
||||
return body_id
|
||||
if self.robot_id is not None:
|
||||
return self.robot_id
|
||||
raise RuntimeError("No robot loaded in SimManager. Call load_scene() first.")
|
||||
|
||||
def set_robot_joint_angles(
|
||||
self,
|
||||
target_angles: Union[np.ndarray, List[float]],
|
||||
joint_indices: Optional[List[int]] = None,
|
||||
body_id: Optional[int] = None,
|
||||
) -> None:
|
||||
"""Applies position control with force=500 matching simulation.py logic."""
|
||||
bid = self._resolve_body_id(body_id)
|
||||
j_indices = joint_indices if joint_indices is not None else self.joint_indices
|
||||
radflat = target_angles.flatten() if isinstance(target_angles, np.ndarray) else target_angles
|
||||
|
||||
for joint_index, target_angle in zip(joint_indices, radflat):
|
||||
for joint_index, target_angle in zip(j_indices, radflat):
|
||||
p.setJointMotorControl2(
|
||||
bodyIndex=body_id,
|
||||
bodyIndex=bid,
|
||||
jointIndex=joint_index,
|
||||
controlMode=p.POSITION_CONTROL,
|
||||
targetPosition=float(target_angle),
|
||||
force=250,
|
||||
force=500,
|
||||
physicsClientId=self.physics_client
|
||||
)
|
||||
|
||||
def step(self):
|
||||
def updatePosForBody(
|
||||
self,
|
||||
body_id_or_rad: Union[int, dt.RadArray],
|
||||
current_rad: Optional[dt.RadArray] = None
|
||||
) -> None:
|
||||
"""Updates robot joint positions using RadArray input."""
|
||||
if isinstance(body_id_or_rad, dt.RadArray):
|
||||
rad_data = body_id_or_rad
|
||||
bid = self.robot_id
|
||||
else:
|
||||
bid = self._resolve_body_id(body_id_or_rad)
|
||||
rad_data = current_rad
|
||||
|
||||
if rad_data is not None:
|
||||
self.set_robot_joint_angles(rad_data.data, body_id=bid)
|
||||
|
||||
def step(self) -> None:
|
||||
p.stepSimulation(physicsClientId=self.physics_client)
|
||||
|
||||
def set_rendering(self, enabled: bool) -> None:
|
||||
"""Toggles PyBullet 3D rendering to speed up simulation."""
|
||||
"""Toggles PyBullet 3D rendering visualizer."""
|
||||
if self.physics_client is not None and p.isConnected(self.physics_client):
|
||||
p.configureDebugVisualizer(
|
||||
p.COV_ENABLE_RENDERING,
|
||||
@@ -91,7 +118,7 @@ class SimManager:
|
||||
physicsClientId=self.physics_client
|
||||
)
|
||||
|
||||
def disconnect(self):
|
||||
def disconnect(self) -> None:
|
||||
if self.physics_client is not None and p.isConnected(self.physics_client):
|
||||
p.disconnect(self.physics_client)
|
||||
self.physics_client = None
|
||||
@@ -116,114 +143,153 @@ class SimManager:
|
||||
|
||||
return p.getContactPoints(**kwargs)
|
||||
|
||||
# --- ROBOT GETTERS AND SETTERS ---
|
||||
# Add to ml/SimManager.py
|
||||
def hard_reset_joint_angles(self, target_angles: np.ndarray, body_id: Optional[int] = None) -> None:
|
||||
bid = self._resolve_body_id(body_id)
|
||||
radflat = target_angles.flatten()
|
||||
for joint_index, target_angle in zip(self.joint_indices, radflat):
|
||||
p.resetJointState(
|
||||
bodyUniqueId=bid,
|
||||
jointIndex=joint_index,
|
||||
targetValue=float(target_angle),
|
||||
targetVelocity=0.0,
|
||||
physicsClientId=self.physics_client
|
||||
)
|
||||
|
||||
# --- SINGLE ROBOT GETTERS & SETTERS ---
|
||||
|
||||
def reset_robot_base(
|
||||
self,
|
||||
body_id: int,
|
||||
position: List[float],
|
||||
body_id_or_pos: Union[int, List[float]],
|
||||
position_or_orn: Optional[List[float]] = None,
|
||||
orientation: Optional[List[float]] = None,
|
||||
linear_velocity: Optional[List[float]] = None,
|
||||
angular_velocity: Optional[List[float]] = None
|
||||
) -> None:
|
||||
"""Resets a robot body's base position, orientation, and velocities."""
|
||||
if orientation is None:
|
||||
orientation = [0.0, 0.0, 0.0, 1.0]
|
||||
if linear_velocity is None:
|
||||
linear_velocity = [0.0, 0.0, 0.0]
|
||||
if angular_velocity is None:
|
||||
angular_velocity = [0.0, 0.0, 0.0]
|
||||
"""Resets the robot base pose and clears linear/angular velocities."""
|
||||
if isinstance(body_id_or_pos, int):
|
||||
bid = body_id_or_pos
|
||||
pos = position_or_orn if position_or_orn is not None else [0.0, 0.0, 0.20]
|
||||
orn = orientation if orientation is not None else [0.0, 0.0, 0.0, 1.0]
|
||||
else:
|
||||
bid = self.robot_id
|
||||
pos = body_id_or_pos
|
||||
orn = position_or_orn if position_or_orn is not None else [0.0, 0.0, 0.0, 1.0]
|
||||
|
||||
lin_v = linear_velocity if linear_velocity is not None else [0.0, 0.0, 0.0]
|
||||
ang_v = angular_velocity if angular_velocity is not None else [0.0, 0.0, 0.0]
|
||||
|
||||
p.resetBasePositionAndOrientation(
|
||||
body_id, position, orientation, physicsClientId=self.physics_client
|
||||
bid, pos, orn, physicsClientId=self.physics_client
|
||||
)
|
||||
p.resetBaseVelocity(
|
||||
body_id, linearVelocity=linear_velocity, angularVelocity=angular_velocity,
|
||||
bid, linearVelocity=lin_v, angularVelocity=ang_v,
|
||||
physicsClientId=self.physics_client
|
||||
)
|
||||
|
||||
def get_robot_pose(self, body_id: int) -> Tuple[List[float], List[float]]:
|
||||
def get_robot_pose(self, body_id: Optional[int] = None) -> Tuple[List[float], List[float]]:
|
||||
"""Returns base position (x, y, z) and orientation quaternion (x, y, z, w)."""
|
||||
pos, orn = p.getBasePositionAndOrientation(body_id, physicsClientId=self.physics_client)
|
||||
bid = self._resolve_body_id(body_id)
|
||||
pos, orn = p.getBasePositionAndOrientation(bid, physicsClientId=self.physics_client)
|
||||
return list(pos), list(orn)
|
||||
|
||||
def get_robot_rpy(self, body_id: int) -> Tuple[float, float, float]:
|
||||
"""Returns roll, pitch, yaw angles in radians for the given robot body."""
|
||||
_, orn = p.getBasePositionAndOrientation(body_id, physicsClientId=self.physics_client)
|
||||
def get_robot_rpy(self, body_id: Optional[int] = None) -> Tuple[float, float, float]:
|
||||
"""Returns roll, pitch, yaw angles in radians for the robot."""
|
||||
bid = self._resolve_body_id(body_id)
|
||||
_, orn = p.getBasePositionAndOrientation(bid, physicsClientId=self.physics_client)
|
||||
roll, pitch, yaw = p.getEulerFromQuaternion(orn)
|
||||
return float(roll), float(pitch), float(yaw)
|
||||
|
||||
def get_foot_link_indices(self, body_id: int) -> list[int]:
|
||||
def get_robot_pose_and_rpy(self, body_id: Optional[int] = None) -> Tuple[List[float], Tuple[float, float, float]]:
|
||||
"""Returns base position and (roll, pitch, yaw) tuple in radians."""
|
||||
bid = self._resolve_body_id(body_id)
|
||||
pos, orn = p.getBasePositionAndOrientation(bid, physicsClientId=self.physics_client)
|
||||
roll, pitch, yaw = p.getEulerFromQuaternion(orn)
|
||||
return list(pos), (float(roll), float(pitch), float(yaw))
|
||||
|
||||
def get_robot_velocity(self, body_id: Optional[int] = None) -> Tuple[List[float], List[float]]:
|
||||
"""Returns linear velocity (vx, vy, vz) and angular velocity (wx, wy, wz)."""
|
||||
bid = self._resolve_body_id(body_id)
|
||||
lin_v, ang_v = p.getBaseVelocity(bid, physicsClientId=self.physics_client)
|
||||
return list(lin_v), list(ang_v)
|
||||
|
||||
def get_robot_joint_angles(
|
||||
self,
|
||||
body_id: Optional[int] = None,
|
||||
joint_indices: Optional[List[int]] = None
|
||||
) -> np.ndarray:
|
||||
"""Returns joint angles as a 1D numpy float32 array."""
|
||||
bid = self._resolve_body_id(body_id)
|
||||
j_indices = joint_indices if joint_indices is not None else self.joint_indices
|
||||
joint_states = p.getJointStates(bid, j_indices, physicsClientId=self.physics_client)
|
||||
return np.array([state[0] for state in joint_states], dtype=np.float32)
|
||||
|
||||
def get_foot_link_indices(self, body_id: Optional[int] = None) -> List[int]:
|
||||
"""Inspects URDF joint structure to extract link IDs for leg tips and tibias."""
|
||||
bid = self._resolve_body_id(body_id)
|
||||
foot_indices = []
|
||||
num_joints = p.getNumJoints(body_id, physicsClientId=self.physics_client)
|
||||
num_joints = p.getNumJoints(bid, physicsClientId=self.physics_client)
|
||||
for j_idx in range(num_joints):
|
||||
info = p.getJointInfo(body_id, j_idx, physicsClientId=self.physics_client)
|
||||
info = p.getJointInfo(bid, j_idx, physicsClientId=self.physics_client)
|
||||
link_name = info[12].decode("utf-8")
|
||||
if "tip" in link_name or "tibia" in link_name:
|
||||
foot_indices.append(j_idx)
|
||||
return foot_indices
|
||||
|
||||
def get_robot_pose_and_rpy(self, body_id: int) -> Tuple[List[float], Tuple[float, float, float]]:
|
||||
"""Returns base position and (roll, pitch, yaw) tuple in radians."""
|
||||
pos, orn = p.getBasePositionAndOrientation(body_id, physicsClientId=self.physics_client)
|
||||
roll, pitch, yaw = p.getEulerFromQuaternion(orn)
|
||||
return list(pos), (float(roll), float(pitch), float(yaw))
|
||||
def set_robot_color(
|
||||
self,
|
||||
body_id_or_rgba: Union[int, List[float]],
|
||||
rgba: Optional[List[float]] = None
|
||||
) -> None:
|
||||
"""Changes visual color RGBA of base link and all joints."""
|
||||
if isinstance(body_id_or_rgba, int):
|
||||
bid = body_id_or_rgba
|
||||
color = rgba
|
||||
else:
|
||||
bid = self.robot_id
|
||||
color = body_id_or_rgba
|
||||
|
||||
def get_robot_velocity(self, body_id: int) -> Tuple[List[float], List[float]]:
|
||||
"""Returns linear velocity (vx, vy, vz) and angular velocity (wx, wy, wz)."""
|
||||
lin_v, ang_v = p.getBaseVelocity(body_id, physicsClientId=self.physics_client)
|
||||
return list(lin_v), list(ang_v)
|
||||
if color is None:
|
||||
color = [1.0, 1.0, 1.0, 1.0]
|
||||
|
||||
def get_robot_joint_angles(self, body_id: int, joint_indices: Optional[List[int]] = None) -> np.ndarray:
|
||||
"""Returns joint angles as a 1D numpy array float32 for specified or registered joint indices."""
|
||||
if joint_indices is None:
|
||||
joint_indices = self.robot_joints.get(body_id, list(range(18)))
|
||||
joint_states = p.getJointStates(body_id, joint_indices, physicsClientId=self.physics_client)
|
||||
return np.array([state[0] for state in joint_states], dtype=np.float32)
|
||||
|
||||
def set_robot_color(self, body_id: int, rgba: List[float]) -> None:
|
||||
"""Changes visual color RGBA of base link and all joints of the specified robot body."""
|
||||
num_joints = p.getNumJoints(body_id, physicsClientId=self.physics_client)
|
||||
p.changeVisualShape(body_id, -1, rgbaColor=rgba, physicsClientId=self.physics_client)
|
||||
num_joints = p.getNumJoints(bid, physicsClientId=self.physics_client)
|
||||
p.changeVisualShape(bid, -1, rgbaColor=color, physicsClientId=self.physics_client)
|
||||
for j in range(num_joints):
|
||||
p.changeVisualShape(body_id, j, rgbaColor=rgba, physicsClientId=self.physics_client)
|
||||
p.changeVisualShape(bid, j, rgbaColor=color, physicsClientId=self.physics_client)
|
||||
|
||||
def measure_robot_heights(self, robot_ids: List[int]) -> List[float]:
|
||||
"""Gets current Z height for all specified robot body IDs."""
|
||||
heights = []
|
||||
for body_id in robot_ids:
|
||||
pos, _ = p.getBasePositionAndOrientation(body_id, physicsClientId=self.physics_client)
|
||||
heights.append(pos[2])
|
||||
return heights
|
||||
def measure_robot_height(self, body_id: Optional[int] = None) -> float:
|
||||
"""Gets current Z height of the robot base."""
|
||||
bid = self._resolve_body_id(body_id)
|
||||
pos, _ = p.getBasePositionAndOrientation(bid, physicsClientId=self.physics_client)
|
||||
return pos[2]
|
||||
|
||||
def settle_and_measure_height(self, robot_ids: List[int], steps: int = 200, fallback_height: float = 0.14) -> float:
|
||||
"""Steps simulation for designated steps so robot settles, then calculates target standing height."""
|
||||
for _ in range(steps):
|
||||
def settle_and_measure_height(
|
||||
self,
|
||||
robot_ids_or_steps: Union[List[int], int] = 200,
|
||||
steps: int = 200,
|
||||
fallback_height: float = 0.122
|
||||
) -> float:
|
||||
"""Steps simulation so robot settles onto floor, then calculates target standing height."""
|
||||
# Intelligently resolve whether the first argument was passed as robot_ids list or step count
|
||||
actual_steps = robot_ids_or_steps if isinstance(robot_ids_or_steps, int) else steps
|
||||
|
||||
for _ in range(actual_steps):
|
||||
self.step()
|
||||
heights = self.measure_robot_heights(robot_ids)
|
||||
mean_height = float(np.mean(heights)) if heights else fallback_height
|
||||
return mean_height if mean_height > 0.0 else fallback_height
|
||||
height = self.measure_robot_height()
|
||||
return height if height > 0.0 else fallback_height
|
||||
|
||||
def apply_external_force(
|
||||
self,
|
||||
body_id: int,
|
||||
force: list[float] | np.ndarray,
|
||||
force: Union[List[float], np.ndarray],
|
||||
body_id: Optional[int] = None,
|
||||
link_index: int = -1,
|
||||
position: list[float] | np.ndarray = (0.0, 0.0, 0.0),
|
||||
position: Union[List[float], np.ndarray] = (0.0, 0.0, 0.0),
|
||||
frame: int = p.WORLD_FRAME,
|
||||
):
|
||||
"""
|
||||
Applies a 3D force vector (in Newtons) to a robot link.
|
||||
|
||||
:param body_id: PyBullet body ID.
|
||||
:param force: [fx, fy, fz] force vector in Newtons.
|
||||
:param link_index: Target link index (-1 refers to the base/torso).
|
||||
:param position: Offset [x, y, z] relative to link center where force is applied.
|
||||
:param frame: p.WORLD_FRAME (global axes) or p.LINK_FRAME (robot's body axes).
|
||||
"""
|
||||
) -> None:
|
||||
"""Applies external force vector to target link (defaults to base link)."""
|
||||
bid = self._resolve_body_id(body_id)
|
||||
p.applyExternalForce(
|
||||
objectUniqueId=body_id,
|
||||
objectUniqueId=bid,
|
||||
linkIndex=link_index,
|
||||
forceObj=list(force),
|
||||
posObj=list(position),
|
||||
|
||||
@@ -5,6 +5,7 @@ import time
|
||||
import math
|
||||
from enum import IntEnum
|
||||
from typing import Optional, Tuple, Dict, Any, List
|
||||
from collections import defaultdict
|
||||
|
||||
import gymnasium as gym
|
||||
from gymnasium import spaces
|
||||
@@ -37,8 +38,10 @@ class JackBotEnv(gym.Env):
|
||||
random_command: bool = True,
|
||||
max_episode_steps: int = 3000,
|
||||
urdf_path: str = cfg.urdf_path,
|
||||
robot_mode: str = "direct",
|
||||
):
|
||||
super().__init__()
|
||||
self.robot_mode = robot_mode
|
||||
self.use_gui = use_gui
|
||||
self.random_command = random_command
|
||||
self.max_episode_steps = max_episode_steps
|
||||
@@ -59,6 +62,10 @@ class JackBotEnv(gym.Env):
|
||||
self._curriculum_advanced = False
|
||||
self._first_reset = True
|
||||
|
||||
# Reward Component Tracking Initialization
|
||||
self.last_reward_components: Dict[str, float] = {}
|
||||
self.episode_reward_components_sum: Dict[str, float] = defaultdict(float)
|
||||
|
||||
# Dynamic Command Resampling Timing (60 Hz control loop)
|
||||
self.control_freq = 60
|
||||
self.min_cmd_hold_steps = int(2.0 * self.control_freq) # 120 steps (2s)
|
||||
@@ -74,16 +81,14 @@ class JackBotEnv(gym.Env):
|
||||
self.plane, pb_robots, robot_joint_indices = self.sim_manager.load_scene(
|
||||
self.urdf_path, 0.0, self._robot_base_position
|
||||
)
|
||||
self.pb_robot = pb_robots[0]
|
||||
self.joint_indices = robot_joint_indices[0]
|
||||
|
||||
# Instantiate Robot Python wrapper (start_pose is managed inside Robot.py)
|
||||
# Instantiate Robot Python wrapper
|
||||
self.robot = Robot(
|
||||
backend_type=PyBulletBackend(self.sim_manager, body_id=self.pb_robot),
|
||||
backend_type=PyBulletBackend(self.sim_manager),
|
||||
urdf_path=self.urdf_path
|
||||
)
|
||||
|
||||
# Action (18 joint deltas) & Observation (18 angles + 4 command dims)
|
||||
action_dim = 18
|
||||
obs_dim = 18 + 4
|
||||
|
||||
@@ -109,7 +114,7 @@ class JackBotEnv(gym.Env):
|
||||
CurriculumPhase.FORWARD: {
|
||||
"survival_steps": 300,
|
||||
"min_avg_height_ratio": 0.88,
|
||||
"max_avg_roll_pitch": 0.18, # ~10 degrees average
|
||||
"max_avg_roll_pitch": 0.18,
|
||||
},
|
||||
CurriculumPhase.TURN_AND_DIRECTION: {
|
||||
"survival_steps": 500,
|
||||
@@ -136,7 +141,7 @@ class JackBotEnv(gym.Env):
|
||||
self.last_time = time.time()
|
||||
|
||||
def _robot_base_position(self, robot_id: int, spacing: float = 0.0) -> list[float]:
|
||||
return [0.0, 0.0, 0.13]
|
||||
return [0.0, 0.0, 0.2]
|
||||
|
||||
def _find_foot_link_indices(self) -> list:
|
||||
return self.sim_manager.get_foot_link_indices(self.pb_robot)
|
||||
@@ -187,12 +192,19 @@ class JackBotEnv(gym.Env):
|
||||
self.episode_roll_sum = 0.0
|
||||
self.episode_pitch_sum = 0.0
|
||||
|
||||
# Reset Component Tracking Dictionary
|
||||
self.last_reward_components = {}
|
||||
self.episode_reward_components_sum = defaultdict(float)
|
||||
|
||||
spawn_pos = self._robot_base_position(0)
|
||||
spawn_orn = [0.0, 0.0, 0.0, 1.0]
|
||||
|
||||
self.sim_manager.reset_robot_base(self.pb_robot, spawn_pos, spawn_orn)
|
||||
self.robot.reset_to_init()
|
||||
|
||||
init_angles = self.robot.current_rad.data.flatten()
|
||||
self.sim_manager.hard_reset_joint_angles(init_angles, self.pb_robot)
|
||||
|
||||
if self.use_gui:
|
||||
self.sim_manager.set_robot_color(self.pb_robot, [1.0, 1.0, 1.0, 1.0])
|
||||
|
||||
@@ -211,8 +223,9 @@ class JackBotEnv(gym.Env):
|
||||
pos, _ = self.sim_manager.get_robot_pose(self.pb_robot)
|
||||
self.start_position = [float(pos[0]), float(pos[1]), float(pos[2])]
|
||||
|
||||
# Fixed single-robot settlement call
|
||||
self.target_height = self.sim_manager.settle_and_measure_height(
|
||||
[self.pb_robot], steps=200, fallback_height=0.122
|
||||
steps=200, fallback_height=0.122
|
||||
)
|
||||
|
||||
self.default_joint_angles = np.array(
|
||||
@@ -226,8 +239,13 @@ class JackBotEnv(gym.Env):
|
||||
|
||||
return self._get_obs(), {}
|
||||
|
||||
def get_reward_component_averages(self) -> Dict[str, float]:
|
||||
"""Calculates step-averaged scores for each sub-reward component."""
|
||||
steps = max(1, self.step_count)
|
||||
return {k: float(v / steps) for k, v in self.episode_reward_components_sum.items()}
|
||||
|
||||
def get_current_robot_metrics(self) -> list:
|
||||
"""Returns metric summary for the callback."""
|
||||
"""Returns metric summary for callbacks."""
|
||||
if self.is_failed:
|
||||
return []
|
||||
|
||||
@@ -263,17 +281,16 @@ class JackBotEnv(gym.Env):
|
||||
random_interval = np.random.randint(self.min_cmd_hold_steps, self.max_cmd_hold_steps + 1)
|
||||
self.next_cmd_resample_step = self.step_count + random_interval
|
||||
|
||||
self.robot.apply_rl_action(action)
|
||||
self.robot.step_with_command(command=self.command, action=action, mode=self.robot_mode)
|
||||
|
||||
if self.step_count % 60 == 0:
|
||||
if self.robot_mode != "kinematics_only" and self.step_count % 60 == 0:
|
||||
random_force = np.random.uniform(-2.0, 2.0, size=2)
|
||||
self.sim_manager.apply_external_force(
|
||||
body_id=self.pb_robot,
|
||||
force=[random_force[0], random_force[1], 0.0]
|
||||
)
|
||||
|
||||
render_freq = 10 # Only draw 1 in every 10 frames
|
||||
|
||||
render_freq = 10
|
||||
if self.use_gui and self.step_count % render_freq != 0:
|
||||
self.sim_manager.set_rendering(False)
|
||||
|
||||
@@ -295,9 +312,15 @@ class JackBotEnv(gym.Env):
|
||||
terminated = self.is_failed
|
||||
truncated = self.step_count >= self.max_episode_steps
|
||||
|
||||
info = {
|
||||
"reward_components": self.last_reward_components.copy()
|
||||
}
|
||||
|
||||
if self.step_count % 120 == 0 and self.use_gui:
|
||||
self._update_hud()
|
||||
return obs, reward, terminated, truncated, {}
|
||||
|
||||
# Gymnasium standard 5-tuple return
|
||||
return obs, reward, terminated, truncated, info
|
||||
|
||||
def _update_distance_metrics(self):
|
||||
pos, _ = self.sim_manager.get_robot_pose(self.pb_robot)
|
||||
@@ -312,22 +335,17 @@ class JackBotEnv(gym.Env):
|
||||
return False
|
||||
|
||||
req = self.curriculum_stage_requirements[next_phase]
|
||||
|
||||
# 1. Survival Check
|
||||
survival_ok = self.max_survival_steps >= req["survival_steps"]
|
||||
|
||||
# 2. Smooth Average Stability Checks (Prevents 1-frame spikes from failing curriculum)
|
||||
avg_roll = self.episode_roll_sum / max(1, self.step_count)
|
||||
avg_pitch = self.episode_pitch_sum / max(1, self.step_count)
|
||||
max_allowed_angle = req.get("max_avg_roll_pitch", 0.20)
|
||||
stability_ok = (avg_roll <= max_allowed_angle) and (avg_pitch <= max_allowed_angle)
|
||||
|
||||
# 3. Average Height Check
|
||||
avg_height = self.episode_height_sum / max(1, self.step_count)
|
||||
required_min_avg_height = self.target_height * req.get("min_avg_height_ratio", 0.85)
|
||||
height_ok = avg_height >= required_min_avg_height
|
||||
|
||||
# 4. Distance and Drift Checks
|
||||
pos, _ = self.sim_manager.get_robot_pose(self.pb_robot)
|
||||
start_x, start_y, _ = self.start_position
|
||||
dx = pos[0] - start_x
|
||||
@@ -373,10 +391,8 @@ class JackBotEnv(gym.Env):
|
||||
|
||||
def _compute_reward(self, action: np.ndarray, previous_action: np.ndarray) -> float:
|
||||
"""
|
||||
Calculates task rewards using normalized Exponential Kernels.
|
||||
Includes a deadband filter for jittering and zero-reward gating when stationary.
|
||||
Calculates task rewards using normalized Exponential Kernels and tracks component terms.
|
||||
"""
|
||||
# 1. Fetch Robot State
|
||||
pos, (roll, pitch, yaw) = self.sim_manager.get_robot_pose_and_rpy(self.pb_robot)
|
||||
linear_vel, angular_vel = self.sim_manager.get_robot_velocity(self.pb_robot)
|
||||
current_joints = np.array(
|
||||
@@ -387,9 +403,8 @@ class JackBotEnv(gym.Env):
|
||||
cmd_vx, cmd_vy, _, cmd_yaw = self.command
|
||||
cmd_norm = math.hypot(cmd_vx, cmd_vy)
|
||||
|
||||
# 2. Velocity Deadband Filtering (Ignores jittering & micro-movements)
|
||||
VEL_DEADBAND = 0.04 # 4 cm/s threshold
|
||||
YAW_DEADBAND = 0.05 # 0.05 rad/s threshold
|
||||
VEL_DEADBAND = 0.04
|
||||
YAW_DEADBAND = 0.05
|
||||
|
||||
raw_speed = math.hypot(linear_vel[0], linear_vel[1])
|
||||
if raw_speed < VEL_DEADBAND:
|
||||
@@ -405,7 +420,6 @@ class JackBotEnv(gym.Env):
|
||||
else:
|
||||
filtered_yaw_rate = angular_vel[2]
|
||||
|
||||
# 3. Posture & Stability Sub-Rewards
|
||||
height_error = pos[2] - self.target_height
|
||||
r_height = math.exp(-150.0 * (height_error ** 2))
|
||||
|
||||
@@ -418,13 +432,19 @@ class JackBotEnv(gym.Env):
|
||||
action_delta = np.mean(np.square(action - previous_action))
|
||||
r_smoothness = math.exp(-0.1 * action_delta)
|
||||
|
||||
# 4. Mode Logic
|
||||
r_lin_vel = 0.0
|
||||
r_ang_vel = 0.0
|
||||
stillness_penalty = 0.0
|
||||
gated_zero = 0.0
|
||||
|
||||
if cmd_norm < 0.05 and abs(cmd_yaw) < 0.05:
|
||||
# STANDING MODE: Reward clean posture, height, and stability
|
||||
# STANDING MODE
|
||||
w_height = 0.35
|
||||
w_stability = 0.35
|
||||
w_pose = 0.20
|
||||
w_smoothness = 0.10
|
||||
w_lin_vel = 0.0
|
||||
w_ang_vel = 0.0
|
||||
|
||||
total_reward = (
|
||||
(w_height * r_height)
|
||||
@@ -436,42 +456,60 @@ class JackBotEnv(gym.Env):
|
||||
# WALKING / TURNING MODE
|
||||
is_moving = (filtered_speed > 0.0) or (abs(filtered_yaw_rate) > 0.0)
|
||||
|
||||
# HARD GATE: If commanded to move but standing still/jittering, reward is strictly 0.0
|
||||
if not is_moving:
|
||||
return 0.0
|
||||
|
||||
target_vx = cmd_vx * self.max_robot_speed
|
||||
target_vy = cmd_vy * self.max_robot_speed
|
||||
target_speed = math.hypot(target_vx, target_vy)
|
||||
|
||||
lin_vel_error = (filtered_vx - target_vx)**2 + (filtered_vy - target_vy)**2
|
||||
r_lin_vel = math.exp(-25.0 * lin_vel_error)
|
||||
if not is_moving:
|
||||
gated_zero = 1.0
|
||||
total_reward = 0.0
|
||||
w_lin_vel, w_ang_vel, w_height, w_stability, w_pose, w_smoothness = 0, 0, 0, 0, 0, 0
|
||||
else:
|
||||
lin_vel_error = (filtered_vx - target_vx)**2 + (filtered_vy - target_vy)**2
|
||||
r_lin_vel = math.exp(-25.0 * lin_vel_error)
|
||||
|
||||
ang_vel_error = (filtered_yaw_rate - cmd_yaw)**2
|
||||
r_ang_vel = math.exp(-15.0 * ang_vel_error)
|
||||
ang_vel_error = (filtered_yaw_rate - cmd_yaw)**2
|
||||
r_ang_vel = math.exp(-15.0 * ang_vel_error)
|
||||
|
||||
# Stillness Check: Commanded to move, but staying virtually still
|
||||
stillness_penalty = 0.0
|
||||
if math.hypot(target_vx, target_vy) > 0.08 and math.hypot(linear_vel[0], linear_vel[1]) < 0.03:
|
||||
r_lin_vel = 0.0 # Strip velocity credit completely
|
||||
stillness_penalty = -0.25
|
||||
if target_speed > 0.08 and raw_speed < 0.03:
|
||||
r_lin_vel = 0.0
|
||||
stillness_penalty = -0.25
|
||||
|
||||
w_lin_vel = 0.55
|
||||
w_ang_vel = 0.15
|
||||
w_height = 0.10
|
||||
w_stability = 0.12
|
||||
w_smoothness = 0.08
|
||||
w_lin_vel = 0.55
|
||||
w_ang_vel = 0.15
|
||||
w_height = 0.10
|
||||
w_stability = 0.12
|
||||
w_pose = 0.0
|
||||
w_smoothness = 0.08
|
||||
|
||||
total_reward = (
|
||||
(w_lin_vel * r_lin_vel)
|
||||
+ (w_ang_vel * r_ang_vel)
|
||||
+ (w_height * r_height)
|
||||
+ (w_stability * r_stability)
|
||||
+ (w_smoothness * r_smoothness)
|
||||
+ stillness_penalty
|
||||
)
|
||||
total_reward = (
|
||||
(w_lin_vel * r_lin_vel)
|
||||
+ (w_ang_vel * r_ang_vel)
|
||||
+ (w_height * r_height)
|
||||
+ (w_stability * r_stability)
|
||||
+ (w_smoothness * r_smoothness)
|
||||
+ stillness_penalty
|
||||
)
|
||||
|
||||
# Scaled reward for policy stability
|
||||
return float(total_reward / 10.0)
|
||||
final_reward = float(total_reward / 10.0)
|
||||
|
||||
comp = {
|
||||
"lin_vel": float((w_lin_vel * r_lin_vel) / 10.0),
|
||||
"ang_vel": float((w_ang_vel * r_ang_vel) / 10.0),
|
||||
"height": float((w_height * r_height) / 10.0),
|
||||
"stability": float((w_stability * r_stability) / 10.0),
|
||||
"pose": float((w_pose * r_pose) / 10.0),
|
||||
"smoothness": float((w_smoothness * r_smoothness) / 10.0),
|
||||
"stillness_penalty": float(stillness_penalty / 10.0),
|
||||
"gated_zero": gated_zero,
|
||||
"total_step_reward": final_reward,
|
||||
}
|
||||
|
||||
self.last_reward_components = comp
|
||||
for key, val in comp.items():
|
||||
self.episode_reward_components_sum[key] += val
|
||||
|
||||
return final_reward
|
||||
|
||||
def _update_hud(self):
|
||||
if not self.use_gui:
|
||||
@@ -579,4 +617,34 @@ class CurriculumCallback(BaseCallback):
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
return True
|
||||
|
||||
|
||||
class RewardLoggerCallback(BaseCallback):
|
||||
"""
|
||||
Logs step-averaged individual reward components to TensorBoard during PPO training.
|
||||
"""
|
||||
|
||||
def __init__(self, verbose=0):
|
||||
super().__init__(verbose)
|
||||
|
||||
def _on_step(self) -> bool:
|
||||
return True
|
||||
|
||||
def _on_rollout_end(self) -> bool:
|
||||
try:
|
||||
vec_env = self.training_env
|
||||
all_comp_averages = vec_env.env_method("get_reward_component_averages")
|
||||
|
||||
if not all_comp_averages:
|
||||
return True
|
||||
|
||||
keys = all_comp_averages[0].keys()
|
||||
for key in keys:
|
||||
avg_val = np.mean([env_comp.get(key, 0.0) for env_comp in all_comp_averages])
|
||||
self.logger.record(f"reward_components/{key}", float(avg_val))
|
||||
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
return True
|
||||
@@ -0,0 +1,76 @@
|
||||
"""
|
||||
ml/run_eval_training.py - Benchmark reward system across all curriculum phases.
|
||||
"""
|
||||
import time
|
||||
import numpy as np
|
||||
from ml.env import JackBotEnv, CurriculumPhase
|
||||
|
||||
def evaluate_kinematics(episode_length: int = 1000):
|
||||
# Use kinematics_only mode so inverse kinematics generates gait motion from commands
|
||||
env = JackBotEnv(
|
||||
use_gui=True,
|
||||
random_command=False,
|
||||
max_episode_steps=episode_length,
|
||||
robot_mode="kinematics_only"
|
||||
)
|
||||
|
||||
print("\n" + "=" * 70)
|
||||
print(" RUNNING MULTI-PHASE REWARD BENCHMARK (KINEMATICS MODE)")
|
||||
print("=" * 70 + "\n")
|
||||
|
||||
# Define test suite covering every curriculum stage
|
||||
phase_configs = [
|
||||
(CurriculumPhase.STAND_ONLY, "STAND ONLY", np.array([0.0, 0.0, 0.0, 0.0], dtype=np.float32)),
|
||||
(CurriculumPhase.FORWARD, "FORWARD GAIT", np.array([0.1, 0.0, 0.0, 0.0], dtype=np.float32)),
|
||||
(CurriculumPhase.TURN_AND_DIRECTION, "FORWARD + YAW TURN", np.array([0.3, 0.0, 0.0, 0.4], dtype=np.float32)),
|
||||
(CurriculumPhase.OMNI_DIRECTION, "STRIDE LATERAL", np.array([0.2, 0.3, 0.0, 0.0], dtype=np.float32)),
|
||||
(CurriculumPhase.FULL_COMMAND, "FULL OMNI COMBINATION",np.array([0.3, 0.2, 0.0, 0.3], dtype=np.float32)),
|
||||
]
|
||||
|
||||
for phase_enum, label, cmd in phase_configs:
|
||||
obs, _ = env.reset()
|
||||
|
||||
# Force specific curriculum phase & target command
|
||||
env.curriculum_phase = phase_enum
|
||||
env.command = cmd.copy()
|
||||
|
||||
done = False
|
||||
total_reward = 0.0
|
||||
step_count = 0
|
||||
|
||||
while not done:
|
||||
# Action array is unused in kinematics_only mode
|
||||
dummy_action = np.zeros(18, dtype=np.float32)
|
||||
|
||||
obs, reward, terminated, truncated, _ = env.step(dummy_action)
|
||||
total_reward += reward
|
||||
step_count += 1
|
||||
done = terminated or truncated
|
||||
|
||||
time.sleep(1.0 / 60.0)
|
||||
|
||||
# Retrieve detailed component averages
|
||||
comp_averages = env.get_reward_component_averages()
|
||||
|
||||
print(f"\n--- Episode Stage: [{phase_enum.name}] ({label}) ---")
|
||||
print(f"Command Applied: vx={cmd[0]:.2f}, vy={cmd[1]:.2f}, vz={cmd[2]:.2f}, yaw={cmd[3]:.2f}")
|
||||
print(f"Total Episode Reward: {total_reward:.4f}")
|
||||
print("Component Step Averages:")
|
||||
for name, value in comp_averages.items():
|
||||
print(f" • {name:<20}: {value:+.5f}")
|
||||
|
||||
metrics = env.get_current_robot_metrics()
|
||||
dist = metrics[0]["distance_from_start"] if metrics else 0.0
|
||||
speed = metrics[0]["speed"] if metrics else 0.0
|
||||
avg_reward = total_reward / max(1, step_count)
|
||||
|
||||
print(f" ├─ Average Reward / Step: {avg_reward:.4f}")
|
||||
print(f" ├─ Distance Travelled: {dist:.2f} m")
|
||||
print(f" ├─ Actual Avg Speed: {speed:.2f} m/s")
|
||||
print(f" └─ Steps Survived: {step_count} / {episode_length}")
|
||||
print("-" * 70)
|
||||
|
||||
env.close()
|
||||
|
||||
if __name__ == "__main__":
|
||||
evaluate_kinematics()
|
||||
+1
-1
@@ -77,7 +77,7 @@ def main():
|
||||
model = PPO(
|
||||
policy="MlpPolicy",
|
||||
env=vec_env,
|
||||
learning_rate=3e-4,
|
||||
learning_rate=1e-4,
|
||||
n_steps=256,
|
||||
batch_size=256,
|
||||
n_epochs=10,
|
||||
|
||||
Reference in New Issue
Block a user