train.py

import os
import torch

from stable_baselines3 import PPO
from utils import parser as parse

from env.werdna_balance import WerdnaBalance
from env.werdna_stand import WerdnaStand

logs_name = ""

def main():

    config = parse.parser("config/config_stand.yaml")

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

    z_bias = config['z']
    dz_bias = config['dz']
    x_bias = config['x']
    v_bias = config['v']
    ent_coef = config['ec']

    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)}") # Get seed value from config, default to None if not specified

    if env_name == "werdna_balance":
        env = WerdnaBalance(
            model=robot_model, 
            render_mode=connect_type,
            z_bias = z_bias,
            x_bias = x_bias,
            v_bias = v_bias
        )
        logs_name = "logs/werdna_balance_tensorboard"
    elif env_name == "werdna_stand":
        env = WerdnaStand(
            model=robot_model,
            render_mode=connect_type,
            z_bias = z_bias,
            z_rate_bias=dz_bias,
            x_bias = x_bias,
            v_bias = v_bias
        )
        logs_name = "logs/werdna_stand_tensorboard"
    else:
        raise ValueError(f"Unknown Environment: {env_name}")

    model = PPO(
        "MlpPolicy",
        env=env,
        verbose=1,
        device=device,
        ent_coef= ent_coef,
        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()