import numpy as np
import os
from datetime import datetime
# from .vs_task_registry import VS_TaskRegistry
from roboimi.envs.vs_rl_task.task.vs_task_registry import VS_TaskRegistry
import wandb

ENVS_DIR = ''
task_registry = VS_TaskRegistry()

def train():
    wandb.init(project='vs_robot', name='vs_robot', entity="zfu")
    wandb.save(ENVS_DIR + "/h1/h1_config.py", policy="now")
    wandb.save(ENVS_DIR + "/h1/h1.py", policy="now")

    env, env_cfg = task_registry.make_vs_env()
    ppo_runner, train_cfg = task_registry.make_vs_alg_runner(env=env, name='vs_robot')
    ppo_runner.learn(num_learning_iterations=train_cfg.runner.max_iterations, init_at_random_ep_len=True)

if __name__ == '__main__':
    train()