|
from typing import Union, Optional, Tuple |
|
import os |
|
from functools import partial |
|
from copy import deepcopy |
|
|
|
import torch |
|
from tensorboardX import SummaryWriter |
|
from torch.utils.data import DataLoader |
|
|
|
from ding.envs import get_vec_env_setting, create_env_manager |
|
from ding.worker import BaseLearner, InteractionSerialEvaluator |
|
from ding.config import read_config, compile_config |
|
from ding.policy import create_policy |
|
from ding.utils import set_pkg_seed |
|
from ding.utils.data.dataset import load_bfs_datasets |
|
|
|
|
|
def serial_pipeline_pc( |
|
input_cfg: Union[str, Tuple[dict, dict]], |
|
seed: int = 0, |
|
model: Optional[torch.nn.Module] = None, |
|
max_iter=int(1e6), |
|
) -> Union['Policy', bool]: |
|
r""" |
|
Overview: |
|
Serial pipeline entry of procedure cloning using BFS as expert policy. |
|
Arguments: |
|
- input_cfg (:obj:`Union[str, Tuple[dict, dict]]`): Config in dict type. \ |
|
``str`` type means config file path. \ |
|
``Tuple[dict, dict]`` type means [user_config, create_cfg]. |
|
- seed (:obj:`int`): Random seed. |
|
- model (:obj:`Optional[torch.nn.Module]`): Instance of torch.nn.Module. |
|
- max_iter (:obj:`Optional[int]`): Max iteration for executing PC training. |
|
Returns: |
|
- policy (:obj:`Policy`): Converged policy. |
|
- convergence (:obj:`bool`): whether the training is converged |
|
""" |
|
if isinstance(input_cfg, str): |
|
cfg, create_cfg = read_config(input_cfg) |
|
else: |
|
cfg, create_cfg = deepcopy(input_cfg) |
|
cfg = compile_config(cfg, seed=seed, auto=True, create_cfg=create_cfg) |
|
|
|
|
|
env_fn, _, evaluator_env_cfg = get_vec_env_setting(cfg.env) |
|
evaluator_env = create_env_manager(cfg.env.manager, [partial(env_fn, cfg=c) for c in evaluator_env_cfg]) |
|
|
|
evaluator_env.seed(cfg.seed, dynamic_seed=False) |
|
set_pkg_seed(cfg.seed, use_cuda=cfg.policy.cuda) |
|
policy = create_policy(cfg.policy, model=model, enable_field=['learn', 'eval']) |
|
|
|
|
|
tb_logger = SummaryWriter(os.path.join('./{}/log/'.format(cfg.exp_name), 'serial')) |
|
train_data, test_data = load_bfs_datasets(train_seeds=cfg.train_seeds) |
|
dataloader = DataLoader(train_data, batch_size=cfg.policy.learn.batch_size, shuffle=True) |
|
test_dataloader = DataLoader(test_data, batch_size=cfg.policy.learn.batch_size, shuffle=True) |
|
learner = BaseLearner(cfg.policy.learn.learner, policy.learn_mode, tb_logger, exp_name=cfg.exp_name) |
|
evaluator = InteractionSerialEvaluator( |
|
cfg.policy.eval.evaluator, evaluator_env, policy.eval_mode, tb_logger, exp_name=cfg.exp_name |
|
) |
|
|
|
|
|
|
|
|
|
learner.call_hook('before_run') |
|
stop = False |
|
iter_cnt = 0 |
|
for epoch in range(cfg.policy.learn.train_epoch): |
|
|
|
criterion = torch.nn.CrossEntropyLoss() |
|
for i, train_data in enumerate(dataloader): |
|
learner.train(train_data) |
|
iter_cnt += 1 |
|
if iter_cnt >= max_iter: |
|
stop = True |
|
break |
|
if epoch % 69 == 0: |
|
policy._optimizer.param_groups[0]['lr'] /= 10 |
|
if stop: |
|
break |
|
losses = [] |
|
acces = [] |
|
|
|
for _, test_data in enumerate(test_dataloader): |
|
observations, bfs_input_maps, bfs_output_maps = test_data['obs'], test_data['bfs_in'].long(), \ |
|
test_data['bfs_out'].long() |
|
states = observations |
|
bfs_input_onehot = torch.nn.functional.one_hot(bfs_input_maps, 5).float() |
|
|
|
bfs_states = torch.cat([ |
|
states, |
|
bfs_input_onehot, |
|
], dim=-1).cuda() |
|
logits = policy._model(bfs_states)['logit'] |
|
logits = logits.flatten(0, -2) |
|
labels = bfs_output_maps.flatten(0, -1).cuda() |
|
|
|
loss = criterion(logits, labels).item() |
|
preds = torch.argmax(logits, dim=-1) |
|
acc = torch.sum((preds == labels)) / preds.shape[0] |
|
|
|
losses.append(loss) |
|
acces.append(acc) |
|
print('Test Finished! Loss: {} acc: {}'.format(sum(losses) / len(losses), sum(acces) / len(acces))) |
|
stop, reward = evaluator.eval(learner.save_checkpoint, learner.train_iter) |
|
learner.call_hook('after_run') |
|
print('final reward is: {}'.format(reward)) |
|
return policy, stop |
|
|