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

env cleanup

Browse Source 
rewards adjustment
pybullet logic contained in SimManager
This commit is contained in:
JackM323
2026-08-04 21:03:53 +02:00
parent 766f2855da
commit 402c20dfb5
4 changed files with 162 additions and 153 deletions
Download Patch File Download Diff File Expand all files Collapse all files
ml/env.py
+62 -135
View File
@@ -8,7 +8,6 @@ from typing import Optional, Tuple, Dict, Any, List
import gymnasium as gym
from gymnasium import spaces
import numpy as np
import pybullet as p
from config import cfg
from Robot import Robot, PyBulletBackend
@@ -79,9 +78,7 @@ class JackBotEnv(gym.Env):
self.collapse_height_fraction = 0.55
self.tilt_failure_rad = 0.9
# Small random exploration pulse to break local optima.
# The bonus is only granted on a rare step and only when the robot
# is still stable enough that a bold move is likely to remain safe.
# Small random exploration pulse settings
self.exploration_bonus_prob = 0.03
self.exploration_bonus_interval = 120
self.exploration_bonus_scale = 0.08
@@ -93,7 +90,7 @@ class JackBotEnv(gym.Env):
self.curriculum_phase = 0
self.curriculum_episode_limit = 150
self.curriculum_stage_requirements = {
1: {"survival_steps": 200, "distance": 0.00, "stability_roll_pitch": 0.35},
1: {"survival_steps": 1000, "distance": 0.00, "stability_roll_pitch": 0.35},
2: {"survival_steps": 350, "distance": 10.0, "stability_roll_pitch": 0.30},
3: {"survival_steps": 550, "distance": 15.00, "stability_roll_pitch": 0.25},
}
@@ -103,13 +100,6 @@ class JackBotEnv(gym.Env):
self.leader_crown = LeaderCrown(physics_client_id=self.sim_manager.physics_client)
self.last_time = time.time()
def _set_robot_color(self, pb_id: int, rgba: List[float]):
"""Helper to change the visual color of a robot body and all its links."""
num_joints = p.getNumJoints(pb_id, physicsClientId=self.sim_manager.physics_client)
p.changeVisualShape(pb_id, -1, rgbaColor=rgba, physicsClientId=self.sim_manager.physics_client)
for j in range(num_joints):
p.changeVisualShape(pb_id, j, rgbaColor=rgba, physicsClientId=self.sim_manager.physics_client)
def _robot_base_position(self, robot_id: int, spacing: float = 0.5) -> list[float]:
return [0.0, 0.0, 0.14]
@@ -148,20 +138,13 @@ class JackBotEnv(gym.Env):
for idx, (pb_id, robot_obj) in enumerate(zip(self.pb_robots, self.robots)):
spawn_pos = self._robot_base_position(idx, self.robot_spacing)
spawn_orn = [0, 0, 0, 1]
p.resetBasePositionAndOrientation(
pb_id, spawn_pos, spawn_orn, physicsClientId=self.sim_manager.physics_client
)
p.resetBaseVelocity(
pb_id, linearVelocity=[0, 0, 0], angularVelocity=[0, 0, 0],
physicsClientId=self.sim_manager.physics_client
)
spawn_orn = [0.0, 0.0, 0.0, 1.0]
self.sim_manager.reset_robot_base(pb_id, spawn_pos, spawn_orn)
robot_obj.reset_to_init()
if self.use_gui:
self._set_robot_color(pb_id, [1.0, 1.0, 1.0, 1.0])
self.sim_manager.set_robot_color(pb_id, [1.0, 1.0, 1.0, 1.0])
self.last_action = np.zeros(self.action_space.shape[0], dtype=np.float32)
self.max_distance_from_start = [0.0]
@@ -178,18 +161,14 @@ class JackBotEnv(gym.Env):
else:
self.commands = np.zeros((1, 4), dtype=np.float32)
for idx, (pb_id, robot_obj) in enumerate(zip(self.pb_robots, self.robots)):
pos, _ = p.getBasePositionAndOrientation(
pb_id, physicsClientId=self.sim_manager.physics_client
)
for idx, pb_id in enumerate(self.pb_robots):
pos, _ = self.sim_manager.get_robot_pose(pb_id)
self.start_positions[idx] = [float(pos[0]), float(pos[1]), float(pos[2])]
for _ in range(200):
self.sim_manager.step()
self.target_height = self._measure_settled_height()
if self.target_height <= 0.0:
self.target_height = 0.14
# Settle for 200 steps after dropping in a standing position, then measure target height
self.target_height = self.sim_manager.settle_and_measure_height(
self.pb_robots, steps=200, fallback_height=0.14
)
if self.use_gui:
self.hud.reset()
@@ -199,83 +178,63 @@ class JackBotEnv(gym.Env):
return self._get_obs(), {}
def get_robot_velocities(self) -> list:
"""Exposes velocities for the SB3 metrics callback."""
vels = []
for pb_id in self.pb_robots:
lin_v, ang_v = p.getBaseVelocity(pb_id, physicsClientId=self.sim_manager.physics_client)
vels.append((lin_v, ang_v))
return vels
"""Exposes velocities for the SB3 metrics callback."""
vels = []
for pb_id in self.pb_robots:
lin_v, ang_v = self.sim_manager.get_robot_velocity(pb_id)
vels.append((lin_v, ang_v))
return vels
def get_robot_distance_metrics(self) -> list:
"""Exposes distance-from-start metrics for logging without affecting reward."""
metrics = []
for idx, pb_id in enumerate(self.pb_robots):
pos, _ = p.getBasePositionAndOrientation(
pb_id, physicsClientId=self.sim_manager.physics_client
)
start_x, start_y, _ = self.start_positions[idx]
dist = float(np.linalg.norm(np.array([pos[0] - start_x, pos[1] - start_y], dtype=np.float32)))
self.max_distance_from_start[idx] = max(self.max_distance_from_start[idx], dist)
metrics.append((dist, self.max_distance_from_start[idx]))
"""Exposes distance-from-start metrics for logging without affecting reward."""
metrics = []
for idx, pb_id in enumerate(self.pb_robots):
pos, _ = self.sim_manager.get_robot_pose(pb_id)
start_x, start_y, _ = self.start_positions[idx]
dist = float(np.linalg.norm(np.array([pos[0] - start_x, pos[1] - start_y], dtype=np.float32)))
self.max_distance_from_start[idx] = max(self.max_distance_from_start[idx], dist)
metrics.append((dist, self.max_distance_from_start[idx]))
self.max_survival_steps = max(self.max_survival_steps, self.step_count)
return metrics
self.max_survival_steps = max(self.max_survival_steps, self.step_count)
return metrics
def get_current_robot_metrics(self) -> list:
"""Returns current per-robot reward, distance, and velocity summaries for alive robots only."""
metrics = []
for idx, pb_id in enumerate(self.pb_robots):
if self.failed_robots_mask[idx]:
continue
"""Returns current per-robot reward, distance, and velocity summaries for alive robots only."""
metrics = []
for idx, pb_id in enumerate(self.pb_robots):
if self.failed_robots_mask[idx]:
continue
pos, _ = p.getBasePositionAndOrientation(
pb_id, physicsClientId=self.sim_manager.physics_client
)
linear_vel, angular_vel = p.getBaseVelocity(
pb_id, physicsClientId=self.sim_manager.physics_client
)
start_x, start_y, _ = self.start_positions[idx]
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]))
metrics.append({
"reward": float(self.robot_rewards[idx]),
"distance_from_start": dist,
"speed": speed,
"yaw_rate": yaw_rate,
"alive": True,
"survival_steps": int(self.step_count),
})
pos, _ = self.sim_manager.get_robot_pose(pb_id)
linear_vel, angular_vel = self.sim_manager.get_robot_velocity(pb_id)
start_x, start_y, _ = self.start_positions[idx]
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]))
metrics.append({
"reward": float(self.robot_rewards[idx]),
"distance_from_start": dist,
"speed": speed,
"yaw_rate": yaw_rate,
"alive": True,
"survival_steps": int(self.step_count),
})
return metrics
return metrics
def get_survival_steps(self) -> int:
"""Returns the current survival length for the environment's current episode."""
return int(self.step_count)
"""Returns the current survival length for the environment's current episode."""
return int(self.step_count)
def _get_obs(self) -> np.ndarray:
obs_list = []
for idx, (pb_id, joint_indices) in enumerate(zip(self.pb_robots, self.robot_joint_indices)):
joint_states = p.getJointStates(
pb_id,
joint_indices,
physicsClientId=self.sim_manager.physics_client
)
joint_angles = np.array([state[0] for state in joint_states], dtype=np.float32)
joint_angles = self.sim_manager.get_robot_joint_angles(pb_id, joint_indices)
robot_obs = np.concatenate([joint_angles, self.commands[idx]])
obs_list.append(robot_obs)
return np.concatenate(obs_list).astype(np.float32)
def _measure_settled_height(self) -> float:
heights = []
for pb_id in self.pb_robots:
pos, _ = p.getBasePositionAndOrientation(
pb_id, physicsClientId=self.sim_manager.physics_client
)
heights.append(pos[2])
return float(np.mean(heights)) if heights else 0.14
def step(self, action: np.ndarray) -> Tuple[np.ndarray, float, bool, bool, Dict[str, Any]]:
self.step_count += 1
self.total_steps += 1
@@ -322,10 +281,7 @@ class JackBotEnv(gym.Env):
survival_ok = self.max_survival_steps >= req["survival_steps"]
distance_ok = self.max_distance_from_start[0] >= req["distance"]
position, orientation = p.getBasePositionAndOrientation(
self.pb_robots[0], physicsClientId=self.sim_manager.physics_client
)
roll, pitch, _ = p.getEulerFromQuaternion(orientation)
position, (roll, pitch, _) = self.sim_manager.get_robot_pose_and_rpy(self.pb_robots[0])
stability_ok = abs(roll) <= req["stability_roll_pitch"] and abs(pitch) <= req["stability_roll_pitch"]
height_ok = position[2] >= max(0.09, self.target_height * 0.85)
@@ -363,78 +319,55 @@ class JackBotEnv(gym.Env):
previous_actions = previous_action.reshape(1, 18)
for idx, pb_id in enumerate(self.pb_robots):
pos, orientation = p.getBasePositionAndOrientation(
pb_id, physicsClientId=self.sim_manager.physics_client
)
linear_vel, angular_vel = p.getBaseVelocity(
pb_id, physicsClientId=self.sim_manager.physics_client
)
roll, pitch, _ = p.getEulerFromQuaternion(orientation)
pos, (roll, pitch, _) = self.sim_manager.get_robot_pose_and_rpy(pb_id)
linear_vel, angular_vel = self.sim_manager.get_robot_velocity(pb_id)
command = self.commands[idx]
cmd_vx = command[0]
cmd_vy = command[1]
cmd_yaw = command[3]
# -------------------------------------------------------------
# 1. LINEAR VECTOR SPEED MAXIMIZATION (Magnitude + Direction)
# -------------------------------------------------------------
# 1. LINEAR VECTOR SPEED MAXIMIZATION
cmd_dir = np.array([cmd_vx, cmd_vy], dtype=np.float32)
cmd_norm = np.linalg.norm(cmd_dir)
if cmd_norm > 0.05:
# Normalize target direction vector
unit_cmd_dir = cmd_dir / cmd_norm
actual_vel_2d = np.array([linear_vel[0], linear_vel[1]], dtype=np.float32)
# Speed aligned with target direction (m/s)
aligned_speed = float(np.dot(actual_vel_2d, unit_cmd_dir))
# Moderate movement reward to encourage directional motion
linear_speed_reward = 2.5 * aligned_speed
# Penalize sideways drift (perpendicular velocity to commanded direction)
perp_vel = actual_vel_2d - aligned_speed * unit_cmd_dir
drift_penalty = 0.35 * float(np.dot(perp_vel, perp_vel))
else:
# If no linear command given, penalize all horizontal movement
linear_speed_reward = 0.0
drift_penalty = 1.0 * (linear_vel[0]**2 + linear_vel[1]**2)
# -------------------------------------------------------------
# 2. TURNING SPEED MAXIMIZATION
# -------------------------------------------------------------
actual_yaw_rate = angular_vel[2] # rad/s in PyBullet Z-axis
actual_yaw_rate = angular_vel[2]
if abs(cmd_yaw) > 0.05:
# Reward turning in the commanded direction, but less aggressively
turning_reward = 1.5 * (actual_yaw_rate * cmd_yaw)
else:
# Penalize unwanted rotation when joystick turn is centered
turning_reward = -0.6 * (actual_yaw_rate ** 2)
# -------------------------------------------------------------
# 3. SMALL ALIVE BONUS + AGE-RAMPED STILLNESS PENALTY
# -------------------------------------------------------------
alive_reward = 0.02
alive_reward = 0.5
age_ratio = min(1.0, self.step_count / max(1, self.curriculum_episode_limit))
still_penalty = 0.0
if (cmd_norm > 0.1 or abs(cmd_yaw) > 0.1) and (abs(linear_vel[0]) < 0.02 and abs(actual_yaw_rate) < 0.05):
still_penalty = 0.35 + 0.85 * age_ratio
# -------------------------------------------------------------
# 4. POSTURE & STABILITY PENALTIES
# -------------------------------------------------------------
height_penalty = 6.0 * ((self.target_height - pos[2]) ** 2) if pos[2] < self.target_height else 0.0
height_penalty = 10.0 * ((self.target_height - pos[2]) ** 2) if pos[2] < self.target_height else 0.0
stability_penalty = 1.5 * (roll**2 + pitch**2)
# Penalize only large control jumps; allow continuous low-amplitude motion
print(-height_penalty)
control_delta = np.abs(current_actions[idx] - previous_actions[idx])
large_delta_mask = control_delta > 0.12
large_delta_penalty = 0.002 * float(np.sum(np.square(control_delta[large_delta_mask]))) if np.any(large_delta_mask) else 0.0
# Combined Step Reward
r_step = (
alive_reward
+ linear_speed_reward
@@ -465,8 +398,7 @@ class JackBotEnv(gym.Env):
rolls = []
pitches = []
for pb_id in self.pb_robots:
pos, orient = p.getBasePositionAndOrientation(pb_id, physicsClientId=self.sim_manager.physics_client)
roll, pitch, _ = p.getEulerFromQuaternion(orient)
pos, (roll, pitch, _) = self.sim_manager.get_robot_pose_and_rpy(pb_id)
heights.append(pos[2])
rolls.append(math.degrees(roll))
pitches.append(math.degrees(pitch))
@@ -490,9 +422,7 @@ class JackBotEnv(gym.Env):
best_idx = int(np.argmax(self.robot_rewards))
leader_pb_id = self.pb_robots[best_idx]
leader_pos, _ = p.getBasePositionAndOrientation(
leader_pb_id, physicsClientId=self.sim_manager.physics_client
)
leader_pos, _ = self.sim_manager.get_robot_pose(leader_pb_id)
self.leader_crown.update(leader_pos)
def _update_robot_failures(self):
@@ -501,10 +431,7 @@ class JackBotEnv(gym.Env):
if self.failed_robots_mask[idx]:
continue
position, orientation = p.getBasePositionAndOrientation(
pb_id, physicsClientId=self.sim_manager.physics_client
)
roll, pitch, _ = p.getEulerFromQuaternion(orientation)
position, (roll, pitch, _) = self.sim_manager.get_robot_pose_and_rpy(pb_id)
collapse_threshold = max(0.06, self.collapse_height_fraction * self.target_height)
is_tilted = abs(roll) > self.tilt_failure_rad or abs(pitch) > self.tilt_failure_rad
@@ -513,7 +440,7 @@ class JackBotEnv(gym.Env):
if is_tilted or is_collapsed:
self.failed_robots_mask[idx] = True
if self.use_gui:
self._set_robot_color(pb_id, COLOR_FAILED)
self.sim_manager.set_robot_color(pb_id, COLOR_FAILED)
def close(self):
self.sim_manager.disconnect()