test_env.py

import argparse
import logging
from pathlib import Path

from algorithm.config_helper import set_logger
from algorithm.sac_main import Main as SACMain
from algorithm.utils.elapse_timer import UnifiedElapsedTimer


class Main(SACMain):
    def __init__(self, root_dir, config_dir, args):
        """
        config_path: the directory of config file
        args: command arguments generated by argparse
        """
        self._logger = logging.getLogger('test_env')

        self._profiler = UnifiedElapsedTimer(self._logger)

        config_abs_dir = self._init_config(root_dir, config_dir, args)

        self.train_mode = False

        self._init_env()
        self._run()

    def _log_episode_info(self, iteration, iter_time):
        for n, mgr in self.ma_manager:
            if len(mgr.agents) == 0:
                continue
            rewards = [a.reward for a in mgr.agents.values()]
            rewards = ", ".join([f"{i:6.1f}" for i in rewards])
            max_step = max([a.steps for a in mgr.agents.values()])
            self._logger.info(f'{n} {iteration}, T {iter_time:.2f}s, S {max_step}, R {rewards}')


if __name__ == '__main__':
    set_logger(debug=True)

    parser = argparse.ArgumentParser()
    parser.add_argument('env')
    parser.add_argument('--config', '-c', help='config file')
    parser.add_argument('--run', action='store_true', help='inference mode')

    parser.add_argument('--render', action='store_true', help='render')
    parser.add_argument('--env_args', default=[], nargs='+', help='additional arguments for environments')
    parser.add_argument('--envs', type=int, help='number of env copies')
    parser.add_argument('--max_iter', type=int, help='max iteration')

    parser.add_argument('--port', '-p', type=int, default=5005, help='UNITY: communication port')
    parser.add_argument('--editor', action='store_true', help='UNITY: running in Unity Editor')
    parser.add_argument('--timescale', type=float, default=None, help='UNITY: timescale')

    args = parser.parse_args()
    args.run_a = []
    args.logger_in_file = False
    args.name = ''
    args.disable_sample = True
    args.use_env_nn = False
    args.device = None
    args.ckpt = None
    args.nn = None

    root_dir = Path(__file__).resolve().parent
    Main(root_dir, f'envs/{args.env}', args)