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

Compare commits

base: JackM323/JackBot:bec157a458d5ed076c6804cfd35e973b7f6d909d
Branches Tags
JackM323/JackBot:main
..
compare: JackM323/JackBot:5d598b4d94479c717789721b12486babc136b018
Branches Tags
JackM323/JackBot:main

2 Commits
bec157a458 .. 5d598b4d94

Author SHA1 Message Date
JackM323 5d598b4d94 Urdf update with collision 2026-07-30 17:23:50 +02:00
JackM323 c5ca79a354 Machine Learning Trainer 
Training environment to make a walk model for the hexapod
generated code that will be checked
 2026-07-30 17:23:36 +02:00
17 changed files with 1131 additions and 89 deletions
Expand all files Collapse all files
.gitignore
+2
View File
@@ -3,6 +3,8 @@ node_modules/
.venv/
.vs/
__pycache__/
ml/checkpoints/
ml/tensorboard/
# Ignore environment files with private passwords/keys
.env
.vscode/extensions.json
Vendored
+3
View File
@@ -0,0 +1,3 @@
{
"recommendations": []
}
GlobalVariables.py
+2 -44
View File
@@ -3,8 +3,8 @@ from ArduinoCommunication import ArduinoCommunication
from EspCommunication import ESP32Communication
from simulation import Simulation
import DataTypes as dt
import kinematics as kin
import config as cfg
from robot_init import init_deg, init90_deg, init_pos, center_points
# Globale Variablen
robotCommunication = None
@@ -36,46 +36,4 @@ leg_state = np.array(["step", "drag", "drag", "drag", "drag", "drag"])
control_pause: bool = False
# Initilize mit Start Position
#init_deg: dt.DegArray = dt.DegArray(
# [
# [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())
## Shared init pose definitions are imported from robot_init.py
JackBotUrdf.urdf
+150
View File
@@ -19,6 +19,12 @@
<color rgba="1 0 0 0.5"/>
</material>
</visual>
<collision>
<geometry>
<box size="0.160 0.120 0.090"/>
</geometry>
<origin xyz="0 0 0" rpy="0 0 0"/>
</collision>
</link>
<!-- LEG 1 -->
@@ -40,6 +46,12 @@
<color rgba="0 1 0 1"/>
</material>
</visual>
<collision>
<geometry>
<cylinder length="0.068" radius="0.005"/>
</geometry>
<origin xyz="0.028 0 0" rpy="0 1.5708 0"/>
</collision>
</link>
<joint name="leg1_coxa_joint" type="revolute">
<parent link="base_link"/>
@@ -68,6 +80,12 @@
<color rgba="0 0 1 1"/>
</material>
</visual>
<collision>
<geometry>
<cylinder length="0.0602" radius="0.005"/>
</geometry>
<origin xyz="0 0 -0.0365" rpy="0 0 0"/>
</collision>
</link>
<joint name="leg1_femur_joint" type="revolute">
<parent link="leg1_coxa"/>
@@ -95,6 +113,12 @@
<color rgba="0 1 1 1"/>
</material>
</visual>
<collision>
<geometry>
<cylinder length="0.07055" radius="0.005"/>
</geometry>
<origin xyz="0.035 0 0" rpy="0 1.5708 0"/>
</collision>
</link>
<joint name="leg1_tibia_joint" type="revolute">
<parent link="leg1_femur"/>
@@ -122,6 +146,12 @@
<color rgba="1 0 1 1"/>
</material>
</visual>
<collision>
<geometry>
<sphere radius="0.01"/>
</geometry>
<origin xyz="0 0 0" rpy="0 0 0"/>
</collision>
</link>
<joint name="leg1_tip_joint" type="fixed">
<parent link="leg1_tibia"/>
@@ -148,6 +178,12 @@
<color rgba="0 1 0 1"/>
</material>
</visual>
<collision>
<geometry>
<cylinder length="0.068" radius="0.005"/>
</geometry>
<origin xyz="0.028 0 0" rpy="0 1.5708 0"/>
</collision>
</link>
<joint name="leg2_coxa_joint" type="revolute">
<parent link="base_link"/>
@@ -176,6 +212,12 @@
<color rgba="0 0 1 1"/>
</material>
</visual>
<collision>
<geometry>
<cylinder length="0.0602" radius="0.005"/>
</geometry>
<origin xyz="0 0 -0.0365" rpy="0 0 0"/>
</collision>
</link>
<joint name="leg2_femur_joint" type="revolute">
<parent link="leg2_coxa"/>
@@ -203,6 +245,12 @@
<color rgba="0 1 1 1"/>
</material>
</visual>
<collision>
<geometry>
<cylinder length="0.07055" radius="0.005"/>
</geometry>
<origin xyz="0.035 0 0" rpy="0 1.5708 0"/>
</collision>
</link>
<joint name="leg2_tibia_joint" type="revolute">
<parent link="leg2_femur"/>
@@ -230,6 +278,12 @@
<color rgba="1 0 1 1"/>
</material>
</visual>
<collision>
<geometry>
<sphere radius="0.01"/>
</geometry>
<origin xyz="0 0 0" rpy="0 0 0"/>
</collision>
</link>
<joint name="leg2_tip_joint" type="fixed">
<parent link="leg2_tibia"/>
@@ -256,6 +310,12 @@
<color rgba="0 1 0 1"/>
</material>
</visual>
<collision>
<geometry>
<cylinder length="0.068" radius="0.005"/>
</geometry>
<origin xyz="0.028 0 0" rpy="0 1.5708 0"/>
</collision>
</link>
<joint name="leg3_coxa_joint" type="revolute">
<parent link="base_link"/>
@@ -284,6 +344,12 @@
<color rgba="0 0 1 1"/>
</material>
</visual>
<collision>
<geometry>
<cylinder length="0.0602" radius="0.005"/>
</geometry>
<origin xyz="0 0 -0.0365" rpy="0 0 0"/>
</collision>
</link>
<joint name="leg3_femur_joint" type="revolute">
<parent link="leg3_coxa"/>
@@ -311,6 +377,12 @@
<color rgba="0 1 1 1"/>
</material>
</visual>
<collision>
<geometry>
<cylinder length="0.07055" radius="0.005"/>
</geometry>
<origin xyz="0.035 0 0" rpy="0 1.5708 0"/>
</collision>
</link>
<joint name="leg3_tibia_joint" type="revolute">
<parent link="leg3_femur"/>
@@ -338,6 +410,12 @@
<color rgba="1 0 1 1"/>
</material>
</visual>
<collision>
<geometry>
<sphere radius="0.01"/>
</geometry>
<origin xyz="0 0 0" rpy="0 0 0"/>
</collision>
</link>
<joint name="leg3_tip_joint" type="fixed">
<parent link="leg3_tibia"/>
@@ -364,6 +442,12 @@
<color rgba="0 1 0 1"/>
</material>
</visual>
<collision>
<geometry>
<cylinder length="0.068" radius="0.005"/>
</geometry>
<origin xyz="0.028 0 0" rpy="0 1.5708 0"/>
</collision>
</link>
<joint name="leg4_coxa_joint" type="revolute">
<parent link="base_link"/>
@@ -392,6 +476,12 @@
<color rgba="0 0 1 1"/>
</material>
</visual>
<collision>
<geometry>
<cylinder length="0.0602" radius="0.005"/>
</geometry>
<origin xyz="0 0 -0.0365" rpy="0 0 0"/>
</collision>
</link>
<joint name="leg4_femur_joint" type="revolute">
<parent link="leg4_coxa"/>
@@ -419,6 +509,12 @@
<color rgba="0 1 1 1"/>
</material>
</visual>
<collision>
<geometry>
<cylinder length="0.07055" radius="0.005"/>
</geometry>
<origin xyz="0.035 0 0" rpy="0 1.5708 0"/>
</collision>
</link>
<joint name="leg4_tibia_joint" type="revolute">
<parent link="leg4_femur"/>
@@ -446,6 +542,12 @@
<color rgba="1 0 1 1"/>
</material>
</visual>
<collision>
<geometry>
<sphere radius="0.01"/>
</geometry>
<origin xyz="0 0 0" rpy="0 0 0"/>
</collision>
</link>
<joint name="leg4_tip_joint" type="fixed">
<parent link="leg4_tibia"/>
@@ -472,6 +574,12 @@
<color rgba="0 1 0 1"/>
</material>
</visual>
<collision>
<geometry>
<cylinder length="0.068" radius="0.005"/>
</geometry>
<origin xyz="0.028 0 0" rpy="0 1.5708 0"/>
</collision>
</link>
<joint name="leg5_coxa_joint" type="revolute">
<parent link="base_link"/>
@@ -500,6 +608,12 @@
<color rgba="0 0 1 1"/>
</material>
</visual>
<collision>
<geometry>
<cylinder length="0.0602" radius="0.005"/>
</geometry>
<origin xyz="0 0 -0.0365" rpy="0 0 0"/>
</collision>
</link>
<joint name="leg5_femur_joint" type="revolute">
<parent link="leg5_coxa"/>
@@ -527,6 +641,12 @@
<color rgba="0 1 1 1"/>
</material>
</visual>
<collision>
<geometry>
<cylinder length="0.07055" radius="0.005"/>
</geometry>
<origin xyz="0.035 0 0" rpy="0 1.5708 0"/>
</collision>
</link>
<joint name="leg5_tibia_joint" type="revolute">
<parent link="leg5_femur"/>
@@ -554,6 +674,12 @@
<color rgba="1 0 1 1"/>
</material>
</visual>
<collision>
<geometry>
<sphere radius="0.01"/>
</geometry>
<origin xyz="0 0 0" rpy="0 0 0"/>
</collision>
</link>
<joint name="leg5_tip_joint" type="fixed">
<parent link="leg5_tibia"/>
@@ -580,6 +706,12 @@
<color rgba="0 1 0 1"/>
</material>
</visual>
<collision>
<geometry>
<cylinder length="0.068" radius="0.005"/>
</geometry>
<origin xyz="0.028 0 0" rpy="0 1.5708 0"/>
</collision>
</link>
<joint name="leg6_coxa_joint" type="revolute">
<parent link="base_link"/>
@@ -608,6 +740,12 @@
<color rgba="0 0 1 1"/>
</material>
</visual>
<collision>
<geometry>
<cylinder length="0.0602" radius="0.005"/>
</geometry>
<origin xyz="0 0 -0.0365" rpy="0 0 0"/>
</collision>
</link>
<joint name="leg6_femur_joint" type="revolute">
<parent link="leg6_coxa"/>
@@ -635,6 +773,12 @@
<color rgba="0 1 1 1"/>
</material>
</visual>
<collision>
<geometry>
<cylinder length="0.07055" radius="0.005"/>
</geometry>
<origin xyz="0.035 0 0" rpy="0 1.5708 0"/>
</collision>
</link>
<joint name="leg6_tibia_joint" type="revolute">
<parent link="leg6_femur"/>
@@ -662,6 +806,12 @@
<color rgba="1 0 1 1"/>
</material>
</visual>
<collision>
<geometry>
<sphere radius="0.01"/>
</geometry>
<origin xyz="0 0 0" rpy="0 0 0"/>
</collision>
</link>
<joint name="leg6_tip_joint" type="fixed">
<parent link="leg6_tibia"/>
README.md
+19 -7
View File
@@ -6,8 +6,7 @@ It uses IKPy for inverse kinematics, PyBullet for optional simulation, and can s
## Requirements
- Python 3.12 recommended on Windows
- Python 3.11 / 3.12 is safest for `pygame` compatibility
- Python 3.12 is safest for `pygame` compatibility
- Required Python packages:
- `numpy`
- `pygame`
@@ -18,22 +17,22 @@ It uses IKPy for inverse kinematics, PyBullet for optional simulation, and can s
## Setup
### Linux / macOS (Bash)
### Linux (Bash)
```bash
python3 -m venv .venv
python3.12 -m venv .venv
source .venv/bin/activate
python -m pip install --upgrade pip setuptools wheel
pip install numpy pygame ikpy pybullet pyserial matplotlib
pip install -r requirements.txt
```
### Windows (PowerShell)
```powershell
python -m venv .venv
python3.12 -m venv .venv
.\.venv\Scripts\Activate.ps1
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:
@@ -51,6 +50,19 @@ From the repository root:
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
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)