diff --git a/.gitignore b/.gitignore index 4d50f7e..3eabf78 100644 --- a/.gitignore +++ b/.gitignore @@ -3,6 +3,8 @@ node_modules/ .venv/ .vs/ __pycache__/ +ml/checkpoints/ +ml/tensorboard/ # Ignore environment files with private passwords/keys .env diff --git a/.vscode/extensions.json b/.vscode/extensions.json new file mode 100644 index 0000000..7e257db --- /dev/null +++ b/.vscode/extensions.json @@ -0,0 +1,3 @@ +{ + "recommendations": [] +} \ No newline at end of file diff --git a/GlobalVariables.py b/GlobalVariables.py index 553c3ef..d7a023e 100644 --- a/GlobalVariables.py +++ b/GlobalVariables.py @@ -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()) \ No newline at end of file +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 diff --git a/README.md b/README.md index 30618da..84e8849 100644 --- a/README.md +++ b/README.md @@ -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: diff --git a/RobotState/ml_walking.py b/RobotState/ml_walking.py new file mode 100644 index 0000000..723e0f5 --- /dev/null +++ b/RobotState/ml_walking.py @@ -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 diff --git a/ml/README.md b/ml/README.md new file mode 100644 index 0000000..7c2fc51 --- /dev/null +++ b/ml/README.md @@ -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. diff --git a/ml/__init__.py b/ml/__init__.py new file mode 100644 index 0000000..952554b --- /dev/null +++ b/ml/__init__.py @@ -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"] diff --git a/ml/env.py b/ml/env.py new file mode 100644 index 0000000..f4f3830 --- /dev/null +++ b/ml/env.py @@ -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 diff --git a/ml/evaluate.py b/ml/evaluate.py new file mode 100644 index 0000000..d67ebd3 --- /dev/null +++ b/ml/evaluate.py @@ -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, + ) diff --git a/ml/model.py b/ml/model.py new file mode 100644 index 0000000..cb3aabd --- /dev/null +++ b/ml/model.py @@ -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 diff --git a/ml/run_eval.py b/ml/run_eval.py new file mode 100644 index 0000000..20aafd9 --- /dev/null +++ b/ml/run_eval.py @@ -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() diff --git a/ml/run_train.py b/ml/run_train.py new file mode 100644 index 0000000..045569f --- /dev/null +++ b/ml/run_train.py @@ -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() diff --git a/ml/sim_manager.py b/ml/sim_manager.py new file mode 100644 index 0000000..fc20567 --- /dev/null +++ b/ml/sim_manager.py @@ -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() \ No newline at end of file diff --git a/ml/train.py b/ml/train.py new file mode 100644 index 0000000..64435b7 --- /dev/null +++ b/ml/train.py @@ -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, + ) diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..512662e --- /dev/null +++ b/requirements.txt @@ -0,0 +1,14 @@ +numpy +pygame +ikpy +pybullet +pyserial +matplotlib + +# Deep learning / RL +torch +stable-baselines3 +gym +gymnasium[box2d] +shimmy +tensorboard diff --git a/robot_init.py b/robot_init.py new file mode 100644 index 0000000..da79016 --- /dev/null +++ b/robot_init.py @@ -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)