<commit_before>from slm_lab.lib import logger, util, viz
from slm_lab.spec import random_baseline
import numpy as np
import os
import pandas as pd
import pydash as ps
import shutil
import torch


METRICS_COLS = [
    'final_return_ma',
    'strength', 'max_strength', 'final_strength',
    'sample_efficiency', 'training_efficiency',
    'stability', 'consistency',
]

logger = logger.get_logger(__name__)


# methods to generate returns (total rewards)

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'''
    epi_start = True
    ckpt_total_reward = np.nan
    total_reward = 0
    vec_dones = False
    state = env.reset()
    done = False
    while not np.all(vec_dones):
        action = agent.act(state)
        state, reward, done, info = env.step(action)
        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
    return np.mean(total_reward)


def gen_avg_return(agent, env):
    '''Generate average return for agent and an env'''
    with util.ctx_lab_mode('eval'):  # enter eval context
        agent.algorithm.update()  # set explore_var etc. to end_val under ctx
    with torch.no_grad():
        ret = gen_return(agent, env)
    # exit eval context, restore variables simply by updating
    agent.algorithm.update()
    return ret


# metrics calculation methods

def calc_strength(mean_returns, mean_rand_returns):
    '''
    Calculate strength for metric
    str &= \frac{1}{N} \sum_{i=0}^N \overline{R}_i - \overline{R}_{rand}
    @param Series:mean_returns A series of mean returns from each checkpoint
    @param float:mean_rand_returns The random baseline
    @returns float:str, Series:local_strs
    '''
    local_strs = mean_returns - mean_rand_returns
    str_ = local_strs.mean()
    return str_, local_strs


def calc_efficiency(local_strs, ts):
    '''<commit_msg>fix eval env ref for inference<commit_after>from slm_lab.lib import logger, util, viz
from slm_lab.spec import random_baseline
import numpy as np
import os
import pandas as pd
import pydash as ps
import shutil
import torch


METRICS_COLS = [
    'final_return_ma',
    'strength', 'max_strength', 'final_strength',
    'sample_efficiency', 'training_efficiency',
    'stability', 'consistency',
]

logger = logger.get_logger(__name__)


# methods to generate returns (total rewards)

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)
        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'''
    with util.ctx_lab_mode('eval'):  # enter eval context
        agent.algorithm.update()  # set explore_var etc. to end_val under ctx
    with torch.no_grad():
        ret = gen_return(agent, env)
    # exit eval context, restore variables simply by updating
    agent.algorithm.update()
    return ret


# metrics calculation methods

def calc_strength(mean_returns, mean_rand_returns):
    '''
    Calculate strength for metric
    str &= \frac{1}{N} \sum_{i=0}^N \overline{R}_i - \overline{R}_{rand}
    @param Series:mean_returns A series of mean returns from each checkpoint
    @param float:mean_rand_returns The random baseline
    @returns float:str, Series:local_strs
    '''
    local_strs = mean_returns - mean_rand_returns
    str_ = local_strs.mean()
    return str_, local_strs


def calc_efficiency(local_strs, ts):
    '''
{'fraction_non_alphanumeric': 0.0892774965485504, 'fraction_numerical': 0.0013805798435342844, 'mean_word_length': 3.644230769230769, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1857588', 'n_tokens_mistral': 1398, 'n_tokens_neox': 1349, 'n_words': 417}