<jupyter_start><jupyter_text>Config<jupyter_code>class Config():

    # Basics
    num_state_features = 5
    num_actions = 3
    num_timesteps = 2

    # Initial state distribution
    P0_mean = np.zeros(num_state_features)
    P0_cov = np.eye(num_state_features)

    # Transition function cov
    P_cov = np.eye(num_state_features)

    # Models
    ps_model_name = 'LinearPS'
    q_model_name = 'LSTDQ'

    # Simulation related params
    start_seed = 2021
    n_simulations = 10
    mc_samples = 10000

    # VSUP SSL
    labeled_data_prop = 0.3
    cv_splits = 5

config = Config()<jupyter_output><empty_output><jupyter_text>MDP<jupyter_code>def init_weights(formula):
    X = patsy.dmatrix(formula, data=balanced_data, return_type='dataframe')
    weights = pd.Series(np.random.uniform(-2, 2, X.shape[1]))
    weights.index = X.columns
    
    return weights
np.random.seed(config.start_seed)

# Define the MDP
# Initial State Distribution
P0 = InitialState(
    num_state_features=config.num_state_features,
    mean=config.P0_mean, 
    cov = config.P0_cov)

balanced_data = pd.concat([
    P0(config.num_actions),
    pd.DataFrame([str(a) for a in range(config.num_actions)], columns=['A'])], axis=1)
levels = list(balanced_data['A'].unique())

# Reward function
reward_formula = f"(0 + C(A, Treatment(reference='0'), levels={levels})) + (0 + C(A, Treatment(reference='0'), levels={levels})):({'+'.join(P0(1).columns)})"
reward_weights = init_weights(reward_formula)
R = Reward(
    formula=reward_formula,
    weights=reward_weights,
)

# Transition function
transition_formula = f"(0 + C(A, Treatment(reference='0'), levels={levels}))"
transition_weights = pd.concat([init_weights(transition_formula)]*config.num_state_features, axis=1)
transition_weights.columns = [f'S{s}' for s in range(config.num_state_features)]
P = Transition(formula=transition_formula, weights=transition_weights, cov=config.P_cov)

# Policy
policy_formula = '+'.join(P0(1).columns)
policy_weights = pd.concat([init_weights(policy_formula) for a in range(0, config.num_actions)], axis=1)
policy_weights.columns = [f'A{a}' for a in range(config.num_actions)]
policy_weights['A0'] = 0
Pi_b = Policy(formula=policy_formula, weights=policy_weights, deterministic=False)

# MDP
M = MDP(P0=P0, R=R, P=P, policy=Pi_b, num_actions=config.num_actions, num_timesteps=config.num_timesteps, num_state_features=config.num_state_features)<jupyter_output><empty_output><jupyter_text>Evaluation Policy<jupyter_code>eval_policy_weights = pd.concat([init_weights(policy_formula) for a in range(0,config.num_actions)], axis=1)
eval_policy_weights.columns = [f'A{a}' for a in range(config.num_actions)]
eval_policy_weights['A0'] = 0
eval_policy = Policy(formula=policy_formula, weights=eval_policy_weights, deterministic=True)
mc_estimates = list()
for i in tqdm(range(config.n_simulations)):
    eval_policy_data = M.sample_trajectories_for_policy(config.mc_samples, eval_policy, 
                                                        seed=config.start_seed + 6782 + i) 
    mc_estimates.append(
        sum(eval_policy_data[t]['R'].mean() for t in range(config.num_timesteps)))
q_model_params = {
    'formula': reward_formula,
    'eval_policy': eval_policy
}

ps_model_params = {
    'formula': policy_formula
}<jupyter_output><empty_output><jupyter_text>Supervised estimator<jupyter_code>vsup_dr_estimates = list()

for i in tqdm(range(config.n_simulations)):

    # Sample trajectories
    data = M.sample_trajectories(config.mc_samples, seed=config.start_seed+123123+i)

    # Models for Q function and Propensity score
    Q_model = construct_Q_model(model_name=config.q_model_name, model_params=q_model_params)
    Q_model.fit(data=data)

    PS_model = construct_PS_model(model_name=config.ps_model_name, model_params=ps_model_params)
    PS_model.fit(data=data)
    
    # Calculate supervised estimator
    vsup_dr_estimates.append(
        VSUP_DR(data=data, Q_model=Q_model, PS_model=PS_model)
    )
sns.distplot(vsup_dr_estimates)
sns.distplot(mc_estimates)<jupyter_output><empty_output><jupyter_text>Semi supervised estimator<jupyter_code>vssl_dr_estimates = list()

for i in tqdm(range(config.n_simulations)):

    # Sample trajectories
    data = M.sample_trajectories(config.mc_samples, seed=config.start_seed+123123+i)

    # Create surrogate variables
    for t in range(config.num_timesteps):
        data[t]['W'] = np.rint(data[t]['R'].values + np.random.normal(loc=0, scale=np.sqrt(2), size=config.mc_samples))

    # Randomly split into labeled and unlabeled data
    labeled_ids = np.random.choice(
        list(range(len(data[0]))),
        size=int(config.labeled_data_prop*len(data[0])), 
        replace=False)

    labeled_data = [df.iloc[labeled_ids].reset_index(drop=True) for df in data]
    unlabeled_data = [df.drop(index=labeled_ids).reset_index(drop=True) for df in data]

    # Models for Q function and Propensity score
    Q_model = construct_Q_model(model_name=config.q_model_name, model_params=q_model_params)
    Q_model.fit(data=labeled_data)

    PS_model = construct_PS_model(model_name=config.ps_model_name, model_params=ps_model_params)
    PS_model.fit(data=data)
    
    # Learn imputation model (Imputation)
    Imputation_model = ImputationLasso(
        cv_splits=config.cv_splits, 
        imputation_features=[f'S{s}' for s in range(config.num_state_features)] + ['A'] + ['W'],
        target='R',
        seed=config.start_seed
    )
    Imputation_model.fit(data=labeled_data)

    # Compute VSSL_DR estimate    
    vssl_dr_estimates.append(
        VSSL_DR(labeled_data=labeled_data, unlabeled_data=unlabeled_data, Q_model=Q_model, PS_model=PS_model, Imputation_model=Imputation_model)
    )

    
# Observed data
obs_data = M.sample_trajectories(config.mc_samples, seed=config.start_seed)

# Create surrogate variables 
for t in range(config.num_timesteps):
    obs_data[t]['W'] = np.rint(obs_data[t]['R'].values + np.random.normal(loc=0, scale=np.sqrt(2), size=obs_data[t].shape[0]))
# Estimate Q function for eval policy
Q_model.fit(obs_data)

# Estimate propensity scores for behavior policy
PS_model.fit(obs_data)
for t in range(config.num_timesteps):
    df = obs_data[t].copy(deep=True)

    if Q_model.eval_policy.deterministic:

        # Predicted action value for eval policy
        df['A'] = Q_model.eval_policy(data=df, timestep=t)
        obs_data[t]['V_eval_(S)'] = Q_model.predict(data=df, timestep=t)

        # Calculate probability of action under eval and behav policy
        obs_data[t]['Pr_eval_(A)'] = (df['A'] == obs_data[t]['A']).astype(int)
        behav_policy_prob_action = PS_model.predict_prob(
            data=obs_data[t], timestep=t)
        behav_policy_prob_action['A'] = obs_data[t]['A']
        obs_data[t]['Pr_(A)'] = behav_policy_prob_action.apply(
            lambda row: row[f'A{row["A"]}'], axis=1
        )
    else:
        raise NotImplementedError

    # Calculate importance sampling ratio
    obs_data[t]['rho'] = obs_data[t]['Pr_eval_(A)'] / obs_data[t]['Pr_(A)']
    if t > 0:
        obs_data[t]['rho'] *= obs_data[t-1]['rho']<jupyter_output><empty_output><jupyter_text>Create surrogate variables<jupyter_code># Randomly split into labeled and unlabeled data
labeled_ids = np.random.choice(list(range(obs_data[0].shape[0])), size=int(config.labeled_data_prop*obs_data[0].shape[0]), replace=False)

labeled_data = [data.iloc[labeled_ids].reset_index(drop=True) for data in obs_data]
unlabeled_data = [data.drop(index=labeled_ids).reset_index(drop=True) for data in obs_data]<jupyter_output><empty_output><jupyter_text>1. Imputation<jupyter_code>from sklearn.linear_model import Lasso, LassoCV
from sklearn.model_selection import KFold
imp_models = list()
imp_features = [f'S{s}' for s in range(config.num_state_features)] + ['A'] + ['W']

for t in range(config.num_timesteps):
    cv = KFold(n_splits=config.cv_splits)
    
    X = labeled_data[t][imp_features]
    y = labeled_data[t]['R']

    R_hat = pd.Series(index=labeled_data[t].index, dtype='float64')
    
    # Model selection
    # Find best parameters using cross validation
    model = LassoCV(cv=config.cv_splits, random_state=config.start_seed, normalize=True).fit(X, y)
    eps = model.eps

    # Getting prediction on every labeled data point
    # Get reward predictions using cross validation (different from earlier)
    imp_models.append([])
    for train_index, test_index in cv.split(X):

        # Train and test val data
        X_val_train, X_val_test = X.iloc[train_index], X.iloc[test_index]
        y_val_train, y_val_test = y.iloc[train_index], y.iloc[test_index]
        
        # Fit model
        model = Lasso(alpha=eps, normalize=True)
        model.fit(X_val_train, y_val_train)

        # Predict reward on test val data
        R_hat.iloc[test_index] = model.predict(X_val_test)

        imp_models[t].append(model)
    
    labeled_data[t]['R_hat'] = R_hat<jupyter_output><empty_output><jupyter_text>2. Refitting<jupyter_code>bias_coef = list()

for t in range(config.num_timesteps):
    bias_coef.append(
        np.mean(labeled_data[t]['rho']*(labeled_data[t]['R'] - labeled_data[t]['R_hat'] ))
    )<jupyter_output><empty_output><jupyter_text>3. Projection<jupyter_code>for t in range(config.num_timesteps):
    R_hat = list()
    for k in range(config.cv_splits):
        R_hat.append(imp_models[t][k].predict(unlabeled_data[t][imp_features]))
    unlabeled_data[t]['R_hat'] = np.vstack(R_hat).mean(axis=0)
    unlabeled_data[t]['mu_hat'] = (unlabeled_data[t]['R_hat']*unlabeled_data[t]['rho']) + bias_coef[t]<jupyter_output><empty_output><jupyter_text>Estimate<jupyter_code>VSSL_DR = unlabeled_data[0]['V_eval_(S)'].mean()
for t in range(config.num_timesteps):
    VSSL_DR += unlabeled_data[t]['mu_hat'].mean()

    if t < (config.num_timesteps-1):
        VSSL_DR -= (unlabeled_data[t]['rho'].clip(0, 1000) *\
            (unlabeled_data[t]['V_eval_(S)'] - unlabeled_data[t+1]['V_eval_(S)'])).mean()
                    
    else:
        VSSL_DR -= (unlabeled_data[t]['rho'].clip(0, 1000) *\
            unlabeled_data[t]['V_eval_(S)']).mean()

VSSL_DR
VSSL_DR
VSUP_DR(unlabeled_data, Q_model, PS_model)
unlabeled_data[t]
Q_function_formula = reward_formula
models = list()
for t in range(config.num_timesteps-1, -1, -1):

    df = obs_data[t]

    # Target
    if t == (config.num_timesteps-1):
        df['y'] = df['R']
    else:
        df['y'] = df['R'] + pred_future_value

    # Fit model
    model = smf.ols(formula='y ~ ' + Q_function_formula, data=df).fit()
    models.append(model)

    # Find the value for the action that would have been taken by the eval policy
    df_copy = df.copy(deep=True)        
    df_copy['A'] = eval_policy(df_copy)
    pred_future_value = model.predict(df_copy)

models = models[::-1]<jupyter_output><empty_output><jupyter_text>Q Learning<jupyter_code># Q_function_formula = reward_formula
# models = list()
# for t in range(config.num_timesteps-1, -1, -1):

#     df = obs_data[t]

#     # Target
#     if t == (config.num_timesteps-1):
#         df['y'] = df['R']
#     else:
#         df['y'] = df['R'] + max_value

#     # Fit model
#     model = smf.ols(formula='y ~ ' + Q_function_formula, data=df).fit()
#     models.append(model)

#     # Find the predicted value of each action
#     df_copy = df.copy(deep=True)        
#     pred_value = []
#     for a in range(config.num_actions):
#         df_copy['A'] = str(a)
#         pred_value.append(model.predict(df_copy))
#     pred_value = pd.concat(pred_value, axis=1)
#     pred_value.columns = columns=[f'A{d}' for d in range(config.num_actions)]

#     # Find the maximum predicted value
#     max_value = pred_value.max(axis=1)

# models = models[::-1]<jupyter_output><empty_output><jupyter_text>Propensity Score Models<jupyter_code>PS_model = policy_formula

ps_models = []
for t in range(0, config.num_timesteps, 1):
    ps_model = smf.mnlogit('A ~ ' + PS_model, data=obs_data[t].astype({'A': int}))
    ps_model = ps_model.fit(method='ncg') 
    ps_models.append(ps_model)

    print(f"Accuracy {t}: {(ps_model.predict(obs_data[t]).values.argmax(axis=1) == obs_data[t]['A'].astype(int)).mean()}")
ps_model.predict(obs_data[0])<jupyter_output><empty_output><jupyter_text>Eval Policy Value<jupyter_code>def calculat_avg_reward(data):
    reward = np.zeros(data[0].shape[0])
    for t in range(len(data)):
        reward += data[t]['R']
    
    return reward.mean()
# optimal_policy = ActionValuePolicy(models=models, num_actions=config.num_actions)<jupyter_output><empty_output><jupyter_text>1. Sampling<jupyter_code># optimal_policy_data = M.sample_trajectories_for_policy(10000, optimal_policy,seed=config.seed + 10) 
eval_policy_data = M.sample_trajectories_for_policy(100000, eval_policy, seed=config.seed + 120) 

print(calculat_avg_reward(eval_policy_data))
# behav_policy_data = M.sample_trajectories(10000, seed=2020) 

# print(calculat_avg_reward(behav_policy_data))<jupyter_output><empty_output><jupyter_text>2. VSUP<jupyter_code>behav_policy = list()
for t in range(config.num_timesteps):
    ps_model_weights = ps_models[t].params
    ps_model_weights.rename({a: f'A{a+1}' for a in ps_model_weights.columns}, axis=1, 
    inplace=True)
    ps_model_weights['A0'] = 0
    ps_model_weights = ps_model_weights[[f'A{a}' for a in range(config.num_actions)]]
    
    behav_policy.append(
        Policy(formula=PS_model, weights=ps_model_weights, deterministic=False)
    )
# Hack
try:
    for t in range(config.num_timesteps):
        for col in [f'Pr(A{a})' for a in range(config.num_actions)]:
            del obs_data[t][col]
except KeyError:
    pass

# # Calculate value of each state under behavior policy
# for t in range(config.num_timesteps):
#     pred_action_value = optimal_policy.calculate_action_value(obs_data[t], t)
#     pred_action_prob = behav_policy[t].calculate_action_prob(obs_data[t], t)
    
#     # pred_value = (pred_action_value.rename({col: col[5:-1] for col in pred_action_value.columns}, axis=1) * pred_action_prob).sum(axis=1)
#     pred_value = pred_action_value.max(axis=1)
    
#     obs_data[t]['V(S)'] = pred_value
#     obs_data[t] = pd.concat([obs_data[t], pred_action_prob.rename({col: f'Pr({col})' for col in pred_action_prob.columns}, axis=1)], axis=1)

# Calculate value of each state under behavior policy
for t in range(config.num_timesteps):
    df = obs_data[t].copy(deep=True)

    # Predicted action value for eval policy
    df['A'] = eval_policy(df, t)
    obs_data[t]['V(S)'] = models[t].predict(df)
    obs_data[t]['eval_A'] = df['A']
    
    # Prob of action taken by eval policy under behav policy
    pred_action_prob = behav_policy[t].calculate_action_prob(obs_data[t], t)
    obs_data[t] = pd.concat([obs_data[t], pred_action_prob.rename({col: f'Pr({col})' for col in pred_action_prob.columns}, axis=1)], axis=1)
# Calculate importance sampling weights
# for t in range(config.num_timesteps):
#     rho = obs_data[t].apply(lambda row: row[f'Pr(A{row["A"]})'], axis=1)
#     rho *= (obs_data[t]['A'] == optimal_policy(obs_data[t], t)).astype(int)

#     if t > 0:
#         rho *= obs_data[t-1]['rho']
    
#     obs_data[t]['rho'] = rho

for t in range(config.num_timesteps):
    rho = (obs_data[t]['A'] == obs_data[t]['eval_A']).astype(int)
    rho /= obs_data[t].apply(lambda row: row[f'Pr(A{row["A"]})'], axis=1)

    if t > 0:
        rho *= obs_data[t-1]['rho']
    try:
        del obs_data[t]['rho'] 
    except KeyError:
        pass
    
    obs_data[t]['rho'] = rho
VSUP_DR = obs_data[0]['V(S)'].mean()
for h in range(0, config.num_timesteps):
    if h < (config.num_timesteps-1):
        VSUP_DR += (obs_data[h]['rho'] * (obs_data[h]['R']-obs_data[h]['V(S)'] +obs_data[h+1]['V(S)'])).mean()
    else:
        VSUP_DR += (obs_data[h]['rho'] * (obs_data[h]['R'] -obs_data[h]['V(S)'])).mean()
print(VSUP_DR)
# Predict probability of each action at every timestep

# Predict Q(s, a) at every timestep

# Calculate V_st
# Timestep 0
t = 0
data_copy = obs_data[t].copy(deep=True)
pred_value = pd.DataFrame()
for a in range(config.num_actions):
    data_copy['A'] = str(a)
    pred_value[f'A{a}'] = models[t].predict(data_copy)

pred_value
ps_models[1]
# Replace Pi by PS_model
ps_models[1]
alt_policy = optimal_policy

VSUP_DR  = 0
for n, i in enumerate(list(np.random.randint(0, obs_data[0].shape[0], 1000))):
    VSUP_DR += models[0].predict(obs_data[0].iloc[[i]]).iloc[0]
    for t in range(config.num_timesteps):
        # VSUP_DR += models[t].predict(obs_data[t].iloc[[i]]).iloc[0]
        rho_h = 1
        for h in range(t):
            # pi_a_s = optimal_policy()

            if alt_policy.deterministic:
                alt_action = alt_policy(obs_data[h].iloc[[i]], h).iloc[0]
                actual_action = obs_data[h].iloc[[i]].iloc[0]['A']
                pi_a_s = 0 if (alt_action != actual_action) else 1
            else:
                raise NotImplementedError

            pib_a_s = Pi.calculate_action_prob(obs_data[h].iloc[i], h)
            pib_a_s = pib_a_s['A' + actual_action].iloc[0]
            
            rho_h *=  pi_a_s / pib_a_s
        

        td_target = obs_data[t].iloc[i]['R']
        # td_target = obs_data[t].iloc[i]['R'] + sum(
            # obs_data[f].iloc[i]['R'] for f in range(t+1, config.num_timesteps))
        try:
            td_target += models[t+1].predict(obs_data[t+1].iloc[[i]]).iloc[0]
        except IndexError:
            pass
        delta_td = td_target - models[t].predict(obs_data[t].iloc[[i]]).iloc[0]
        VSUP_DR += rho_h*delta_td

    
    if n % 50 == 0:    
        print(VSUP_DR/(n+1))
 obs_data[h].iloc[[i]]
obs_data[h].iloc[i]['A']
pib_a_s
from typing import Dict

import numpy as np
from numpy.random import multinomial, multivariate_normal, uniform, normal
from scipy.special import softmax
import pandas as pd
import patsy
import statsmodels.formula.api as smf
import statsmodels.api as sm
from sklearn.preprocessing import OneHotEncoder

seed = 2021
np.random.seed(seed)
# Config
num_timesteps = 5
num_actions = 3
state_dim = 5
size = 10000

# Mean of MVN for drawing state at timestep t
mu0 = np.zeros(state_dim)
sigma = np.eye(state_dim)

# Coefficient for each action in the transition function
gamma = uniform(-2, 2, size=num_actions)

# Regression coefficients for outcome
beta = uniform(-2, 2, size=(state_dim+1)*num_actions*num_timesteps).reshape(num_timesteps, num_actions, state_dim+1)
# beta_a = uniform(-2, 2, size=num_timesteps*num_actions).reshape(num_timesteps, num_actions)
# beta = uniform(-2, 2, size=state_dim*num_actions).reshape(num_actions, state_dim)
# beta_a = uniform(-2, 2, size=num_actions)

# Regression coefficients for treatment 
# phi = uniform(-5, 2, size=(state_dim+1)*(num_actions-1)*num_timesteps).reshape(num_timesteps, num_actions-1, state_dim+1)
# phi = np.concatenate([np.zeros((num_timesteps, 1, state_dim+1)), phi], 1)
phi = uniform(-2, 2, size=(state_dim+1)*(num_actions-1)).reshape(num_actions-1, state_dim+1)
phi = np.concatenate([phi, np.zeros((1, state_dim+1))], 0)

print("gamma", gamma.shape)
print("beta", beta.shape)
print("phi", phi.shape)
print("mu0", mu0.shape)
print("sigma", sigma.shape)
data = list()

for t in range(num_timesteps):
    
    # Generate state
    if t == 0:
        S_t = multivariate_normal(mu0, sigma, size=size)
    else:
        S_t = np.vstack([multivariate_normal(S_tminus1[i] + A_tminus1[i]@gamma, sigma, size=1) for i in range(size)])
    S_tdash = np.concatenate([np.ones((size, 1)), S_t], 1)

    # Generate action
    pvals = softmax(S_tdash @ phi.T, axis=1)
    A_t = np.vstack([multinomial(n=1, pvals=pvals[i], size=1) for i in range(size)])
    
    # Generate reward
    mean_R_tplus1 = np.vstack([(S_tdash@beta[t, a]) * (A_t[:, a]) for a in range(num_actions)]).sum(axis=0)
    R_tplus1 = normal(loc=mean_R_tplus1, scale=1)
    
    # Convert to dataframe
    data.append(
        pd.concat([
            pd.DataFrame(S_t, columns=[f'S{d}' for d in range(1, state_dim + 1)]),
            pd.DataFrame(np.argmax(A_t, axis=1), columns=['A']).astype(object),
            pd.DataFrame(R_tplus1, columns=['R'])
        ], axis=1)
    )

    # Transition
    A_tminus1 = A_t
    S_tminus1 = S_t

    print(A_t.sum(axis=0))

# Ensure one action of each type!
# TO DO #<jupyter_output>[2728 4323 2949]
[4199 3948 1853]
[4855 3679 1466]
[5230 3435 1335]
[5526 3311 1163]<jupyter_text>Q-learning<jupyter_code>Q_function_model = f"A*({'+'.join(f'S{d}' for d in range(1, state_dim+1))})"
print(Q_function_model)
models = list()

for t in range(num_timesteps-1, -1, -1):

    df = data[t]

    # Target
    if t == (num_timesteps-1):
        df['y'] = df['R']
    else:
        df['y'] = df['R'] + max_value

    # Fit model
    model = smf.ols(formula='y ~ ' + Q_function_model, data=df).fit()
    models.append(model)

    # Find the predicted value of each action
    df_copy = df.copy(deep=True)        
    pred_value = []
    for a in range(num_actions):
        df_copy['A'] = a
        pred_value.append(model.predict(df_copy))
    pred_value = pd.concat(pred_value, axis=1)
    pred_value.columns = columns=[f'A{d}' for d in range(num_actions)]

    # Find the maximum predicted value
    max_value = pred_value.max(axis=1)

models = models[::-1]<jupyter_output><empty_output><jupyter_text>Propensity score models<jupyter_code>PS_model = f"{'+'.join(f'S{d}' for d in range(1, state_dim+1))}"

ps_models = []
for t in range(0, num_timesteps, 1):
    ps_model = smf.mnlogit('A ~ ' + PS_model, data=data[t].astype({'A': int}))
    ps_model = ps_model.fit(method='ncg')
    ps_models.append(ps_model)

    print(f"Accuracy {t}: {(ps_model.predict(data[t]).values.argmax(axis=1) == data[t]['A']).mean()}")<jupyter_output>Optimization terminated successfully.
         Current function value: 0.577084
         Iterations: 8
         Function evaluations: 9
         Gradient evaluations: 16
         Hessian evaluations: 8
Accuracy 0: 0.7515
Optimization terminated successfully.
         Current function value: 0.311172
         Iterations: 9
         Function evaluations: 10
         Gradient evaluations: 18
         Hessian evaluations: 9
Accuracy 1: 0.8717
Optimization terminated successfully.
         Current function value: 0.201014
         Iterations: 11
         Function evaluations: 12
         Gradient evaluations: 22
         Hessian evaluations: 11
Accuracy 2: 0.9147
Optimization terminated successfully.
         Current function value: 0.149221
         Iterations: 11
         Function evaluations: 12
         Gradient evaluations: 22
         Hessian evaluations: 11
Accuracy 3: 0.9384
Optimization terminated successfully.
         Current function value: 0.116937
         Iterations: 12
     [...]<jupyter_text>Optimal value function estimation<jupyter_code>def get_dummies(x, num_levels):
    dummy_data = np.zeros((x.shape[0], num_levels))
    for i, xi in enumerate(x):
        dummy_data[i, xi] = 1
    return dummy_data
optimal_policy_data = list()

for t in range(num_timesteps):
    
    # Generate state
    if t == 0:
        S_t = multivariate_normal(mu0, sigma, size=size)
    else:
        S_t = np.vstack([multivariate_normal(S_tminus1[i] + A_tminus1[i]@gamma, sigma, size=1) for i in range(size)])
    S_tdash = np.concatenate([np.ones((size, 1)), S_t], 1)
    
    # Choose action
    # Find the predicted value of each action
    df_copy = pd.DataFrame(S_t, columns=[f'S{d}' for d in range(1, state_dim + 1)])
    pred_value = []
    for a in range(num_actions):
        df_copy['A'] = a
        pred_value.append(models[t].predict(df_copy))
    pred_value = pd.concat(pred_value, axis=1)
    pred_value.columns = columns=[f'A{d}' for d in range(num_actions)]
    # Find the best action
    A_t = get_dummies(pred_value.values.argmax(axis=1), num_actions)
    
    # Generate reward
    mean_R_tplus1 = np.vstack([(S_tdash@beta[t, a]) * (A_t[:, a]) for a in range(num_actions)]).sum(axis=0)
    R_tplus1 = normal(loc=mean_R_tplus1, scale=1)
    
    # Convert to dataframe
    optimal_policy_data.append(
        pd.concat([
            pd.DataFrame(S_t, columns=[f'S{d}' for d in range(1, state_dim + 1)]),
            pd.DataFrame(np.argmax(A_t, axis=1), columns=['A']).astype(object),
            pd.DataFrame(R_tplus1, columns=['R'])
        ], axis=1)
    )

    # Transition
    A_tminus1 = A_t
    S_tminus1 = S_t

# Ensure one action of each type!
# TO DO #
reward = np.zeros(size)
for t in range(num_timesteps):
    reward += optimal_policy_data[t]['R']
print(reward.mean())
reward = np.zeros(size)
for t in range(num_timesteps):
    reward += data[t]['R']
print(reward.mean())
model.predict(df_copy)<jupyter_output><empty_output>
StarcoderdataJupyterStructuredCleanDedup