-
Notifications
You must be signed in to change notification settings - Fork 0
/
train_data_collection.py
58 lines (47 loc) · 1.77 KB
/
train_data_collection.py
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
54
55
56
57
import os
import argparse
import torch
import numpy as np
from learner import Learner
from torchkit.pytorch_utils import set_gpu_mode
from data_collection_config import args_cheetah_vel, args_cheetah_dir, args_ant_dir
def main():
parser = argparse.ArgumentParser()
# parser.add_argument('--env-type', default='gridworld')
# parser.add_argument('--env-type', default='point_robot_sparse')
# parser.add_argument('--env-type', default='cheetah_vel')
parser.add_argument('--env-type', default='ant_semicircle_sparse')
args, rest_args = parser.parse_known_args()
env = args.env_type
# --- GridWorld ---
if env == 'gridworld':
args = args_gridworld.get_args(rest_args)
elif env == 'gridworld_block':
args = args_gridworld_block.get_args(rest_args)
# --- PointRobot ---
elif env == 'point_robot':
args = args_point_robot.get_args(rest_args)
elif env == 'point_robot_v1':
args = args_point_robot_v1.get_args(rest_args)
# --- Mujoco ---
elif env == 'cheetah_vel':
args = args_cheetah_vel.get_args(rest_args)
elif env == 'cheetah_dir':
args = args_cheetah_dir.get_args(rest_args)
elif env == 'ant_semicircle':
args = args_ant_semicircle.get_args(rest_args)
elif env == 'ant_dir':
args = args_ant_dir.get_args(rest_args)
elif env == 'hopper_param':
args = args_hopper_param.get_args(rest_args)
elif env == 'walker_param':
args = args_walker_param.get_args(rest_args)
else:
raise NotImplementedError
set_gpu_mode(torch.cuda.is_available())
if hasattr(args, 'save_buffer') and args.save_buffer:
os.makedirs(args.main_save_dir, exist_ok=True)
learner = Learner(args)
learner.train()
if __name__ == '__main__':
main()