def run(args):
    '\n    testing option will:\n    - use argmax, ie disable stochastic draws\n    - not run optimizers\n    - not save model\n    '
    if args.wandb:
        if args.wandb_offline:
            os.environ['WANDB_MODE'] = 'dryrun'
        wandb.init(project='ecn', name=args.name, dir=f'{args.savedir}', group=args.wandb_group)
        wandb.config.update(args)
        wandb.config.update(FLAGS)
    flags_dict = {flag.name: flag.value for flag in FLAGS.flags_by_module_dict()['main.py']}
    args_dict = args.__dict__
    pprint(args_dict)
    pprint(flags_dict)
    os.makedirs(args.model_dir, exist_ok=True)
    os.makedirs(args.logdir, exist_ok=True)
    if (args.seed is not None):
        np.random.seed(args.seed)
        torch.manual_seed(args.seed)
        train_r = np.random.RandomState(args.seed)
    else:
        train_r = np.random
    test_r = np.random.RandomState(args.test_seed)
    test_batches = generate_test_batches(batch_size=args.batch_size, num_batches=5, random_state=test_r)
    test_hashes = hash_batches(test_batches)
    episode = 0
    start_time = time.time()
    agent_models = []
    agent_opts = []
    agent_name = ['A', 'B']
    for i in range(2):
        model = AgentModel(name=agent_name[i], term_entropy_reg=args.term_entropy_reg, utterance_entropy_reg=args.utterance_entropy_reg, proposal_entropy_reg=args.proposal_entropy_reg).to(FLAGS.device)
        agent_models.append(model)
        agent_opts.append(optim.Adam(params=agent_models[i].parameters()))
    if args.wandb:
        wandb.watch(agent_models)
    if (path.isfile(args.model_file) and (not args.no_load)):
        (episode, start_time) = load_model(model_file=args.model_file, agent_models=agent_models, agent_opts=agent_opts)
        print('loaded model')
    elif args.testing:
        print('')
        print('ERROR: must have loadable model to use --testing option')
        print('')
        return
    last_print = time.time()
    rewards_sum = torch.zeros(3, device=FLAGS.device)
    steps_sum = 0
    count_sum = 0
    f_log = open(args.log_file, 'w')
    all_args = {**args_dict, **flags_dict}
    f_log.write(('meta: %s\n' % json.dumps(all_args)))
    last_save = time.time()
    baseline = torch.zeros(3, device=FLAGS.device)
    term_matches_argmax_count = 0
    num_policy_runs = 0
    utt_matches_argmax_count = 0
    utt_stochastic_draws = 0
    prop_matches_argmax_count = 0
    prop_stochastic_draws = 0
    utt_mask_count = np.array([0, 0])
    prop_mask_count = np.array([0, 0])
    while (episode < args.episodes):
        render = ((episode % args.render_every_episode) == 0)
        split = (2 if FLAGS.randomize_first else 1)
        agent_losses = [0, 0]
        both_rewards = []
        for i in range(2):
            agent_opts[i].zero_grad()
        for initial_agent in range(split):
            batch = generate_training_batch(batch_size=(args.batch_size // split), test_hashes=test_hashes, random_state=train_r)
            (actions, rewards, steps, alive_masks, entropy_loss_by_agent, _term_matches_argmax_count, _num_policy_runs, _utt_matches_argmax_count, _utt_stochastic_draws, _prop_matches_argmax_count, _prop_stochastic_draws, _utt_mask_count, _prop_mask_count) = run_episode(batch=batch, agent_models=agent_models, batch_size=(args.batch_size // split), render=render, initial_agent=initial_agent, testing=args.testing)
            term_matches_argmax_count += _term_matches_argmax_count
            utt_matches_argmax_count += _utt_matches_argmax_count
            utt_stochastic_draws += _utt_stochastic_draws
            num_policy_runs += _num_policy_runs
            prop_matches_argmax_count += _prop_matches_argmax_count
            prop_stochastic_draws += _prop_stochastic_draws
            utt_mask_count += _utt_mask_count
            prop_mask_count += _prop_mask_count
            if (not args.testing):
                reward_loss_by_agent = [0, 0]
                baselined_rewards = (rewards - baseline)
                rewards_by_agent = []
                for i in range(2):
                    if FLAGS.prosocial:
                        rewards_by_agent.append(baselined_rewards[:, 2])
                    else:
                        rewards_by_agent.append(baselined_rewards[:, i])
                sieve_playback = SievePlayback(alive_masks)
                for (t, global_idxes) in sieve_playback:
                    agent = ((initial_agent + t) % 2)
                    if (len(actions[t]) > 0):
                        for action in actions[t]:
                            _rewards = rewards_by_agent[agent]
                            _reward = _rewards[global_idxes].float().contiguous().view(sieve_playback.batch_size, 1)
                            _reward_loss = (- (action * _reward))
                            _reward_loss = _reward_loss.sum()
                            reward_loss_by_agent[agent] += _reward_loss
                for i in range(2):
                    loss = (entropy_loss_by_agent[i] + reward_loss_by_agent[i])
                    loss.backward()
            rewards_sum += rewards.detach().sum(0)
            steps_sum += steps.sum()
            count_sum += (args.batch_size // split)
            both_rewards.append(rewards)
        for i in range(2):
            agent_opts[i].step()
        rewards = torch.cat(both_rewards).detach()
        baseline = ((0.7 * baseline) + (0.3 * rewards.mean(0).detach()))
        if render:
            '\n            run the test batches, print the results\n            '
            test_rewards_sum = np.zeros(3)
            test_count_sum = (len(test_batches) * args.batch_size)
            test_num_policy_runs = 0
            test_utt_mask_count = [0, 0]
            test_prop_mask_count = [0, 0]
            test_utt_mask_count = np.array([0, 0])
            test_prop_mask_count = np.array([0, 0])
            for test_batch in test_batches:
                (actions, test_rewards, steps, alive_masks, entropy_loss_by_agent, _term_matches_argmax_count, _test_num_policy_runs, _utt_matches_argmax_count, _utt_stochastic_draws, _prop_matches_argmax_count, _prop_stochastic_draws, _test_utt_mask_count, _test_prop_mask_count) = run_episode(batch=test_batch, agent_models=agent_models, batch_size=args.batch_size, render=True, testing=True)
                test_rewards_sum += test_rewards.sum(0).cpu().numpy()
                test_num_policy_runs += _test_num_policy_runs
                test_utt_mask_count += _test_utt_mask_count
                test_prop_mask_count += _test_prop_mask_count
            time_since_last = (time.time() - last_print)
            rewards_str = ('%.2f,%.2f,%.2f' % ((rewards_sum[0] / count_sum), (rewards_sum[1] / count_sum), (rewards_sum[2] / count_sum)))
            test_rewards_str = ('%.2f,%.2f,%.2f' % ((test_rewards_sum[0] / test_count_sum), (test_rewards_sum[1] / test_count_sum), (test_rewards_sum[2] / test_count_sum)))
            baseline_str = ('%.2f,%.2f,%.2f' % (baseline[0], baseline[1], baseline[2]))
            utt_mask_pct = (utt_mask_count / (3 * count_sum))
            test_utt_mask_pct = (test_utt_mask_count / (3 * test_count_sum))
            prop_mask_pct = (prop_mask_count / (3 * count_sum))
            test_prop_mask_pct = (test_prop_mask_count / (3 * test_count_sum))
            print('test  {}'.format(test_rewards_str))
            print('train {}'.format(rewards_str))
            print('base  {}'.format(baseline_str))
            print('ep {}, {} games/sec, {:2.2f} avg steps'.format(episode, int((count_sum / time_since_last)), (steps_sum.item() / count_sum)))
            print('argmaxp term={:4.4f} utt={:4.4f} prop={:4.4f}'.format((term_matches_argmax_count / num_policy_runs), safe_div(utt_matches_argmax_count, utt_stochastic_draws), (prop_matches_argmax_count / prop_stochastic_draws)))
            if FLAGS.force_masking_comm:
                print('utt mask % {:2.2f},{:2.2f} test % {:2.2f},{:2.2f}'.format(*utt_mask_pct, *test_utt_mask_pct))
                print('prop mask % {:2.2f},{:2.2f} test % {:2.2f},{:2.2f}'.format(*prop_mask_pct, *test_prop_mask_pct))
            episode_log = {'episode': episode, 'avg_reward_A': (rewards_sum[0] / count_sum).item(), 'avg_reward_B': (rewards_sum[1] / count_sum).item(), 'avg_reward_0': (rewards_sum[2] / count_sum).item(), 'test_reward_A': (test_rewards_sum[0] / test_count_sum).item(), 'test_reward_B': (test_rewards_sum[1] / test_count_sum).item(), 'test_reward': (test_rewards_sum[2] / test_count_sum).item(), 'avg_steps': torch.true_divide(steps_sum, count_sum).item(), 'games_sec': (count_sum / time_since_last), 'elapsed': (time.time() - start_time), 'argmaxp_term': (term_matches_argmax_count / num_policy_runs), 'argmaxp_utt': safe_div(utt_matches_argmax_count, utt_stochastic_draws), 'argmaxp_prop': (prop_matches_argmax_count / prop_stochastic_draws), 'utt_unmasked_A': utt_mask_pct[0], 'utt_unmasked_B': utt_mask_pct[1], 'prop_unmasked_A': prop_mask_pct[0], 'prop_unmasked_B': prop_mask_pct[1], 'test_utt_unmasked_A': test_utt_mask_pct[0], 'test_utt_unmasked_B': test_utt_mask_pct[1], 'test_prop_unmasked_A': test_prop_mask_pct[0], 'test_prop_unmasked_B': test_prop_mask_pct[1]}
            f_log.write((json.dumps(episode_log) + '\n'))
            f_log.flush()
            if args.wandb:
                wandb.log(episode_log)
            last_print = time.time()
            steps_sum = 0
            rewards_sum.fill_(0)
            term_matches_argmax_count = 0
            num_policy_runs = 0
            utt_matches_argmax_count = 0
            utt_stochastic_draws = 0
            prop_matches_argmax_count = 0
            prop_stochastic_draws = 0
            count_sum = 0
            utt_mask_count.fill(0)
            prop_mask_count.fill(0)
        if ((not args.testing) and (not args.no_save) and (episode > 0) and ((episode % args.save_every_episode) == 0)):
            save_model(model_file=args.model_file, agent_models=agent_models, agent_opts=agent_opts, start_time=start_time, episode=episode)
            print('saved model')
        episode += 1
    if ((not args.no_save) and (not args.testing)):
        save_model(model_file=args.model_file, agent_models=agent_models, agent_opts=agent_opts, start_time=start_time, episode=episode)
        print('saved model')
    f_log.close()
testing option will:
- use argmax, ie disable stochastic draws
- not run optimizers
- not save model
def run(args):
    '\n    testing option will:\n    - use argmax, ie disable stochastic draws\n    - not run optimizers\n    - not save model\n    '
    if args.wandb:
        if args.wandb_offline:
            os.environ['WANDB_MODE'] = 'dryrun'
        wandb.init(project='ecn', name=args.name, dir=f'{args.savedir}', group=args.wandb_group)
        wandb.config.update(args)
        wandb.config.update(FLAGS)
    flags_dict = {flag.name: flag.value for flag in FLAGS.flags_by_module_dict()['main.py']}
    args_dict = args.__dict__
    pprint(args_dict)
    pprint(flags_dict)
    os.makedirs(args.model_dir, exist_ok=True)
    os.makedirs(args.logdir, exist_ok=True)
    if (args.seed is not None):
        np.random.seed(args.seed)
        torch.manual_seed(args.seed)
        train_r = np.random.RandomState(args.seed)
    else:
        train_r = np.random
    test_r = np.random.RandomState(args.test_seed)
    test_batches = generate_test_batches(batch_size=args.batch_size, num_batches=5, random_state=test_r)
    test_hashes = hash_batches(test_batches)
    episode = 0
    start_time = time.time()
    agent_models = []
    agent_opts = []
    agent_name = ['A', 'B']
    for i in range(2):
        model = AgentModel(name=agent_name[i], term_entropy_reg=args.term_entropy_reg, utterance_entropy_reg=args.utterance_entropy_reg, proposal_entropy_reg=args.proposal_entropy_reg).to(FLAGS.device)
        agent_models.append(model)
        agent_opts.append(optim.Adam(params=agent_models[i].parameters()))
    if args.wandb:
        wandb.watch(agent_models)
    if (path.isfile(args.model_file) and (not args.no_load)):
        (episode, start_time) = load_model(model_file=args.model_file, agent_models=agent_models, agent_opts=agent_opts)
        print('loaded model')
    elif args.testing:
        print()
        print('ERROR: must have loadable model to use --testing option')
        print()
        return
    last_print = time.time()
    rewards_sum = torch.zeros(3, device=FLAGS.device)
    steps_sum = 0
    count_sum = 0
    f_log = open(args.log_file, 'w')
    all_args = {**args_dict, **flags_dict}
    f_log.write(('meta: %s\n' % json.dumps(all_args)))
    last_save = time.time()
    baseline = torch.zeros(3, device=FLAGS.device)
    term_matches_argmax_count = 0
    num_policy_runs = 0
    utt_matches_argmax_count = 0
    utt_stochastic_draws = 0
    prop_matches_argmax_count = 0
    prop_stochastic_draws = 0
    utt_mask_count = np.array([0, 0])
    prop_mask_count = np.array([0, 0])
    while (episode < args.episodes):
        render = ((episode % args.render_every_episode) == 0)
        split = (2 if FLAGS.randomize_first else 1)
        agent_losses = [0, 0]
        both_rewards = []
        for i in range(2):
            agent_opts[i].zero_grad()
        for initial_agent in range(split):
            batch = generate_training_batch(batch_size=(args.batch_size // split), test_hashes=test_hashes, random_state=train_r)
            (actions, rewards, steps, alive_masks, entropy_loss_by_agent, _term_matches_argmax_count, _num_policy_runs, _utt_matches_argmax_count, _utt_stochastic_draws, _prop_matches_argmax_count, _prop_stochastic_draws, _utt_mask_count, _prop_mask_count) = run_episode(batch=batch, agent_models=agent_models, batch_size=(args.batch_size // split), render=render, initial_agent=initial_agent, testing=args.testing)
            term_matches_argmax_count += _term_matches_argmax_count
            utt_matches_argmax_count += _utt_matches_argmax_count
            utt_stochastic_draws += _utt_stochastic_draws
            num_policy_runs += _num_policy_runs
            prop_matches_argmax_count += _prop_matches_argmax_count
            prop_stochastic_draws += _prop_stochastic_draws
            utt_mask_count += _utt_mask_count
            prop_mask_count += _prop_mask_count
            if (not args.testing):
                reward_loss_by_agent = [0, 0]
                baselined_rewards = (rewards - baseline)
                rewards_by_agent = []
                for i in range(2):
                    if FLAGS.prosocial:
                        rewards_by_agent.append(baselined_rewards[:, 2])
                    else:
                        rewards_by_agent.append(baselined_rewards[:, i])
                sieve_playback = SievePlayback(alive_masks)
                for (t, global_idxes) in sieve_playback:
                    agent = ((initial_agent + t) % 2)
                    if (len(actions[t]) > 0):
                        for action in actions[t]:
                            _rewards = rewards_by_agent[agent]
                            _reward = _rewards[global_idxes].float().contiguous().view(sieve_playback.batch_size, 1)
                            _reward_loss = (- (action * _reward))
                            _reward_loss = _reward_loss.sum()
                            reward_loss_by_agent[agent] += _reward_loss
                for i in range(2):
                    loss = (entropy_loss_by_agent[i] + reward_loss_by_agent[i])
                    loss.backward()
            rewards_sum += rewards.detach().sum(0)
            steps_sum += steps.sum()
            count_sum += (args.batch_size // split)
            both_rewards.append(rewards)
        for i in range(2):
            agent_opts[i].step()
        rewards = torch.cat(both_rewards).detach()
        baseline = ((0.7 * baseline) + (0.3 * rewards.mean(0).detach()))
        if render:
            '\n            run the test batches, print the results\n            '
            test_rewards_sum = np.zeros(3)
            test_count_sum = (len(test_batches) * args.batch_size)
            test_num_policy_runs = 0
            test_utt_mask_count = [0, 0]
            test_prop_mask_count = [0, 0]
            test_utt_mask_count = np.array([0, 0])
            test_prop_mask_count = np.array([0, 0])
            for test_batch in test_batches:
                (actions, test_rewards, steps, alive_masks, entropy_loss_by_agent, _term_matches_argmax_count, _test_num_policy_runs, _utt_matches_argmax_count, _utt_stochastic_draws, _prop_matches_argmax_count, _prop_stochastic_draws, _test_utt_mask_count, _test_prop_mask_count) = run_episode(batch=test_batch, agent_models=agent_models, batch_size=args.batch_size, render=True, testing=True)
                test_rewards_sum += test_rewards.sum(0).cpu().numpy()
                test_num_policy_runs += _test_num_policy_runs
                test_utt_mask_count += _test_utt_mask_count
                test_prop_mask_count += _test_prop_mask_count
            time_since_last = (time.time() - last_print)
            rewards_str = ('%.2f,%.2f,%.2f' % ((rewards_sum[0] / count_sum), (rewards_sum[1] / count_sum), (rewards_sum[2] / count_sum)))
            test_rewards_str = ('%.2f,%.2f,%.2f' % ((test_rewards_sum[0] / test_count_sum), (test_rewards_sum[1] / test_count_sum), (test_rewards_sum[2] / test_count_sum)))
            baseline_str = ('%.2f,%.2f,%.2f' % (baseline[0], baseline[1], baseline[2]))
            utt_mask_pct = (utt_mask_count / (3 * count_sum))
            test_utt_mask_pct = (test_utt_mask_count / (3 * test_count_sum))
            prop_mask_pct = (prop_mask_count / (3 * count_sum))
            test_prop_mask_pct = (test_prop_mask_count / (3 * test_count_sum))
            print('test  {}'.format(test_rewards_str))
            print('train {}'.format(rewards_str))
            print('base  {}'.format(baseline_str))
            print('ep {}, {} games/sec, {:2.2f} avg steps'.format(episode, int((count_sum / time_since_last)), (steps_sum.item() / count_sum)))
            print('argmaxp term={:4.4f} utt={:4.4f} prop={:4.4f}'.format((term_matches_argmax_count / num_policy_runs), safe_div(utt_matches_argmax_count, utt_stochastic_draws), (prop_matches_argmax_count / prop_stochastic_draws)))
            if FLAGS.force_masking_comm:
                print('utt mask % {:2.2f},{:2.2f} test % {:2.2f},{:2.2f}'.format(*utt_mask_pct, *test_utt_mask_pct))
                print('prop mask % {:2.2f},{:2.2f} test % {:2.2f},{:2.2f}'.format(*prop_mask_pct, *test_prop_mask_pct))
            episode_log = {'episode': episode, 'avg_reward_A': (rewards_sum[0] / count_sum).item(), 'avg_reward_B': (rewards_sum[1] / count_sum).item(), 'avg_reward_0': (rewards_sum[2] / count_sum).item(), 'test_reward_A': (test_rewards_sum[0] / test_count_sum).item(), 'test_reward_B': (test_rewards_sum[1] / test_count_sum).item(), 'test_reward': (test_rewards_sum[2] / test_count_sum).item(), 'avg_steps': torch.true_divide(steps_sum, count_sum).item(), 'games_sec': (count_sum / time_since_last), 'elapsed': (time.time() - start_time), 'argmaxp_term': (term_matches_argmax_count / num_policy_runs), 'argmaxp_utt': safe_div(utt_matches_argmax_count, utt_stochastic_draws), 'argmaxp_prop': (prop_matches_argmax_count / prop_stochastic_draws), 'utt_unmasked_A': utt_mask_pct[0], 'utt_unmasked_B': utt_mask_pct[1], 'prop_unmasked_A': prop_mask_pct[0], 'prop_unmasked_B': prop_mask_pct[1], 'test_utt_unmasked_A': test_utt_mask_pct[0], 'test_utt_unmasked_B': test_utt_mask_pct[1], 'test_prop_unmasked_A': test_prop_mask_pct[0], 'test_prop_unmasked_B': test_prop_mask_pct[1]}
            f_log.write((json.dumps(episode_log) + '\n'))
            f_log.flush()
            if args.wandb:
                wandb.log(episode_log)
            last_print = time.time()
            steps_sum = 0
            rewards_sum.fill_(0)
            term_matches_argmax_count = 0
            num_policy_runs = 0
            utt_matches_argmax_count = 0
            utt_stochastic_draws = 0
            prop_matches_argmax_count = 0
            prop_stochastic_draws = 0
            count_sum = 0
            utt_mask_count.fill(0)
            prop_mask_count.fill(0)
        if ((not args.testing) and (not args.no_save) and (episode > 0) and ((episode % args.save_every_episode) == 0)):
            save_model(model_file=args.model_file, agent_models=agent_models, agent_opts=agent_opts, start_time=start_time, episode=episode)
            print('saved model')
        episode += 1
    if ((not args.no_save) and (not args.testing)):
        save_model(model_file=args.model_file, agent_models=agent_models, agent_opts=agent_opts, start_time=start_time, episode=episode)
        print('saved model')
    f_log.close()