ding.framework.middleware.functional.ctx_helper

ding.framework.middleware.functional.ctx_helper

Full Source Code

../ding/framework/middleware/functional/ctx_helper.py

1from typing import TYPE_CHECKING, Callable 2import os 3import pickle 4import dataclasses 5from ding.framework import task 6if TYPE_CHECKING: 7 from ding.framework import Context 8 9 10def final_ctx_saver(name: str) -> Callable: 11 12 def _save(ctx: "Context"): 13 if task.finish: 14 # make sure the items to be recorded are all kept in the context 15 with open(os.path.join(name, 'result.pkl'), 'wb') as f: 16 final_data = { 17 'total_step': ctx.total_step, 18 'train_iter': ctx.train_iter, 19 'last_eval_iter': ctx.last_eval_iter, 20 'eval_value': ctx.last_eval_value, 21 } 22 if ctx.has_attr('env_step'): 23 final_data['env_step'] = ctx.env_step 24 final_data['env_episode'] = ctx.env_episode 25 if ctx.has_attr('trained_env_step'): 26 final_data['trained_env_step'] = ctx.trained_env_step 27 pickle.dump(final_data, f) 28 29 return _save