File size: 1,597 Bytes
037569d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 | import os
import torch
from stable_baselines3 import PPO
from stable_baselines3.common.callbacks import CheckpointCallback # <--- NEW IMPORT
from kuka_env import KukaReachingEnv
from utils import MetricsCallback
if __name__ == "__main__":
os.makedirs("models", exist_ok=True)
os.makedirs("models/checkpoints", exist_ok=True) # <--- NEW FOLDER
os.makedirs("logs", exist_ok=True)
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Training on device: {device.upper()}")
print("Initializing Environment...")
env = KukaReachingEnv(render_mode=None)
print("Setting up PPO with Convolutional Neural Network (CNN)...")
model = PPO(
"CnnPolicy",
env,
verbose=1,
tensorboard_log="./logs/",
device=device,
batch_size=128,
n_steps=2048
)
# 1. Our custom metric tracker
metrics_callback = MetricsCallback()
# 2. NEW: The Checkpoint Saver!
# Saves a backup of the model every 50,000 steps
checkpoint_callback = CheckpointCallback(
save_freq=50_000,
save_path="./models/checkpoints/",
name_prefix="ppo_vision"
)
print("Starting Deep Vision Training (3,000,000 steps)...")
print("Backups will be saved automatically to models/checkpoints/")
# Pass BOTH callbacks as a list
model.learn(
total_timesteps=3_000_000,
callback=[metrics_callback, checkpoint_callback]
)
print("Training complete! Saving the final model...")
model.save("models/ppo_vision_kuka")
env.close() |