File size: 1,039 Bytes
079c32c |
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 |
from typing import TYPE_CHECKING, Callable
import os
import pickle
import dataclasses
from ding.framework import task
if TYPE_CHECKING:
from ding.framework import Context
def final_ctx_saver(name: str) -> Callable:
def _save(ctx: "Context"):
if task.finish:
# make sure the items to be recorded are all kept in the context
with open(os.path.join(name, 'result.pkl'), 'wb') as f:
final_data = {
'total_step': ctx.total_step,
'train_iter': ctx.train_iter,
'last_eval_iter': ctx.last_eval_iter,
'eval_value': ctx.last_eval_value,
}
if ctx.has_attr('env_step'):
final_data['env_step'] = ctx.env_step
final_data['env_episode'] = ctx.env_episode
if ctx.has_attr('trained_env_step'):
final_data['trained_env_step'] = ctx.trained_env_step
pickle.dump(final_data, f)
return _save
|