train_ground_truth.py

import os
import torch

from stable_baselines3 import PPO
from utils import parser as parse

from env.werdna_ground_truth import WerdnaGT

logs_name = ""

def main():
    config = parse.parser("config/config_ground_truth.yaml")

    robot_model = config['robot_model']
    env_name = config['environment']
    connect_type = config['connect_type']
    filename = config['filename']
    timesteps = config['timesteps']
    device_type = config['device']
    ec = config['ec']

    # Parse the biases from the configuration
    biases = config['biases']

    # Create a tb_log_name by appending each bias value
    tb_log_name = f"{filename}_" + "_".join([f"{key}{value}" for key, value in biases.items()])

    print(f"Robot Model: {robot_model}")
    print(f"Environment: {env_name}")

    if device_type == "cpu":
        device = "cpu"
        print("Training on CPU")
    else:
        device = torch.device("cuda")
        print(f"Training on GPU: {torch.cuda.get_device_name(0)}")

    env = WerdnaGT(
        model=robot_model, 
        render_mode=connect_type,
        **biases  # Pass biases directly as keyword arguments
    )
    logs_name = "logs/werdna_gT_tensorboard"

    model = PPO(
        "MlpPolicy",
        env=env,
        verbose=1,
        device=device,
        ent_coef=ec,
        use_sde=False,
        tensorboard_log=logs_name
    )

    model.learn(total_timesteps=timesteps, tb_log_name=tb_log_name)

    # Create results directory if it doesn't exist
    results_dir = os.path.join("results", tb_log_name)
    os.makedirs(results_dir, exist_ok=True)

    # Save the model
    print(f"Saving model to {os.path.join(results_dir, filename)}")
    model.save(os.path.join(results_dir, filename))

    env.close()

if __name__ == "__main__":
    main()