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

Machine Learning Trainer

Browse Source 
Training environment to make a walk model for the hexapod
generated code that will be checked
This commit is contained in:
JackM323
2026-07-30 17:23:36 +02:00
parent bec157a458
commit c5ca79a354
16 changed files with 980 additions and 88 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.gitignore
+2
View File
@@ -3,6 +3,8 @@ node_modules/
.venv/ .venv/
.vs/ .vs/
__pycache__/ __pycache__/
ml/checkpoints/
ml/tensorboard/
# Ignore environment files with private passwords/keys # Ignore environment files with private passwords/keys
.env .env
.vscode/extensions.json
Vendored
+3
View File
@@ -0,0 +1,3 @@
{
"recommendations": []
}
GlobalVariables.py
+39 -81
View File
@@ -1,81 +1,39 @@
import numpy as np import numpy as np
from ArduinoCommunication import ArduinoCommunication from ArduinoCommunication import ArduinoCommunication
from EspCommunication import ESP32Communication from EspCommunication import ESP32Communication
from simulation import Simulation from simulation import Simulation
import DataTypes as dt import DataTypes as dt
import kinematics as kin import config as cfg
import config as cfg from robot_init import init_deg, init90_deg, init_pos, center_points
# Globale Variablen # Globale Variablen
robotCommunication = None robotCommunication = None
shared_sim = None shared_sim = None
emote = None emote = None
display_text = "" display_text = ""
if cfg.sim: if cfg.sim:
shared_sim = Simulation() shared_sim = Simulation()
else: else:
if cfg.arduinoConnection: if cfg.arduinoConnection:
robotCommunication = ArduinoCommunication() robotCommunication = ArduinoCommunication()
else: else:
robotCommunication = ESP32Communication() robotCommunication = ESP32Communication()
vector_dirmov = [0, 0, 0] # Direction Movement [vx, vy, omega] vector_dirmov = [0, 0, 0] # Direction Movement [vx, vy, omega]
current_rad: dt.RadArray # Current Rad current_rad: dt.RadArray # Current Rad
current_pos: dt.PosArray # Current Position current_pos: dt.PosArray # Current Position
# current_deg: dt.DegArray # Current Degrees # current_deg: dt.DegArray # Current Degrees
# legarray: dt.DegArray # Working Leg Array # legarray: dt.DegArray # Working Leg Array
# target_pos: dt.PosArray # Target Position # target_pos: dt.PosArray # Target Position
robot_state = "idle" # Current State robot_state = "idle" # Current State
# Init for 6 Legs # Init for 6 Legs
#leg_state = np.array(["step", "drag", "step", "drag", "step", "drag"]) #leg_state = np.array(["step", "drag", "step", "drag", "step", "drag"])
# Init for 4 Legs # Init for 4 Legs
leg_state = np.array(["step", "drag", "drag", "drag", "drag", "drag"]) leg_state = np.array(["step", "drag", "drag", "drag", "drag", "drag"])
control_pause: bool = False control_pause: bool = False
# Initilize mit Start Position # Initilize mit Start Position
#init_deg: dt.DegArray = dt.DegArray( ## Shared init pose definitions are imported from robot_init.py
# [
# [90, 30, 115],
# [90, 30, 115],
# [90, 30, 115],
# [90, 150, 65],
# [90, 150, 65],
# [90, 150, 65],
# ]
#)
#init_deg: dt.DegArray = dt.DegArray(
# [
# [90, 30, 95],
# [90, 30, 95],
# [90, 30, 95],
# [90, 150, 85],
# [90, 150, 85],
# [90, 150, 85],
# ]
#)
init_deg: dt.DegArray = dt.DegArray(
[
[90, 45, 140],
[90, 45, 140],
[90, 45, 140],
[90, 135, 40],
[90, 135, 40],
[90, 135, 40],
]
)
init90_deg: dt.DegArray = dt.DegArray(
[
[90, 90, 90],
[90, 90, 90],
[90, 90, 90],
[90, 90, 90],
[90, 90, 90],
[90, 90, 90],
]
)
center_points: dt.PosArray = kin.ikpyForward(init_deg.to_rad())
README.md
+19 -7
View File
@@ -6,8 +6,7 @@ It uses IKPy for inverse kinematics, PyBullet for optional simulation, and can s
## Requirements ## Requirements
- Python 3.12 recommended on Windows - Python 3.12 is safest for `pygame` compatibility
- Python 3.11 / 3.12 is safest for `pygame` compatibility
- Required Python packages: - Required Python packages:
- `numpy` - `numpy`
- `pygame` - `pygame`
@@ -18,22 +17,22 @@ It uses IKPy for inverse kinematics, PyBullet for optional simulation, and can s
## Setup ## Setup
### Linux / macOS (Bash) ### Linux (Bash)
```bash ```bash
python3 -m venv .venv python3.12 -m venv .venv
source .venv/bin/activate source .venv/bin/activate
python -m pip install --upgrade pip setuptools wheel python -m pip install --upgrade pip setuptools wheel
pip install numpy pygame ikpy pybullet pyserial matplotlib pip install -r requirements.txt
``` ```
### Windows (PowerShell) ### Windows (PowerShell)
```powershell ```powershell
python -m venv .venv python3.12 -m venv .venv
.\.venv\Scripts\Activate.ps1 .\.venv\Scripts\Activate.ps1
python -m pip install --upgrade pip setuptools wheel python -m pip install --upgrade pip setuptools wheel
pip install numpy pygame ikpy pybullet pyserial matplotlib pip install -r requirements.txt
``` ```
If PowerShell blocks activation: If PowerShell blocks activation:
@@ -51,6 +50,19 @@ From the repository root:
python main.py python main.py
``` ```
If Linux breaks with
```
ExampleBrowserThreadFunc started
X11 functions dynamically loaded using dlopen/dlsym OK!
cannot connect to X server
```
run:
```
export DISPLAY=:0
python main.py
```
## Configuration ## Configuration
Edit `config.py` before running: Edit `config.py` before running:
RobotState/ml_walking.py
+63
View File
@@ -0,0 +1,63 @@
import numpy as np
import math
import torch
import config as cfg
import GlobalVariables as gv
import kinematics as kin
import DataTypes as dt
class MLWalkingState:
def __init__(self, model_path: str | None = None):
self.model_path = model_path or "ml/checkpoints/ppo_joint_command.zip"
self.model = None
self._load_model()
self.step_count = 0
def _load_model(self):
try:
from stable_baselines3 import PPO
except ImportError:
print("stable-baselines3 not installed: ML walking will not be available.")
self.model = None
return
try:
self.model = PPO.load(self.model_path)
print(f"Loaded ML walking model from {self.model_path}")
except Exception as exc:
print(f"Failed to load ML walking model: {exc}")
self.model = None
def infer_joint_commands(self, current_rad: dt.RadArray, direction: np.ndarray) -> dt.RadArray:
if self.model is None:
return current_rad
observation = np.concatenate([current_rad.data.flatten(), direction]).astype(np.float32)
action, _ = self.model.predict(observation, deterministic=True)
action = np.clip(action, -1.0, 1.0).astype(np.float32)
new_rad = np.clip(
current_rad.data.flatten() + action * math.radians(5.0),
-math.pi,
math.pi,
).reshape((6, 3))
return dt.RadArray(new_rad)
def update(self, ctx, intent, dt_step):
if not intent.walk:
return "idle"
direction = np.array([intent.move_vector.x, intent.move_vector.y, 0.0, intent.turn], dtype=np.float32)
target_rad = self.infer_joint_commands(ctx.current_rad, direction)
if ctx.robotCommunication:
ctx.robotCommunication.send_motion(target_rad)
if ctx.shared_sim:
ctx.shared_sim.updatePos(target_rad)
ctx.shared_sim.step()
ctx.current_rad = target_rad
self.step_count += 1
return None
ml/README.md
+59
View File
@@ -0,0 +1,59 @@
JackBot ML training and evaluation
Quick-start
1. Create a Python virtualenv and activate it:
```bash
python -m venv .venv
source .venv/bin/activate
```
2. Install dependencies (GPU users should install `torch` appropriate for their CUDA):
```bash
pip install -r requirements.txt
```
3. Quick training (short, for smoke test):
```bash
python ml/run_train.py --timesteps 50000 --model ml/checkpoints/ppo_joint_command
```
Use multi-robot training with `--num-robots` and `--robot-spacing`:
```bash
python ml/run_train.py --timesteps 50000 --model ml/checkpoints/ppo_joint_command --num-robots 3 --robot-spacing 0.75
```
Choose the start pose at reset with `--start-pose`:
```bash
python ml/run_train.py --timesteps 50000 --model ml/checkpoints/ppo_joint_command --start-pose init_deg
```
If you want to watch the agent train in the PyBullet window, add `--gui`:
```bash
python ml/run_train.py --timesteps 50000 --model ml/checkpoints/ppo_joint_command --gui
```
4. Evaluate the trained model in GUI mode:
```bash
python ml/run_eval.py --model ml/checkpoints/ppo_joint_command.zip --episodes 3 --gui
```
For evaluation with multiple robots and start pose:
```bash
python ml/run_eval.py --model ml/checkpoints/ppo_joint_command.zip --episodes 3 --gui --num-robots 2 --robot-spacing 0.6 --start-pose init_deg
```
Notes
- `ml/env.py` exposes a `JackBotEnv` gym environment that uses your `JackBotUrdf.urdf`.
- The observation is `[joint_angles..., vx, vy, vz, omega]` and the action is per-joint delta in [-1,1].
- Start with small timesteps and in `use_gui=False` to speed up iteration. Increase timesteps and tune the reward for better walking quality.
- Keep the trained models off hardware until they behave well in simulation.
ml/__init__.py
+6
View File
@@ -0,0 +1,6 @@
from .env import JackBotEnv
from .model import ActorCritic
from .train import train
from .evaluate import evaluate
__all__ = ["JackBotEnv", "ActorCritic", "train", "evaluate"]
ml/env.py
+317
View File
@@ -0,0 +1,317 @@
import math
import os
import random
import numpy as np
import pybullet as p
import pybullet_data
import gymnasium as gym
from gymnasium import spaces
import config as cfg
import robot_init as ri
class JackBotEnv(gym.Env):
"""Gymnasium environment for joint-command learning with the JackBot URDF."""
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",
):
"""
Initialize the JackBot environment.
Args:
urdf_path: Path to the robot's URDF file. Defaults to config value.
use_gui: If True, uses p.GUI (shows window), else p.DIRECT (headless).
frame_skip: Number of physics simulation steps per agent step.
Higher values mean the agent takes actions less frequently (e.g., 4 steps at 240Hz = 60Hz control).
random_command: If True, assigns a random velocity command (vx, vy, omega) to each robot every reset.
max_episode_steps: Maximum number of steps (interactions) before the episode is truncated.
num_robots: Number of robots to spawn for parallel training in the same scene.
robot_spacing: Distance in meters between robots when spawning multiple.
start_pose: The initial joint configuration name (e.g., 'init_deg' or 'init90_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)
self.physics_client = None
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.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 _connect_sim(self):
if self.physics_client is not None:
p.disconnect(self.physics_client)
flags = p.GUI if self.use_gui else p.DIRECT
self.physics_client = p.connect(flags)
p.setAdditionalSearchPath(pybullet_data.getDataPath())
p.setGravity(0, 0, -9.81)
self.plane = p.loadURDF("plane.urdf")
self.robots = []
self.robot_joint_indices = []
for robot_id in range(self.num_robots):
base_pos = self._robot_base_position(robot_id)
robot = p.loadURDF(self.urdf_path, basePosition=base_pos, useFixedBase=False)
self.robots.append(robot)
joint_indices = [
i
for i in range(p.getNumJoints(robot))
if p.getJointInfo(robot, i)[2] == p.JOINT_REVOLUTE
]
self.robot_joint_indices.append(joint_indices)
joint_count = len(self.robot_joint_indices[0])
if any(len(indices) != joint_count for indices in self.robot_joint_indices):
raise ValueError("All robots must have the same number of revolute joints.")
lower_limits = []
upper_limits = []
for joint_index in self.robot_joint_indices[0]:
info = p.getJointInfo(self.robots[0], joint_index)
lower = info[8]
upper = info[9]
if lower >= upper:
lower = -math.pi
upper = math.pi
lower_limits.append(lower)
upper_limits.append(upper)
limits = np.array([lower_limits, upper_limits], dtype=np.float32)
self.joint_lower = np.tile(limits[0], (self.num_robots, 1))
self.joint_upper = np.tile(limits[1], (self.num_robots, 1))
pose = ri.init_deg if self.start_pose == "init_deg" else ri.init90_deg
initial_pose = pose.to_rad().data.flatten()
self.initial_angles = np.tile(initial_pose, (self.num_robots, 1)).astype(np.float32)
self.initial_angles = np.clip(
self.initial_angles,
self.joint_lower,
self.joint_upper,
)
self.last_action = np.zeros(self.num_robots * joint_count, dtype=np.float32)
self.commands = np.zeros((self.num_robots, 4), dtype=np.float32)
def reset(self, command: np.ndarray | None = None, seed: int | None = None, options: dict | None = None):
if self._first_reset:
self.step_count = 0
self.last_action = np.zeros(self.num_robots * len(self.robot_joint_indices[0]), dtype=np.float32)
if seed is not None:
self.seed(seed)
if command is not None:
self.commands = np.array(command, dtype=np.float32).reshape(self.num_robots, 4)
elif self.random_command:
self.commands = np.stack([self.sample_command() for _ in range(self.num_robots)])
else:
self.commands = np.zeros((self.num_robots, 4), dtype=np.float32)
self._reset_pose()
self._first_reset = False
return self._get_obs(), {}
p.resetSimulation()
p.setGravity(0, 0, -9.81)
p.setAdditionalSearchPath(pybullet_data.getDataPath())
self.plane = p.loadURDF("plane.urdf")
self.robots = []
self.robot_joint_indices = []
for robot_id in range(self.num_robots):
base_pos = self._robot_base_position(robot_id)
robot = p.loadURDF(self.urdf_path, basePosition=base_pos, useFixedBase=False)
self.robots.append(robot)
joint_indices = [
i
for i in range(p.getNumJoints(robot))
if p.getJointInfo(robot, i)[2] == p.JOINT_REVOLUTE
]
self.robot_joint_indices.append(joint_indices)
self.step_count = 0
self.last_action = np.zeros(self.num_robots * len(self.robot_joint_indices[0]), dtype=np.float32)
if seed is not None:
self.seed(seed)
if command is not None:
self.commands = np.array(command, dtype=np.float32).reshape(self.num_robots, 4)
elif self.random_command:
self.commands = np.stack([self.sample_command() for _ in range(self.num_robots)])
else:
self.commands = np.zeros((self.num_robots, 4), dtype=np.float32)
self._reset_pose()
return self._get_obs(), {}
def seed(self, seed: int | None = None):
"""Set RNG seed for reproducibility (Gym compatibility)."""
if seed is None:
seed = int(np.random.randint(0, 2**31 - 1))
random.seed(seed)
np.random.seed(seed)
self._seed = seed
return [seed]
def _reset_pose(self):
for robot, joint_indices, init_angles in zip(self.robots, self.robot_joint_indices, self.initial_angles):
for joint_index, target_angle in zip(joint_indices, init_angles):
p.resetJointState(robot, joint_index, target_angle)
p.setJointMotorControl2(
bodyIndex=robot,
jointIndex=joint_index,
controlMode=p.POSITION_CONTROL,
targetPosition=target_angle,
force=250,
)
for _ in range(400):
p.stepSimulation()
self.joint_angles = self.initial_angles.copy()
def sample_command(self) -> np.ndarray:
vx = np.random.uniform(-1.0, 1.0)
vy = np.random.uniform(-0.5, 0.5)
vz = 0.0
omega = np.random.uniform(-1.0, 1.0)
return np.array([vx, vy, vz, omega], dtype=np.float32)
def _get_obs(self) -> np.ndarray:
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,
)
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,
)
for _ in range(self.frame_skip):
p.stepSimulation()
self.step_count += 1
observation = self._get_obs()
reward = self._compute_reward()
done = self._is_done()
if self.step_count <= self._min_steps_before_done:
done = False
info = {"step": self.step_count}
# Gymnasium-style return: (obs, reward, terminated, truncated, info)
if done:
terminated = True
truncated = False
else:
terminated = False
truncated = False
return observation, float(reward), terminated, truncated, info
def _compute_reward(self) -> float:
rewards = []
for robot_id, robot in enumerate(self.robots):
linear_vel, angular_vel = p.getBaseVelocity(robot)
_, orientation = p.getBasePositionAndOrientation(robot)
roll, pitch, _ = p.getEulerFromQuaternion(orientation)
command = self.commands[robot_id]
forward_reward = command[0] * linear_vel[0] + command[1] * linear_vel[1]
rotation_reward = command[3] * angular_vel[2]
stability_penalty = abs(roll) + abs(pitch)
action_penalty = float(np.sum(np.square(self.last_action.reshape(self.num_robots, -1)[robot_id]))) * 0.01
rewards.append(0.1 + forward_reward + rotation_reward - 0.2 * stability_penalty - action_penalty)
return float(np.sum(rewards))
def _is_done(self) -> bool:
for robot in self.robots:
_, orientation = p.getBasePositionAndOrientation(robot)
roll, pitch, _ = p.getEulerFromQuaternion(orientation)
if abs(roll) > 0.7 or abs(pitch) > 0.7:
return True
position, _ = p.getBasePositionAndOrientation(robot)
if position[2] < 0.02:
return True
if self.step_count >= self.max_episode_steps:
return True
return False
def render(self, mode="human"):
if self.use_gui:
return None
raise NotImplementedError("Render is only supported with use_gui=True")
def _robot_base_position(self, robot_id: int) -> list[float]:
cols = int(math.sqrt(self.num_robots - 1)) + 1
row = robot_id // cols
col = robot_id % cols
x = (col - (cols - 1) / 2.0) * self.robot_spacing
y = (row - (cols - 1) / 2.0) * self.robot_spacing
return [x, y, 0.1]
def close(self):
if self.physics_client is not None:
p.disconnect(self.physics_client)
self.physics_client = None
ml/evaluate.py
+78
View File
@@ -0,0 +1,78 @@
import argparse
import numpy as np
from .env import JackBotEnv
def evaluate(
model_path: str,
episodes: int = 5,
use_gui: bool = False,
num_robots: int = 1,
robot_spacing: float = 0.5,
start_pose: str = "init_deg",
):
try:
from stable_baselines3 import PPO
except ImportError as exc:
raise ImportError(
"stable-baselines3 is required for evaluation. Install with: pip install stable-baselines3"
) from exc
env = JackBotEnv(
use_gui=use_gui,
random_command=False,
num_robots=num_robots,
robot_spacing=robot_spacing,
start_pose=start_pose,
)
model = PPO.load(model_path)
for episode in range(episodes):
reset_res = env.reset()
# handle Gym / Gymnasium compatibility: reset may return (obs, info)
if isinstance(reset_res, tuple) and len(reset_res) == 2:
obs, _ = reset_res
else:
obs = reset_res
done = False
episode_reward = 0.0
while not done:
# pass only the observation to the policy
action, _ = model.predict(obs, deterministic=True)
step_res = env.step(action)
# Gymnasium-style: (obs, reward, terminated, truncated, info)
if isinstance(step_res, tuple) and len(step_res) == 5:
obs, reward, terminated, truncated, info = step_res
done = bool(terminated or truncated)
else:
# legacy Gym: (obs, reward, done, info)
obs, reward, done, info = step_res
episode_reward += float(reward)
print(f"Episode {episode + 1}: reward={episode_reward:.2f}")
env.close()
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="Evaluate a trained JackBot policy.")
parser.add_argument("--model-path", type=str, required=True)
parser.add_argument("--episodes", type=int, default=5)
parser.add_argument("--gui", action="store_true")
parser.add_argument("--num-robots", type=int, default=1, help="Number of robots in the environment")
parser.add_argument("--robot-spacing", type=float, default=0.5, help="Spacing between robots in meters")
parser.add_argument("--start-pose", type=str, choices=["init_deg", "init90_deg"], default="init_deg", help="Initial robot pose at reset")
args = parser.parse_args()
evaluate(
args.model_path,
episodes=args.episodes,
use_gui=args.gui,
num_robots=args.num_robots,
robot_spacing=args.robot_spacing,
start_pose=args.start_pose,
)
ml/model.py
+49
View File
@@ -0,0 +1,49 @@
import torch
import torch.nn as nn
from torch.distributions import Normal
class ActorCritic(nn.Module):
def __init__(self, obs_dim: int, action_dim: int, hidden_sizes=(256, 256)):
super().__init__()
self.backbone = nn.Sequential(
nn.Linear(obs_dim, hidden_sizes[0]),
nn.ReLU(),
nn.Linear(hidden_sizes[0], hidden_sizes[1]),
nn.ReLU(),
)
self.mean_head = nn.Linear(hidden_sizes[1], action_dim)
self.value_head = nn.Linear(hidden_sizes[1], 1)
self.log_std = nn.Parameter(torch.zeros(action_dim, dtype=torch.float32))
def forward(self, obs: torch.Tensor):
x = self.backbone(obs)
mean = self.mean_head(x)
std = self.log_std.exp()
value = self.value_head(x).squeeze(-1)
return mean, std, value
def get_action(self, obs: torch.Tensor):
mean, std, value = self.forward(obs)
dist = Normal(mean, std)
action = dist.sample()
log_prob = dist.log_prob(action).sum(-1)
return action, log_prob, value
def evaluate_actions(self, obs: torch.Tensor, actions: torch.Tensor):
mean, std, value = self.forward(obs)
dist = Normal(mean, std)
log_prob = dist.log_prob(actions).sum(-1)
entropy = dist.entropy().sum(-1)
return value, log_prob, entropy
def save(self, path: str):
torch.save(self.state_dict(), path)
@classmethod
def load(cls, path: str, obs_dim: int, action_dim: int, hidden_sizes=(256, 256)):
model = cls(obs_dim, action_dim, hidden_sizes)
model.load_state_dict(torch.load(path, map_location=torch.device("cpu")))
model.eval()
return model
ml/run_eval.py
+39
View File
@@ -0,0 +1,39 @@
"""Run a trained policy in the PyBullet sim for quick inspection.
Usage:
python ml/run_eval.py --model ml/checkpoints/ppo_joint_command.zip --episodes 3 --gui
"""
import argparse
import sys
from pathlib import Path
ROOT_DIR = Path(__file__).resolve().parent.parent
if str(ROOT_DIR) not in sys.path:
sys.path.insert(0, str(ROOT_DIR))
from ml.evaluate import evaluate
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--model", type=str, required=True, help="Path to the trained model file (.zip)")
parser.add_argument("--episodes", type=int, default=3, help="Number of 'rounds' to run. One episode lasts from reset until the robot falls over or the time limit is reached.")
parser.add_argument("--gui", action="store_true", help="Show the PyBullet GUI during evaluation")
parser.add_argument("--num-robots", type=int, default=1, help="Number of robots in the evaluation environment")
parser.add_argument("--robot-spacing", type=float, default=0.5, help="Spacing between robots in meters")
parser.add_argument("--start-pose", type=str, choices=["init_deg", "init90_deg"], default="init_deg", help="Initial robot pose at reset")
args = parser.parse_args()
evaluate(
model_path=args.model,
episodes=args.episodes,
use_gui=args.gui,
num_robots=args.num_robots,
robot_spacing=args.robot_spacing,
start_pose=args.start_pose,
)
if __name__ == "__main__":
main()
ml/run_train.py
+47
View File
@@ -0,0 +1,47 @@
"""Minimal training launcher for quick experiments.
Usage:
python ml/run_train.py --timesteps 50000 --model ml/checkpoints/ppo_joint_command
This is a convenience wrapper around `ml.train.train` with friendly defaults
for interactive experimentation.
"""
import argparse
import sys
from pathlib import Path
ROOT_DIR = Path(__file__).resolve().parent.parent
if str(ROOT_DIR) not in sys.path:
sys.path.insert(0, str(ROOT_DIR))
from ml.train import train
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--timesteps", type=int, default=50000, help="Total number of 'practice steps'.")
parser.add_argument("--model", type=str, default="ml/checkpoints/ppo_joint_command", help="Path to save the trained model")
parser.add_argument("--seed", type=int, default=0, help="Random seed for reproducibility")
parser.add_argument("--device", type=str, default="auto", help="Device to use: 'cpu', 'cuda', or 'auto'")
parser.add_argument("--gui", action="store_true", help="Show the PyBullet GUI during training")
parser.add_argument("--num-robots", type=int, default=1, help="Number of robots in the training environment")
parser.add_argument("--robot-spacing", type=float, default=0.5, help="Spacing between robots in meters")
parser.add_argument("--start-pose", type=str, choices=["init_deg", "init90_deg"], default="init_deg", help="Initial robot pose at reset")
args = parser.parse_args()
Path(args.model).parent.mkdir(parents=True, exist_ok=True)
train(
total_timesteps=args.timesteps,
model_path=args.model,
seed=args.seed,
device=args.device,
use_gui=args.gui,
num_robots=args.num_robots,
robot_spacing=args.robot_spacing,
start_pose=args.start_pose,
)
if __name__ == "__main__":
main()
ml/sim_manager.py
+105
View File
@@ -0,0 +1,105 @@
from typing import Dict, Any, List, Tuple
import pybullet as p
import pybullet_data
import numpy as np
class SimManager:
"""
Manages the PyBullet simulation lifecycle and live HUD overlays.
"""
def __init__(self, use_gui: bool = True):
self.use_gui = use_gui
self.physics_client = None
self.debug_text_ids: Dict[str, int] = {}
def connect(self):
"""Connects to PyBullet and sets up the basic physics world."""
if self.physics_client is not None:
p.disconnect(self.physics_client)
flags = p.GUI if self.use_gui else p.DIRECT
self.physics_client = p.connect(flags)
p.setAdditionalSearchPath(pybullet_data.getDataPath())
p.setGravity(0, 0, -9.81)
if self.use_gui:
# Disable unnecessary side panels for a clean UI
p.configureDebugVisualizer(p.COV_ENABLE_GUI, 0)
p.configureDebugVisualizer(p.COV_ENABLE_SEGMENTATION_MARK_PREVIEW, 0)
p.configureDebugVisualizer(p.COV_ENABLE_DEPTH_BUFFER_PREVIEW, 0)
p.configureDebugVisualizer(p.COV_ENABLE_RGB_OUTPUT_PREVIEW, 0)
def load_scene(
self, urdf_path: str, num_robots: int, robot_spacing: float, base_pos_fn
) -> Tuple[int, List[int], List[List[int]]]:
"""Loads plane and robots into the active PyBullet simulation."""
plane_id = p.loadURDF("plane.urdf")
robots = []
robot_joint_indices = []
for robot_id in range(num_robots):
base_pos = base_pos_fn(robot_id, num_robots, robot_spacing)
robot = p.loadURDF(urdf_path, basePosition=base_pos, useFixedBase=False)
robots.append(robot)
joint_indices = [
i
for i in range(p.getNumJoints(robot))
if p.getJointInfo(robot, i)[2] == p.JOINT_REVOLUTE
]
robot_joint_indices.append(joint_indices)
return plane_id, robots, robot_joint_indices
def disconnect(self):
"""Safely disconnects from the simulation."""
if self.physics_client is not None:
p.disconnect(self.physics_client)
self.physics_client = None
self.debug_text_ids.clear()
def update_hud(self, stats: Dict[str, Any]):
"""
Renders live training metrics onto the PyBullet 3D viewport.
Uses replaceItemUniqueId to prevent flicker.
"""
if not self.use_gui or self.physics_client is None:
return
# Fixed position near the top-left of the origin in 3D world coordinates
x_pos, y_pos, z_start = -1.2, -1.2, 1.6
line_height = 0.10
for i, (label, value) in enumerate(stats.items()):
if isinstance(value, float):
display_text = f"{label}: {value:.3f}"
elif isinstance(value, (list, np.ndarray)):
formatted_vals = ", ".join(f"{v:.2f}" for v in np.atleast_1d(value))
display_text = f"{label}: [{formatted_vals}]"
else:
display_text = f"{label}: {value}"
color = [0, 0, 0] # Black text for clear visibility against the light plane
if label in self.debug_text_ids:
p.addUserDebugText(
display_text,
[x_pos, y_pos, z_start - i * line_height],
textColorRGB=color,
textSize=1.1,
replaceItemUniqueId=self.debug_text_ids[label],
)
else:
self.debug_text_ids[label] = p.addUserDebugText(
display_text,
[x_pos, y_pos, z_start - i * line_height],
textColorRGB=color,
textSize=1.1,
)
def reset_hud(self):
"""Clears debug text tracking."""
self.debug_text_ids.clear()
ml/train.py
+113
View File
@@ -0,0 +1,113 @@
import argparse
import os
from pathlib import Path
from .env import JackBotEnv
def parse_args():
parser = argparse.ArgumentParser(description="Train a joint-command policy for JackBot.")
parser.add_argument("--timesteps", type=int, default=500_000, help="Total training timesteps")
parser.add_argument("--model-path", type=str, default="ml/checkpoints/ppo_joint_command", help="Where to save the trained model")
parser.add_argument("--seed", type=int, default=0, help="Random seed")
parser.add_argument("--device", type=str, default="auto", help="Device to use: 'cpu', 'cuda', or 'auto' to autodetect")
parser.add_argument("--use-gui", action="store_true", help="Enable PyBullet GUI during training")
parser.add_argument("--num-robots", type=int, default=1, help="Number of robots in the training environment")
parser.add_argument("--robot-spacing", type=float, default=0.5, help="Spacing between robots in meters")
parser.add_argument("--start-pose", type=str, choices=["init_deg", "init90_deg"], default="init_deg", help="Initial robot pose at reset")
return parser.parse_args()
def train(
total_timesteps: int,
model_path: str,
seed: int = 0,
device: str = "auto",
use_gui: bool = False,
num_robots: int = 1,
robot_spacing: float = 0.5,
start_pose: str = "init_deg",
):
try:
from stable_baselines3 import PPO
from stable_baselines3.common.vec_env import DummyVecEnv
except ImportError as exc:
raise ImportError(
"stable-baselines3 and gym are required for training. "
"Install them with: pip install stable-baselines3 gym"
) from exc
env = DummyVecEnv([
lambda: JackBotEnv(
use_gui=use_gui,
random_command=True,
num_robots=num_robots,
robot_spacing=robot_spacing,
start_pose=start_pose,
)
])
def resolve_device(requested_device: str) -> str:
try:
import torch
except ImportError:
if requested_device != "cpu":
raise RuntimeError(
"PyTorch is not installed in the active environment. "
"Install torch with a GPU-enabled build before using --device cuda."
)
return "cpu"
hip_supported = getattr(torch.version, "hip", None) is not None
cuda_available = torch.cuda.is_available()
hip_available = hip_supported and getattr(torch.backends, "hip", None) is not None and torch.backends.hip.is_available()
if requested_device == "auto":
if hip_available or cuda_available:
return "cuda"
return "cpu"
if requested_device in {"cuda", "gpu", "hip"}:
if hip_available or cuda_available:
return "cuda"
raise RuntimeError(
f"GPU device requested ({requested_device}) but no CUDA/ROCm-capable PyTorch is available. "
f"Installed torch build: {torch.__version__} (hip={getattr(torch.version, 'hip', None)}, cuda={cuda_available})"
)
if requested_device == "cpu":
return "cpu"
raise ValueError(
f"Unsupported device '{requested_device}'. Use 'cpu', 'cuda', or 'auto'."
)
device = resolve_device(device)
model = PPO(
"MlpPolicy",
env,
verbose=1,
seed=seed,
device=device,
tensorboard_log=str(Path(__file__).resolve().parent / "tensorboard"),
)
model.learn(total_timesteps=total_timesteps)
Path(model_path).parent.mkdir(parents=True, exist_ok=True)
model.save(model_path)
env.close()
if __name__ == "__main__":
args = parse_args()
train(
args.timesteps,
args.model_path,
seed=args.seed,
device=args.device,
use_gui=args.use_gui,
num_robots=args.num_robots,
robot_spacing=args.robot_spacing,
start_pose=args.start_pose,
)
requirements.txt
+14
View File
@@ -0,0 +1,14 @@
numpy
pygame
ikpy
pybullet
pyserial
matplotlib
# Deep learning / RL
torch
stable-baselines3
gym
gymnasium[box2d]
shimmy
tensorboard
robot_init.py
+27
View File
@@ -0,0 +1,27 @@
import kinematics as kin
import DataTypes as dt
init_deg: dt.DegArray = dt.DegArray(
[
[90, 45, 140],
[90, 45, 140],
[90, 45, 140],
[90, 135, 40],
[90, 135, 40],
[90, 135, 40],
]
)
init90_deg: dt.DegArray = dt.DegArray(
[
[90, 90, 90],
[90, 90, 90],
[90, 90, 90],
[90, 90, 90],
[90, 90, 90],
[90, 90, 90],
]
)
init_pos: dt.RadArray = init_deg.to_rad()
center_points: dt.PosArray = kin.ikpyForward(init_pos)