# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 2
#     language: python
#     name: python2
# ---

# +
# %matplotlib inline

import gym
import matplotlib
import numpy as np
import sys

from collections import defaultdict
if "../" not in sys.path:
  sys.path.append("../") 
from lib.envs.blackjack import BlackjackEnv
from lib import plotting

matplotlib.style.use('ggplot')
# -

env = BlackjackEnv()


def create_random_policy(nA):
    """
    Creates a random policy function.
    
    Args:
        nA: Number of actions in the environment.
    
    Returns:
        A function that takes an observation as input and returns a vector
        of action probabilities
    """
    A = np.ones(nA, dtype=float) / nA
    def policy_fn(observation):
        return A
    return policy_fn


def create_greedy_policy(Q):
    """
    Creates a greedy policy based on Q values.
    
    Args:
        Q: A dictionary that maps from state -> action values
        
    Returns:
        A function that takes an observation as input and returns a vector
        of action probabilities.
    """
    
    def policy_fn(observation):
        A = np.zeros( len(Q[observation]), dtype=float)
        best_action = np.argmax(Q[observation])
        A[best_action] = 1.0

    return policy_fn


def mc_control_importance_sampling(env, num_episodes, behavior_policy, discount_factor=1.0):
    """
    Monte Carlo Control Off-Policy Control using Weighted Importance Sampling.
    Finds an optimal greedy policy.
    
    Args:
        env: OpenAI gym environment.
        num_episodes: Nubmer of episodes to sample.
        behavior_policy: The behavior to follow while generating episodes.
            A function that given an observation returns a vector of probabilities for each action.
        discount_factor: Lambda discount factor.
    
    Returns:
        A tuple (Q, policy).
        Q is a dictionary mapping state -> action values.
        policy is a function that takes an observation as an argument and returns
        action probabilities. This is the optimal greedy policy.
        
        Off-policy every-visit MC control (returns π ≈ π∗)
Initialize, for all s ∈ S, a ∈ A(s):
Q(s, a) ← arbitrary
C(s,a) ← 0
π(s) ← a deterministic policy that is greedy with respect to Q
Repeat forever:
Generate an episode using any soft policy μ:
S0,A0,R1,...,ST−1,AT−1,RT ,ST G←0
W←1
Fort=T−1,T−2,... downto0:
G ← γG + Rt+1
C(St, At) ← C(St, At) + W
￼Q(St, At) ← Q(St, At) + W/C(St,At) [G − Q(St, At)]
π(St) ← argmaxa Q(St,a) (with ties broken consistently)
￼If At ̸= π(St)

W←W 1/μ(At |St )
    """
    
    # The final action-value function.
    # A dictionary that maps state -> action values
    Q = defaultdict(lambda: np.zeros(env.action_space.n))
    C = defaultdict(lambda: np.zeros(env.action_space.n))

    # Our greedily policy we want to learn
    target_policy = create_greedy_policy(Q)
    actions = np.array(xrange(env.action_space.n))

    for ep in range(num_episodes):
        observation = env.reset()
        episode = []
        done = False
        if num_episodes <= 5:
            print "starting episode %d" % ep
            
        while not done:
            action_dist = behavior_policy(observation)
            action = np.random.choice(actions, p=action_dist)
            next_obs, reward, done, _ = env.step(action)
            episode.append( (observation, action, reward) )
            observation = next_obs
            
        if num_episodes <= 5:
            print "length: %d" % len(episode)
        G = 0.0
        W = 1.0
        for t in range( len(episode)-1, -1, -1 ):
            observation, action, reward = episode[t]
            if num_episodes <= 5:
                print "%d %f" % (action, reward)
            G = (discount_factor * G) + reward
            C[observation][action] += W
            Q[observation][action] += (W / C[observation][action]) * (G - Q[observation][action])
            if action != np.argmax(target_policy(observation)):
                break
            W = W / behavior_policy(observation)[action]
        
    return Q, target_policy


random_policy = create_random_policy(env.action_space.n)
Q, policy = mc_control_importance_sampling(env, num_episodes=5, behavior_policy=random_policy)

random_policy = create_random_policy(env.action_space.n)
Q, policy = mc_control_importance_sampling(env, num_episodes=500000, behavior_policy=random_policy)

# For plotting: Create value function from action-value function
# by picking the best action at each state
V = defaultdict(float)
for state, action_values in Q.items():
    action_value = np.max(action_values)
    V[state] = action_value
plotting.plot_value_function(V, title="Optimal Value Function")



StarcoderdataJupyterScriptsDedupFiltered