JackM323
c5ca79a354
Training environment to make a walk model for the hexapod generated code that will be checked
114 lines
3.9 KiB
Python
114 lines
3.9 KiB
Python
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,
|
|
)
|