Compare commits
2 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
 JackM323
|
5d598b4d94 | ||
|
JackM323
|
c5ca79a354 |
@@ -3,6 +3,8 @@ node_modules/
|
||||
.venv/
|
||||
.vs/
|
||||
__pycache__/
|
||||
ml/checkpoints/
|
||||
ml/tensorboard/
|
||||
|
||||
# Ignore environment files with private passwords/keys
|
||||
.env
|
||||
|
||||
Vendored
+3
@@ -0,0 +1,3 @@
|
||||
{
|
||||
"recommendations": []
|
||||
}
|
||||
+39
-81
@@ -1,81 +1,39 @@
|
||||
import numpy as np
|
||||
from ArduinoCommunication import ArduinoCommunication
|
||||
from EspCommunication import ESP32Communication
|
||||
from simulation import Simulation
|
||||
import DataTypes as dt
|
||||
import kinematics as kin
|
||||
import config as cfg
|
||||
|
||||
# Globale Variablen
|
||||
robotCommunication = None
|
||||
shared_sim = None
|
||||
emote = None
|
||||
display_text = ""
|
||||
|
||||
if cfg.sim:
|
||||
shared_sim = Simulation()
|
||||
else:
|
||||
if cfg.arduinoConnection:
|
||||
robotCommunication = ArduinoCommunication()
|
||||
else:
|
||||
robotCommunication = ESP32Communication()
|
||||
|
||||
vector_dirmov = [0, 0, 0] # Direction Movement [vx, vy, omega]
|
||||
current_rad: dt.RadArray # Current Rad
|
||||
current_pos: dt.PosArray # Current Position
|
||||
|
||||
# current_deg: dt.DegArray # Current Degrees
|
||||
# legarray: dt.DegArray # Working Leg Array
|
||||
# target_pos: dt.PosArray # Target Position
|
||||
|
||||
robot_state = "idle" # Current State
|
||||
# Init for 6 Legs
|
||||
#leg_state = np.array(["step", "drag", "step", "drag", "step", "drag"])
|
||||
# Init for 4 Legs
|
||||
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())
|
||||
import numpy as np
|
||||
from ArduinoCommunication import ArduinoCommunication
|
||||
from EspCommunication import ESP32Communication
|
||||
from simulation import Simulation
|
||||
import DataTypes as dt
|
||||
import config as cfg
|
||||
from robot_init import init_deg, init90_deg, init_pos, center_points
|
||||
|
||||
# Globale Variablen
|
||||
robotCommunication = None
|
||||
shared_sim = None
|
||||
emote = None
|
||||
display_text = ""
|
||||
|
||||
if cfg.sim:
|
||||
shared_sim = Simulation()
|
||||
else:
|
||||
if cfg.arduinoConnection:
|
||||
robotCommunication = ArduinoCommunication()
|
||||
else:
|
||||
robotCommunication = ESP32Communication()
|
||||
|
||||
vector_dirmov = [0, 0, 0] # Direction Movement [vx, vy, omega]
|
||||
current_rad: dt.RadArray # Current Rad
|
||||
current_pos: dt.PosArray # Current Position
|
||||
|
||||
# current_deg: dt.DegArray # Current Degrees
|
||||
# legarray: dt.DegArray # Working Leg Array
|
||||
# target_pos: dt.PosArray # Target Position
|
||||
|
||||
robot_state = "idle" # Current State
|
||||
# Init for 6 Legs
|
||||
#leg_state = np.array(["step", "drag", "step", "drag", "step", "drag"])
|
||||
# Init for 4 Legs
|
||||
leg_state = np.array(["step", "drag", "drag", "drag", "drag", "drag"])
|
||||
control_pause: bool = False
|
||||
|
||||
# Initilize mit Start Position
|
||||
## Shared init pose definitions are imported from robot_init.py
|
||||
|
||||
+151
-1
@@ -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"/>
|
||||
@@ -669,4 +819,4 @@
|
||||
<origin xyz="0.07055 0 0" rpy="0 0 0"/>
|
||||
</joint>
|
||||
|
||||
</robot>
|
||||
</robot>
|
||||
@@ -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:
|
||||
|
||||
@@ -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
|
||||
@@ -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.
|
||||
@@ -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"]
|
||||
@@ -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
|
||||
@@ -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,
|
||||
)
|
||||
+49
@@ -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
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
+113
@@ -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,
|
||||
)
|
||||
@@ -0,0 +1,14 @@
|
||||
numpy
|
||||
pygame
|
||||
ikpy
|
||||
pybullet
|
||||
pyserial
|
||||
matplotlib
|
||||
|
||||
# Deep learning / RL
|
||||
torch
|
||||
stable-baselines3
|
||||
gym
|
||||
gymnasium[box2d]
|
||||
shimmy
|
||||
tensorboard
|
||||
@@ -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)
|
||||
Reference in New Issue
Block a user