JackM323/JackBot
1
0
Fork 0
Code Issues 3 Pull Requests Actions Packages Projects 2 Releases Wiki Activity

simulation physics fixed

Browse Source 
fuck them physics
This commit is contained in:
JackM323
2026-08-06 15:18:06 +02:00
parent 523a4aea89
commit 346ec9e949
8 changed files with 457 additions and 900 deletions
Download Patch File Download Diff File Expand all files Collapse all files
ml/env.py
+116 -372
View File
@@ -10,14 +10,12 @@ from collections import defaultdict
import gymnasium as gym
from gymnasium import spaces
import numpy as np
from stable_baselines3.common.callbacks import BaseCallback
from config import cfg
from simulation import Simulation
from Robot import Robot, PyBulletBackend
from ml.SimManager import SimManager
from ml.MetricsOverlay import MetricsHUD
# Color Palette RGBA for Terminated/Failed Robots
COLOR_FAILED = [0.3, 0.3, 0.3, 0.6]
@@ -30,7 +28,7 @@ class CurriculumPhase(IntEnum):
class JackBotEnv(gym.Env):
"""Gymnasium environment wrapping a single JackBot hexapod."""
"""Gymnasium environment wrapping JackBot hexapod simulation."""
def __init__(
self,
@@ -38,7 +36,7 @@ class JackBotEnv(gym.Env):
random_command: bool = True,
max_episode_steps: int = 3000,
urdf_path: str = cfg.urdf_path,
robot_mode: str = "direct",
robot_mode: str = "direct", # "direct", "residual", or "kinematics"
):
super().__init__()
self.robot_mode = robot_mode
@@ -52,7 +50,6 @@ class JackBotEnv(gym.Env):
self.total_steps = 0
self.cumulative_reward = 0.0
self.robot_reward = 0.0
self.consecutive_still_steps = 0
self.is_failed = False
self.episode_count = 0
@@ -62,94 +59,55 @@ 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)
self.max_cmd_hold_steps = int(6.0 * self.control_freq) # 360 steps (6s)
self.min_cmd_hold_steps = int(2.0 * self.control_freq)
self.max_cmd_hold_steps = int(6.0 * self.control_freq)
self.next_cmd_resample_step = 0
self.initial_stand_steps = 120 # Mandatory 2s standing window at episode reset
self.initial_stand_steps = 120
# Initialize Simulation Manager
self.sim_manager = SimManager(use_gui=self.use_gui)
self.sim_manager.connect()
# --- INITIALIZE PYBULLET SIMULATION ENGINE ---
self.sim = Simulation(urdf_path=self.urdf_path, use_gui=self.use_gui)
self.plane_id, self.pb_robot, self.revolute_joints = self.sim.load_scene()
# Connect physics world & load single robot
self.plane, pb_robots, robot_joint_indices = self.sim_manager.load_scene(
self.urdf_path, 0.0, self._robot_base_position
)
self.joint_indices = robot_joint_indices[0]
# Instantiate Robot Python wrapper
self.robot = Robot(
backend_type=PyBulletBackend(self.sim_manager),
urdf_path=self.urdf_path
)
# --- INITIALIZE ROBOT WITH BACKEND ---
self.backend = PyBulletBackend(self.sim)
self.robot = Robot(backend_type=self.backend, urdf_path=self.urdf_path, mode=self.robot_mode)
action_dim = 18
obs_dim = 18 + 4
obs_dim = 18 + 3 # 18 Joint Angles + 3 Command Inputs [vx, vy, omega]
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)
self.command = np.zeros(4, dtype=np.float32)
self.command = np.zeros(3, dtype=np.float32) # [vx, vy, omega]
self.last_action = np.zeros(action_dim, dtype=np.float32)
self.target_height = 0.122
self.collapse_height_fraction = 0.55
self.tilt_failure_rad = 0.9
self.foot_link_indices = self._find_foot_link_indices()
self.start_position = [0.0, 0.0, 0.0]
self.max_distance_from_start = 0.0
self.max_survival_steps = 0
self.default_joint_angles = np.zeros(18, dtype=np.float32)
# Curriculum Initialization via Enum
self.curriculum_phase = CurriculumPhase.STAND_ONLY
self.curriculum_stage_requirements = {
CurriculumPhase.FORWARD: {
"survival_steps": 300,
"min_avg_height_ratio": 0.88,
"max_avg_roll_pitch": 0.18,
},
CurriculumPhase.TURN_AND_DIRECTION: {
"survival_steps": 500,
"min_forward_distance": 2.5,
"max_lateral_drift": 0.8,
"min_avg_height_ratio": 0.85,
"stability_roll_pitch": 0.25,
},
CurriculumPhase.OMNI_DIRECTION: {
"survival_steps": 600,
"min_distance": 5.0,
"min_avg_height_ratio": 0.85,
"stability_roll_pitch": 0.25,
},
CurriculumPhase.FULL_COMMAND: {
"survival_steps": 750,
"min_distance": 8.0,
"min_avg_height_ratio": 0.85,
"stability_roll_pitch": 0.20,
},
CurriculumPhase.FORWARD: {"survival_steps": 300, "min_avg_height_ratio": 0.88, "max_avg_roll_pitch": 0.18},
CurriculumPhase.TURN_AND_DIRECTION: {"survival_steps": 500, "min_forward_distance": 2.5, "max_lateral_drift": 0.8, "min_avg_height_ratio": 0.85, "stability_roll_pitch": 0.25},
CurriculumPhase.OMNI_DIRECTION: {"survival_steps": 600, "min_distance": 5.0, "min_avg_height_ratio": 0.85, "stability_roll_pitch": 0.25},
CurriculumPhase.FULL_COMMAND: {"survival_steps": 750, "min_distance": 8.0, "min_avg_height_ratio": 0.85, "stability_roll_pitch": 0.20},
}
self.hud = MetricsHUD(physics_client_id=self.sim_manager.physics_client)
self.hud = MetricsHUD(physics_client_id=self.sim.physics_client)
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.2]
def _find_foot_link_indices(self) -> list:
return self.sim_manager.get_foot_link_indices(self.pb_robot)
def sample_command(self) -> np.ndarray:
"""Samples a command vector [vx, vy, omega] based on active curriculum phase."""
phase = self.curriculum_phase
stand_probabilities = {
stand_probs = {
CurriculumPhase.STAND_ONLY: 1.0,
CurriculumPhase.FORWARD: 0.25,
CurriculumPhase.TURN_AND_DIRECTION: 0.20,
@@ -157,28 +115,19 @@ class JackBotEnv(gym.Env):
CurriculumPhase.FULL_COMMAND: 0.15,
}
if np.random.random() < stand_probabilities.get(phase, 0.15):
return np.zeros(4, dtype=np.float32)
if np.random.random() < stand_probs.get(phase, 0.15):
return np.zeros(3, dtype=np.float32)
if phase == CurriculumPhase.FORWARD:
vx = np.random.uniform(0.15, 0.50)
vy, vz, omega = 0.0, 0.0, 0.0
vx, vy, omega = np.random.uniform(0.15, 0.50), 0.0, 0.0
elif phase == CurriculumPhase.TURN_AND_DIRECTION:
vx = np.random.uniform(-0.8, 0.8)
vy, vz = 0.0, 0.0
omega = np.random.uniform(-0.8, 0.8)
vx, vy, omega = np.random.uniform(-0.8, 0.8), 0.0, np.random.uniform(-0.8, 0.8)
elif phase == CurriculumPhase.OMNI_DIRECTION:
vx = np.random.uniform(-0.8, 0.8)
vy = np.random.uniform(-0.5, 0.5)
vz = 0.0
omega = np.random.uniform(-0.8, 0.8)
vx, vy, omega = np.random.uniform(-0.8, 0.8), np.random.uniform(-0.5, 0.5), np.random.uniform(-0.8, 0.8)
else:
vx = np.random.uniform(-1.0, 1.0)
vy = np.random.uniform(-1.0, 1.0)
vz = 0.0
omega = np.random.uniform(-1.0, 1.0)
vx, vy, omega = np.random.uniform(-1.0, 1.0), np.random.uniform(-1.0, 1.0), np.random.uniform(-1.0, 1.0)
return np.array([vx, vy, vz, omega], dtype=np.float32)
return np.array([vx, vy, omega], dtype=np.float32)
def reset(self, seed: Optional[int] = None, options: Optional[Dict[str, Any]] = None):
super().reset(seed=seed)
@@ -191,24 +140,21 @@ class JackBotEnv(gym.Env):
self.episode_height_sum = 0.0
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_pos = [0.0, 0.0, 0.20]
spawn_orn = [0.0, 0.0, 0.0, 1.0]
self.sim_manager.reset_robot_base(self.pb_robot, spawn_pos, spawn_orn)
# 1. Reset base pose and velocities
self.sim.reset_robot_base(spawn_pos, spawn_orn)
# 2. Reset internal kinematics & hard reset joints in PyBullet
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])
self.sim.set_robot_color([1.0, 1.0, 1.0, 1.0])
self.consecutive_still_steps = 0
self.last_action = np.zeros(self.action_space.shape[0], dtype=np.float32)
self.max_distance_from_start = 0.0
self.max_survival_steps = 0
@@ -217,21 +163,15 @@ class JackBotEnv(gym.Env):
self.curriculum_phase = CurriculumPhase.STAND_ONLY
self._first_reset = False
self.command = np.zeros(4, dtype=np.float32)
self.command = np.zeros(3, dtype=np.float32)
self.next_cmd_resample_step = self.initial_stand_steps
pos, _ = self.sim_manager.get_robot_pose(self.pb_robot)
self.start_position = [float(pos[0]), float(pos[1]), float(pos[2])]
pos, _ = self.sim.get_robot_pose()
self.start_position = list(pos)
# Fixed single-robot settlement call
self.target_height = self.sim_manager.settle_and_measure_height(
steps=200, fallback_height=0.122
)
self.default_joint_angles = np.array(
self.sim_manager.get_robot_joint_angles(self.pb_robot, self.joint_indices),
dtype=np.float32
)
# Drop settlement
self.target_height = self.sim.settle_and_measure_height(steps=200, fallback_height=0.122)
self.default_joint_angles = self.sim.get_robot_joint_angles()
if self.use_gui:
self.hud.reset()
@@ -239,36 +179,8 @@ 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 callbacks."""
if self.is_failed:
return []
pos, _ = self.sim_manager.get_robot_pose(self.pb_robot)
linear_vel, angular_vel = self.sim_manager.get_robot_velocity(self.pb_robot)
start_x, start_y, _ = self.start_position
dist = float(np.linalg.norm(np.array([pos[0] - start_x, pos[1] - start_y], dtype=np.float32)))
speed = float(np.linalg.norm(np.array([linear_vel[0], linear_vel[1]], dtype=np.float32)))
yaw_rate = float(abs(angular_vel[2]))
return [{
"reward": float(self.robot_reward),
"distance_from_start": dist,
"speed": speed,
"yaw_rate": yaw_rate,
"alive": True,
"survival_steps": int(self.step_count),
"phase_name": self.curriculum_phase.name,
}]
def _get_obs(self) -> np.ndarray:
joint_angles = self.sim_manager.get_robot_joint_angles(self.pb_robot, self.joint_indices)
return np.concatenate([joint_angles, self.command]).astype(np.float32)
return self.robot.get_observation(command=self.command)
def step(self, action: np.ndarray) -> Tuple[np.ndarray, float, bool, bool, Dict[str, Any]]:
self.step_count += 1
@@ -276,28 +188,25 @@ class JackBotEnv(gym.Env):
previous_action = self.last_action.copy()
self.last_action = action.copy()
# Command resampling
if self.random_command and (self.step_count >= self.next_cmd_resample_step or self._curriculum_advanced):
self.command = self.sample_command()
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.step_with_command(command=self.command, action=action, mode=self.robot_mode)
# Extract [vx, vy, omega]
cmd_vx, cmd_vy, cmd_omega = self.command
if self.robot_mode != "kinematics_only" and self.step_count % 60 == 0:
# Mirror main.py input resolution logic: update robot_state and vector_dirmov directly
self.robot.robot_state = "walking" if (abs(cmd_vx) > 0.01 or abs(cmd_vy) > 0.01 or abs(cmd_omega) > 0.01) else "idle"
self.robot.vector_dirmov = [float(cmd_vx), float(cmd_vy), float(cmd_omega)]
# Delegate execution tick to Robot instance
self.robot.tick(action=action)
if self.robot_mode != "kinematics" 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
if self.use_gui and self.step_count % render_freq != 0:
self.sim_manager.set_rendering(False)
self.sim_manager.step()
if self.use_gui and self.step_count % render_freq == 0:
self.sim_manager.set_rendering(True)
self.sim.apply_external_force(force=[random_force[0], random_force[1], 0.0])
self._update_robot_failure()
self._update_distance_metrics()
@@ -305,25 +214,21 @@ class JackBotEnv(gym.Env):
obs = self._get_obs()
reward = self._compute_reward(action, previous_action)
self.cumulative_reward += reward
self.robot_reward += reward
terminated = self.is_failed
truncated = self.step_count >= self.max_episode_steps
info = {
"reward_components": self.last_reward_components.copy()
}
info = {"reward_components": self.last_reward_components.copy()}
if self.step_count % 120 == 0 and self.use_gui:
self._update_hud()
# 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)
pos, _ = self.sim.get_robot_pose()
self.episode_height_sum += float(pos[2])
start_x, start_y, _ = self.start_position
dist = float(np.linalg.norm(np.array([pos[0] - start_x, pos[1] - start_y], dtype=np.float32)))
@@ -336,7 +241,6 @@ class JackBotEnv(gym.Env):
req = self.curriculum_stage_requirements[next_phase]
survival_ok = self.max_survival_steps >= req["survival_steps"]
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)
@@ -346,10 +250,8 @@ class JackBotEnv(gym.Env):
required_min_avg_height = self.target_height * req.get("min_avg_height_ratio", 0.85)
height_ok = avg_height >= required_min_avg_height
pos, _ = self.sim_manager.get_robot_pose(self.pb_robot)
start_x, start_y, _ = self.start_position
dx = pos[0] - start_x
dy = pos[1] - start_y
pos, _ = self.sim.get_robot_pose()
dx, dy = pos[0] - self.start_position[0], pos[1] - self.start_position[1]
dist_2d = math.hypot(dx, dy)
distance_ok = True
@@ -358,12 +260,9 @@ class JackBotEnv(gym.Env):
elif "min_distance" in req:
distance_ok = dist_2d >= req["min_distance"]
drift_ok = True
if "max_lateral_drift" in req:
drift_ok = abs(dy) <= req["max_lateral_drift"]
drift_ok = abs(dy) <= req["max_lateral_drift"] if "max_lateral_drift" in req else True
return survival_ok and height_ok and stability_ok and distance_ok and drift_ok
def _update_curriculum(self):
self._curriculum_advanced = False
forced_forward = (self.curriculum_phase < CurriculumPhase.FORWARD) and (self.step_count >= 2500)
@@ -371,174 +270,117 @@ class JackBotEnv(gym.Env):
if self.curriculum_phase < CurriculumPhase.FORWARD and (self._phase_progress_ready(CurriculumPhase.FORWARD) or forced_forward):
self.curriculum_phase = CurriculumPhase.FORWARD
self._curriculum_advanced = True
reason = "FORCED (2500 steps)" if forced_forward else "MET"
print(f"[Curriculum] Phase {self.curriculum_phase.name} unlocked [{reason}] at step {self.step_count}")
elif self.curriculum_phase < CurriculumPhase.TURN_AND_DIRECTION and self._phase_progress_ready(CurriculumPhase.TURN_AND_DIRECTION):
self.curriculum_phase = CurriculumPhase.TURN_AND_DIRECTION
self._curriculum_advanced = True
print(f"[Curriculum] Phase {self.curriculum_phase.name} unlocked at total step {self.total_steps}")
elif self.curriculum_phase < CurriculumPhase.OMNI_DIRECTION and self._phase_progress_ready(CurriculumPhase.OMNI_DIRECTION):
self.curriculum_phase = CurriculumPhase.OMNI_DIRECTION
self._curriculum_advanced = True
print(f"[Curriculum] Phase {self.curriculum_phase.name} unlocked at total step {self.total_steps}")
elif self.curriculum_phase < CurriculumPhase.FULL_COMMAND and self._phase_progress_ready(CurriculumPhase.FULL_COMMAND):
self.curriculum_phase = CurriculumPhase.FULL_COMMAND
self._curriculum_advanced = True
print(f"[Curriculum] Phase {self.curriculum_phase.name} unlocked at total step {self.total_steps}")
def _compute_reward(self, action: np.ndarray, previous_action: np.ndarray) -> float:
"""
Calculates task rewards using normalized Exponential Kernels and tracks component terms.
"""
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(
self.sim_manager.get_robot_joint_angles(self.pb_robot, self.joint_indices),
dtype=np.float32
)
pos, (roll, pitch, yaw) = self.sim.get_robot_pose_and_rpy()
linear_vel, angular_vel = self.sim.get_robot_velocity()
current_joints = self.sim.get_robot_joint_angles()
cmd_vx, cmd_vy, _, cmd_yaw = self.command
cmd_vx, cmd_vy, cmd_yaw = self.command
cmd_norm = math.hypot(cmd_vx, cmd_vy)
VEL_DEADBAND = 0.04
YAW_DEADBAND = 0.05
raw_speed = math.hypot(linear_vel[0], linear_vel[1])
if raw_speed < VEL_DEADBAND:
filtered_vx, filtered_vy = 0.0, 0.0
filtered_speed = 0.0
else:
filtered_vx, filtered_vy = linear_vel[0], linear_vel[1]
filtered_speed = raw_speed
filtered_vx, filtered_vy = (linear_vel[0], linear_vel[1]) if raw_speed >= 0.04 else (0.0, 0.0)
filtered_speed = raw_speed if raw_speed >= 0.04 else 0.0
raw_yaw_rate = abs(angular_vel[2])
if raw_yaw_rate < YAW_DEADBAND:
filtered_yaw_rate = 0.0
else:
filtered_yaw_rate = angular_vel[2]
filtered_yaw_rate = angular_vel[2] if raw_yaw_rate >= 0.05 else 0.0
height_error = pos[2] - self.target_height
r_height = math.exp(-150.0 * (height_error ** 2))
r_stability = math.exp(-25.0 * (roll**2 + pitch**2))
r_pose = math.exp(-2.0 * np.mean(np.square(current_joints - self.default_joint_angles)))
r_smoothness = math.exp(-0.1 * np.mean(np.square(action - previous_action)))
orientation_error = roll**2 + pitch**2
r_stability = math.exp(-25.0 * orientation_error)
joint_error = np.mean(np.square(current_joints - self.default_joint_angles))
r_pose = math.exp(-2.0 * joint_error)
action_delta = np.mean(np.square(action - previous_action))
r_smoothness = math.exp(-0.1 * action_delta)
r_lin_vel = 0.0
r_ang_vel = 0.0
stillness_penalty = 0.0
gated_zero = 0.0
r_lin_vel, r_ang_vel, stillness_penalty = 0.0, 0.0, 0.0
if cmd_norm < 0.05 and abs(cmd_yaw) < 0.05:
# 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)
+ (w_stability * r_stability)
+ (w_pose * r_pose)
+ (w_smoothness * r_smoothness)
)
w_height, w_stability, w_pose, w_smoothness = 0.35, 0.35, 0.20, 0.10
total_reward = (w_height * r_height) + (w_stability * r_stability) + (w_pose * r_pose) + (w_smoothness * r_smoothness)
else:
# WALKING / TURNING MODE
is_moving = (filtered_speed > 0.0) or (abs(filtered_yaw_rate) > 0.0)
target_vx = cmd_vx * self.max_robot_speed
target_vy = cmd_vy * self.max_robot_speed
target_vx, target_vy = cmd_vx * self.max_robot_speed, cmd_vy * self.max_robot_speed
target_speed = math.hypot(target_vx, target_vy)
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)
r_ang_vel = math.exp(-15.0 * ((filtered_yaw_rate - cmd_yaw)**2))
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_pose = 0.0
w_smoothness = 0.08
w_lin_vel, w_ang_vel, w_height, w_stability, w_smoothness = 0.55, 0.15, 0.10, 0.12, 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
(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
)
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 = {
"height": float(r_height),
"stability": float(r_stability),
"pose": float(r_pose),
"smoothness": float(r_smoothness),
"lin_vel": float(r_lin_vel),
"ang_vel": float(r_ang_vel),
"total": final_reward,
}
self.last_reward_components = comp
for key, val in comp.items():
self.episode_reward_components_sum[key] += val
for k, v in self.last_reward_components.items():
self.episode_reward_components_sum[k] += v
return final_reward
def get_reward_component_averages(self) -> Dict[str, float]:
steps = max(1, self.step_count)
return {k: v / steps for k, v in self.episode_reward_components_sum.items()}
def get_current_robot_metrics(self) -> List[Dict[str, Any]]:
linear_vel, angular_vel = self.sim.get_robot_velocity()
speed = float(math.hypot(linear_vel[0], linear_vel[1]))
yaw_rate = float(abs(angular_vel[2]))
metrics = {
"alive": not self.is_failed,
"phase_name": self.curriculum_phase.name,
"reward": float(self.cumulative_reward),
"speed": speed,
"yaw_rate": yaw_rate,
"distance_from_start": float(self.max_distance_from_start),
"survival_steps": int(self.max_survival_steps),
}
return [metrics]
def _update_hud(self):
if not self.use_gui:
return
now = time.time()
fps = 1.0 / max(now - self.last_time, 1e-5)
self.last_time = now
pos, (roll, pitch, _) = self.sim_manager.get_robot_pose_and_rpy(self.pb_robot)
pos, (roll, pitch, _) = self.sim.get_robot_pose_and_rpy()
self.hud.update(
episode=self.episode_count,
step=self.total_steps,
robot_rewards=[self.robot_reward],
cmd_vel=self.command,
fps=fps,
height=pos[2],
roll_pitch=(math.degrees(roll), math.degrees(pitch))
episode=self.episode_count, step=self.total_steps, robot_rewards=[self.robot_reward],
cmd_vel=self.command, fps=fps, height=pos[2], roll_pitch=(math.degrees(roll), math.degrees(pitch))
)
def _update_robot_failure(self):
if self.is_failed:
if self.is_failed or self.step_count < 15:
return
if self.step_count < 15:
return
position, (roll, pitch, _) = self.sim_manager.get_robot_pose_and_rpy(self.pb_robot)
position, (roll, pitch, _) = self.sim.get_robot_pose_and_rpy()
collapse_threshold = max(0.04, self.collapse_height_fraction * self.target_height)
is_tilted = abs(roll) > self.tilt_failure_rad or abs(pitch) > self.tilt_failure_rad
is_collapsed = position[2] < collapse_threshold
@@ -546,105 +388,7 @@ class JackBotEnv(gym.Env):
if is_tilted or is_collapsed:
self.is_failed = True
if self.use_gui:
self.sim_manager.set_robot_color(self.pb_robot, COLOR_FAILED)
self.sim.set_robot_color(COLOR_FAILED)
def close(self):
self.sim_manager.disconnect()
class CurriculumCallback(BaseCallback):
"""Logs curriculum phase breakdown and best performance metrics to TensorBoard."""
def __init__(self, verbose=0):
super().__init__(verbose)
self.best_speed = 0.0
self.best_yaw_rate = 0.0
self.best_distance = 0.0
self.best_survival_steps = 0.0
self.best_reward = -float('inf')
def _on_step(self) -> bool:
return True
def _on_rollout_end(self) -> bool:
try:
vec_env = self.training_env
alive_metrics = vec_env.env_method("get_current_robot_metrics")
self.best_speed = 0.0
self.best_yaw_rate = 0.0
self.best_distance = 0.0
self.best_survival_steps = 0.0
self.best_reward = -float('inf')
phase_counts = {
"stand_only": 0,
"forward": 0,
"turn_and_direction": 0,
"omni_direction": 0,
"full_command": 0,
}
for worker_res in alive_metrics:
for metrics in worker_res:
if not metrics.get("alive", False):
continue
phase_key = metrics.get("phase_name", "STAND_ONLY").lower()
if phase_key in phase_counts:
phase_counts[phase_key] += 1
if metrics["reward"] > self.best_reward:
self.best_reward = float(metrics["reward"])
if metrics["speed"] > self.best_speed:
self.best_speed = float(metrics["speed"])
if metrics["yaw_rate"] > self.best_yaw_rate:
self.best_yaw_rate = float(metrics["yaw_rate"])
if metrics["distance_from_start"] > self.best_distance:
self.best_distance = float(metrics["distance_from_start"])
if metrics["survival_steps"] > self.best_survival_steps:
self.best_survival_steps = float(metrics["survival_steps"])
for phase_name, count in phase_counts.items():
self.logger.record(f"phase/{phase_name}", count)
self.logger.record("custom/best_reward", float(self.best_reward) if np.isfinite(self.best_reward) else 0.0)
self.logger.record("custom/best_survival_steps", float(self.best_survival_steps))
self.logger.record("custom/best_distance_from_start_m", float(self.best_distance))
self.logger.record("custom/best_speed_mps", float(self.best_speed))
self.logger.record("custom/best_yaw_rate_rads", float(self.best_yaw_rate))
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
self.sim.disconnect()