38 lines
933 B
Python
38 lines
933 B
Python
import hydra
|
|
import torch
|
|
from termcolor import colored
|
|
|
|
from lerobot.common.envs.factory import make_env
|
|
from lerobot.common.tdmpc import TDMPC
|
|
|
|
from ..common.utils import set_seed
|
|
|
|
|
|
@hydra.main(version_base=None, config_name="default", config_path="../configs")
|
|
def train(cfg: dict):
|
|
assert torch.cuda.is_available()
|
|
set_seed(cfg.seed)
|
|
print(colored("Work dir:", "yellow", attrs=["bold"]), cfg.log_dir)
|
|
|
|
env = make_env(cfg)
|
|
agent = TDMPC(cfg)
|
|
# ckpt_path = "/home/rcadene/code/fowm/logs/xarm_lift/all/default/2/models/offline.pt"
|
|
ckpt_path = "/home/rcadene/code/fowm/logs/xarm_lift/all/default/2/models/final.pt"
|
|
agent.load(ckpt_path)
|
|
|
|
# online training
|
|
|
|
eval_metrics = train_agent(
|
|
env,
|
|
agent,
|
|
num_episodes=10,
|
|
save_video=True,
|
|
video_dir=Path("tmp/2023_01_29_xarm_lift_final"),
|
|
)
|
|
|
|
print(eval_metrics)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
train()
|