714e0afa64f8b6e8f81ea9e0d3ecf0fea7154e96,slm_lab/experiment/analysis.py,,gen_return,#Any#Any#,24

Before Change

    """Generate return for an agent and an env in eval mode"""
    state = env.reset()
    done = False
    total_reward = 0
    while not done:
        action = agent.act(state)
        state, reward, done, info = env.step(action)
        total_reward += reward
    return total_reward


def gen_avg_return(agent, env, num_eval=NUM_EVAL):
    """Generate average return for agent and an env"""

After Change

def gen_return(agent, env):
    """Generate return for an agent and an env in eval mode. eval_env should be a vec env with NUM_EVAL instances"""
    // stats variables
    epi_start = True
    ckpt_total_reward = np.nan
    total_reward = 0
    vec_dones = False
    // swap ref to allow inference based on body.env
    main_env = agent.body.env
    agent.body.env = env
    // start eval loop
    state = env.reset()
    done = False
    while not np.all(vec_dones):
        action = agent.act(state)
        state, reward, done, info = env.step(action)
        if hasattr(env.u_env, "raw_reward"):  // use raw_reward if reward is preprocessed
            reward = env.u_env.raw_reward
        ckpt_total_reward, total_reward, epi_start = util.update_total_reward(ckpt_total_reward, total_reward, epi_start, reward, done)
        vec_dones = np.logical_or(vec_dones, done)  // wait till every vec slot done turns True
    // restore swapped ref
    agent.body.env = main_env
    return np.mean(total_reward)


def gen_avg_return(agent, env):
    """Generate average return for agent and an env"""

Instances