48 lines
1.8 KiB
Python
48 lines
1.8 KiB
Python
"""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-workers", type=int, default=1, help="Number of training environment with one robot each")
|
|
parser.add_argument("--robot-spacing", type=float, default=1.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_workers=args.num_workers,
|
|
robot_spacing=args.robot_spacing,
|
|
start_pose=args.start_pose,
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|