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

added Robot kinematic use for training

Browse Source 
HUGE BUG -> SimManager physics broken (at least with Robot kinematics)
This commit is contained in:
JackM323
2026-08-06 13:48:08 +02:00
parent c93c524a10
commit 523a4aea89
5 changed files with 502 additions and 177 deletions
Download Patch File Download Diff File Expand all files Collapse all files
ml/env.py
+123 -55
View File
@@ -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