if training_config['prioritized_replay']: buf = PrioritizedReplayBuffer(training_config['buffer_size'], alpha=training_config['alpha'], beta=training_config['beta']) else: buf = ReplayBuffer(training_config['buffer_size']) policy.set_eps(1) train_collector = Collector(policy, train_envs, buf) train_collector.collect(n_step=1) train_fn = lambda e: [ policy.set_eps( max( 0.05, 1 - e / training_config['epoch'] / training_config[ 'exploration_ratio'])), torch.save(policy.state_dict(), os.path.join(save_path, 'policy_%d.pth' % (e))) ] result = offpolicy_trainer(policy, train_collector, training_config['epoch'], training_config['step_per_epoch'], training_config['collect_per_step'], training_config['batch_size'], update_per_step=training_config['update_per_step'], train_fn=train_fn, writer=writer) env.close()
def test_dqn(args=get_args()): env = LimitWrapper(StateBonus(ImgObsWrapper(gym.make(args.task)))) # print(env.observation_space.spaces['image']) # args.state_shape = env.observation_space.spaces['image'].shape args.state_shape = env.observation_space.shape args.action_shape = env.env.action_space.shape or env.env.action_space.n # train_envs = gym.make(args.task) train_envs = SubprocVectorEnv([ lambda: LimitWrapper(StateBonus(ImgObsWrapper(gym.make(args.task)))) for _ in range(args.training_num) ]) # test_envs = gym.make(args.task) test_envs = SubprocVectorEnv([ lambda: LimitWrapper(StateBonus(ImgObsWrapper(gym.make(args.task)))) for _ in range(args.test_num) ]) # seed np.random.seed(args.seed) torch.manual_seed(args.seed) train_envs.seed(args.seed) test_envs.seed(args.seed) # model net = DQN(args.state_shape[0], args.state_shape[1], args.action_shape, args.device) net = net.to(args.device) optim = torch.optim.Adam(net.parameters(), lr=args.lr) policy = DQNPolicy(net, optim, args.gamma, args.n_step, use_target_network=args.target_update_freq > 0, target_update_freq=args.target_update_freq) # collector train_collector = Collector(policy, train_envs, ReplayBuffer(args.buffer_size)) test_collector = Collector(policy, test_envs) # policy.set_eps(1) train_collector.collect(n_step=args.batch_size * 4) print(len(train_collector.buffer)) # log writer = SummaryWriter(args.logdir + '/' + 'dqn') def stop_fn(x): # if env.env.spec.reward_threshold: # return x >= env.spec.reward_threshold # else: # return False return False def train_fn(x): policy.set_eps(args.eps_train) def test_fn(x): policy.set_eps(args.eps_test) # trainer result = offpolicy_trainer(policy, train_collector, test_collector, args.epoch, args.step_per_epoch, args.collect_per_step, args.test_num, args.batch_size, train_fn=train_fn, test_fn=test_fn, stop_fn=stop_fn, writer=writer, task=args.task) train_collector.close() test_collector.close() if __name__ == '__main__': pprint.pprint(result) # Let's watch its performance! env = LimitWrapper(StateBonus(ImgObsWrapper(gym.make(args.task)))) collector = Collector(policy, env) result = collector.collect(n_episode=1, render=args.render) print(f'Final reward: {result["rew"]}, length: {result["len"]}') collector.close() torch.save(policy.state_dict(), 'dqn.pth')