# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     name: python3
# ---

# + [markdown] id="view-in-github" colab_type="text"
# <a href="https://colab.research.google.com/github/dandanelbaz/agents/blob/master/ai_week.ipynb" target="_parent"><img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab"/></a>

# + [markdown] id="ongEoQjaTJQk" colab_type="text"
# # Intstall coach
# Just use pip

# + id="sRS7Z_lfGrs6" colab_type="code" colab={}
pip install rl_coach

# + [markdown] colab_type="text" id="9qn7Qf9gHZLk"
# # AI Week Workshop 

# + [markdown] colab_type="text" id="kEcvEIxqHZLt"
# ### ***Add new environmen***t
#
# In this section we will implement the short corridor environment from Sutton & Barto Book.
#
# ![alt text](https://drive.google.com/uc?id=1rYLI9dC92sfpF0BVxVENF964MfWJkxZq)
#
# *   Three non terminal states- The location of the agent
#
# *   The observations are one-hot encoding of the states
# *   Actions are reversed in the second state
#
#
# *   Reward is -1 for each time step
#
#
#
#
#

# + [markdown] id="G2__PdzbNrPi" colab_type="text"
# ##### ***Helper function*** 
# The following code snippet contains some defines and an one-hot encoding helper function.

# + id="lnMACvXo_y3U" colab_type="code" colab={"base_uri": "https://localhost:8080/", "height": 35} outputId="54c9c4f8-cac8-4086-c9ce-a6f5860d7c56"
# %%writefile short_corridor_env_helpper.py
import numpy as np

LEFT = 0
RIGHT = 1
START_STATE = 0
GOAL_STATE = 3
NUM_STATES = 4
REVERSE_STATE = 1

def to_one_hot(state):
    observation = np.zeros((NUM_STATES,))
    observation[state] = 1
    return observation


# + [markdown] colab_type="text" id="RBxwrlSzHZLu"
# ##### ***Write short corridor environment*** 
# Compete the following functions:
#  function and the step function
#
# 1.   is_done - will return a boolean . True only at termination state
#
# 2.   reset - Resets environment to initial state
# 3.   step - Returns the next observation, reward, and the boolean flag done
#
#
#
#
#
# * **complete code**
#

# + colab_type="code" outputId="02dae064-339e-4997-a111-b64fced6448e" id="uxwi5S1vHZLw" colab={"base_uri": "https://localhost:8080/", "height": 35}
# %%writefile short_corridor_env.py
import numpy as np
import gym
from gym import spaces
from  short_corridor_env_helpper import *


class ShortCorridorEnv(gym.Env):

    def __init__(self):
        # Class constructor- Initializes class variables and sets initial state
        self.observation_space = spaces.Box(0, 1, shape=(NUM_STATES,))
        self.action_space = spaces.Discrete(2)
        self.reset()

    def reset(self):
        '''
        Resets the environment to start state
        '''
        # Boolean. True only if the goal state is reached
        self.goal_reached = ???
        # An integer representing the state. Number between zero and three
        self.current_state = ???
        observation = to_one_hot(???)
        return observation

    def _is_done(self, current_state):
        '''
        return done a Boolean- True only if we reached the goal state
        '''
        # ???
        return done

    def step(self, action):
        '''
        Returns the next observation, reward, and the boolean flag done
        '''

        if action ==LEFT:
          step = -1
        elif action == RIGHT:
          step = ???

        if self.current_state == REVERSE_STATE:
        ### Replace step = -1 with step = 1 and vise versa
            # ???

        self.current_state += step
        self.current_state = max(0, self.current_state)

        observation = to_one_hot(self.current_state)
        reward = ???
        done = self._is_done(self.current_state)

        return observation, reward, done, {}



# + [markdown] colab_type="text" id="tO6bmhG1HZL0"
# ##### ***Write preset to run existing agent on the new environment***
# *We will use the same preset from DQN example*.
#
# Since our environment is already using Gym API we are almost good to go.
#
# When selecting the environment parametes in the preset use **GymEnvironmentParameters** and pass the path of the environment source code using the level parameter

# + colab_type="code" id="wVmJ2CfbHZL1" outputId="b0efca7e-119e-4f0c-9026-1c7bbe706796" colab={"base_uri": "https://localhost:8080/", "height": 35}
# %%writefile short_corridor_dqn_preset.py
from rl_coach.environments.gym_environment import GymEnvironmentParameters
from rl_coach.filters.filter import NoInputFilter, NoOutputFilter
from rl_coach.agents.dqn_agent import DQNAgentParameters
from rl_coach.graph_managers.basic_rl_graph_manager import BasicRLGraphManager
from rl_coach.graph_managers.graph_manager import SimpleSchedule
from rl_coach.memories.memory import MemoryGranularity


####################
# Graph Scheduling #
####################
schedule_params = SimpleSchedule()


#########
# Agent #
#########
agent_params = DQNAgentParameters()
agent_params.input_filter = NoInputFilter()
agent_params.output_filter = NoOutputFilter()
# DQN params
# ER size
agent_params.memory.max_size = (MemoryGranularity.Transitions, 40000)


###############
# Environment #
###############
env_params = GymEnvironmentParameters(level='short_corridor_env:ShortCorridorEnv')


#################
# Graph Manager #
#################
graph_manager = BasicRLGraphManager(agent_params=agent_params,
                                    env_params=env_params,
                                    schedule_params=schedule_params)


# + [markdown] colab_type="text" id="mCwoQwnGHZL5"
# ##### ***Run new preset***

# + colab_type="code" id="4db-w4ETHZL6" colab={}
# !coach -p /content/short_corridor_dqn_preset.py:graph_manager


# + [markdown] colab_type="text" id="aMs0JnUBHZL8"
# ### ***Add new agent***
# Coach modularity makes adding an agent a clean and simple task.
# Tipicaly consists of four parts:
#
#
# 1.   Implement an agent spesific network head (and loss)
# 2.   Implement exploration policy (optional)
# 3.   Define new parametes class that extends `AgentParametes`
# 4.   Implement a preset to run the agent on some environment
#
#

# + [markdown] id="MtfJT8XbmsUo" colab_type="text"
# ##### ***Write stochastic output layer***
# We use stochastic policy, meaning that we only produce the probability of going left and going right.
# This layer takes in the input from previous layer, the middleware, and outputs two numbers. 

# + id="KIYd8t8GnCSu" colab_type="code" outputId="76bb7100-76d1-4be3-cc55-22e4cc7acbd8" colab={"base_uri": "https://localhost:8080/", "height": 35}
# %%writefile probabilistic_layer.py
import tensorflow as tf
from rl_coach.architectures.tensorflow_components.layers import Dense

class ProbabilisticLayer(object):
    def __init__(self, input_layer, num_actions):
        super().__init__()
        scores = Dense(num_actions)(input_layer, name='logit')
        self.event_probs = tf.nn.softmax(scores, name="policy")
        # define the distributions for the policy and the old policy
        self.policy_distribution = tf.contrib.distributions.Categorical(probs=self.event_probs)

    def log_prob(self, action):
        return self.policy_distribution.log_prob(action)

    def layer_output(self):
        return self.event_probs


# + [markdown] colab_type="text" id="wNelbGV7HZME"
# ##### ***Implement network head i.e. implement the loss***
# The Head needs to inherit from the base class `Head`.
#
# Inorder to maximize the sum of rewards, we want to go in the following direction $-\Sigma_i A_i \nabla_Wlog(\pi(a_i|x_i))$
#
# $- A_i \nabla_Wlog[\pi(a_i|x_i)]$
#
# `Complete code`
#
#

# + colab_type="code" outputId="dd33ee5d-2eba-4e45-a230-009bd4b37fd2" id="hUnA2ZJHHZMF" colab={"base_uri": "https://localhost:8080/", "height": 35}
# %%writefile simple_pg_head.py
import tensorflow as tf
from rl_coach.architectures.tensorflow_components.heads.head import Head
from rl_coach.base_parameters import AgentParameters
from rl_coach.spaces import SpacesDefinition
from probabilistic_layer import ProbabilisticLayer


class SimplePgHead(Head):
    def __init__(self, agent_parameters: AgentParameters,
                 spaces: SpacesDefinition, network_name: str,
                 head_idx: int = 0, is_local: bool = True):
        super().__init__(agent_parameters, spaces, network_name)

        self.exploration_policy = agent_parameters.exploration

    def _build_module(self, input_layer):
        # Define inputs
        actions = tf.placeholder(tf.int32, [None], name="actions")
        advantages = tf.placeholder(tf.float32, [None], name="advantages")

        # Two actions, left or right
        policy_distribution = ProbabilisticLayer(input_layer, num_actions=2)

        # calculate loss
        log_prob = policy_distribution.log_prob(???)
        mudulated_log_prob = ???
        expected_mudulated_log_prob = tf.reduce_mean(mudulated_log_prob)

        ### Coach bookkeeping
        # List of placeholders for additional inputs to the head 
        #(except from the middleware input)
        self.input.append(???)
        # The output of the head, which is also the output of the network.
        self.output.append(???)
        # Placeholder for the target that we will use to train the network
        self.target = ???
        # The loss that we will use to train the network
        self.loss = ???
        tf.losses.add_loss(self.loss)



# + [markdown] colab_type="text" id="3Eqej3v50a91"
# ##### ***Define exploration policy*** 
# Every iteration we want to sample from the network output distribution i.e. toss a bias coin to get the agent actual move
#
# **`Complete code`**

# + colab_type="code" outputId="4efed15d-3956-4ff1-f072-4ec6592da0e0" id="zuYBPgnI0a-A" colab={"base_uri": "https://localhost:8080/", "height": 35}
# %%writefile simple_pg_exploration.py

import numpy as np
from rl_coach.exploration_policies.exploration_policy import ExplorationPolicy, ExplorationParameters
from rl_coach.spaces import ActionSpace


class DiscreteExplorationParameters(ExplorationParameters):
    @property
    def path(self):
        return 'simple_pg_exploration:DiscreteExploration'


class DiscreteExploration(ExplorationPolicy):
    """
    Discrete exploration policy is intended for discrete action spaces. It expects the action values to
    represent a probability distribution over the action
    """
    def __init__(self, action_space: ActionSpace):
        """
        :param action_space: the action space used by the environment
        """
        super().__init__(action_space)

    def get_action(self, probabilities):
        # choose actions according to the probabilities
        action = np.random.choice(self.action_space.actions, p=???)
        return chosen_action, probabilities



# + [markdown] colab_type="text" id="XG5y-G43HZL9"
# ##### ***Define new agent parameters***
# Coach is modular!
#
# Each class in Coach has a complementary parameters class which defines its constructor. 
# This is also true for the agent. The agent has a complementary `AgentParameters` class. This class enable to select the paramenters of the agent sub modules.
#
# It consists of the following four parts:
#
#
#
# 1.   algorithm
# 2.   exploration
# 3.   memory
# 4.   Networks
#
#

# + colab_type="code" outputId="047d9045-f825-4456-8340-328db9c78bac" id="bTAcPbnsHZL-" colab={"base_uri": "https://localhost:8080/", "height": 35}
# %%writefile simple_pg_params.py
from rl_coach.architectures.embedder_parameters import InputEmbedderParameters
from rl_coach.architectures.head_parameters import HeadParameters
from rl_coach.architectures.middleware_parameters import FCMiddlewareParameters
from rl_coach.base_parameters import NetworkParameters, AlgorithmParameters, \
    AgentParameters

from rl_coach.exploration_policies.additive_noise import AdditiveNoiseParameters
from rl_coach.exploration_policies.categorical import CategoricalParameters
from rl_coach.memories.episodic.single_episode_buffer import SingleEpisodeBufferParameters
from rl_coach.spaces import DiscreteActionSpace, BoxActionSpace
from rl_coach.agents.policy_optimization_agent import PolicyGradientRescaler
from simple_pg_exploration import DiscreteExplorationParameters

class SimplePgAgentParameters(AgentParameters):
    def __init__(self):
        super().__init__(algorithm=SimplePGAlgorithmParameters(),
                         #exploration=CategoricalParameters(),
                         exploration=DiscreteExplorationParameters(),
                         memory=SingleEpisodeBufferParameters(),
                         networks={"main": SimplePgTopology()})
    @property
    def path(self):
        #return 'simple_pg_agent:SimplePgAgent'
        return 'rl_coach.agents.policy_gradients_agent:PolicyGradientsAgent'

        
    
# Since we are adding a new head we need to tell coach the heads path
class SimplePgHeadParams(HeadParameters):
    def __init__(self):
        super().__init__(parameterized_class_name="AiWeekHead")

    @property
    def path(self):
        return 'simple_pg_head:SimplePgHead'


class SimplePgTopology(NetworkParameters):
    def __init__(self):
        super().__init__()
        self.input_embedders_parameters = {'observation': InputEmbedderParameters()}
        self.middleware_parameters = FCMiddlewareParameters()
        self.heads_parameters = [SimplePgHeadParams()]


class SimplePGAlgorithmParameters(AlgorithmParameters):
    """
    :param num_steps_between_gradient_updates: (int)
        The number of steps between calculating gradients for the collected data. In the A3C paper, this parameter is
        called t_max. Since this algorithm is on-policy, only the steps collected between each two gradient calculations
        are used in the batch.
    """
    def __init__(self):
        super().__init__()
        # TOTAL_RETURN
        # FUTURE_RETURN
        # FUTURE_RETURN_NORMALIZED_BY_EPISODE 
        # FUTURE_RETURN_NORMALIZED_BY_TIMESTEP
        # Q_VALUE
        # A_VALUE
        # TD_RESIDUAL
        # DISCOUNTED_TD_RESIDUAL
        # GAE
        self.policy_gradient_rescaler = PolicyGradientRescaler.FUTURE_RETURN
        self.num_steps_between_gradient_updates = 20000  # this is called t_max in all the papers






# + [markdown] colab_type="text" id="_8RT_QFdHZMH"
# ##### ***Write preset to run new agent on short corridor***
# complete code
# * **complete code**
# * **Hint: look at DQN preset**
#

# + colab_type="code" outputId="ce90a661-5691-4b91-8bf0-cc9af56df38f" id="D1PLQuNaHZMI" colab={"base_uri": "https://localhost:8080/", "height": 35}
# %%writefile short_corridor_new_agent_preset.py
from rl_coach.base_parameters import VisualizationParameters
from rl_coach.core_types import EnvironmentEpisodes, EnvironmentSteps
from rl_coach.environments.gym_environment import GymEnvironmentParameters
from rl_coach.filters.filter import NoInputFilter, NoOutputFilter
from rl_coach.graph_managers.basic_rl_graph_manager import BasicRLGraphManager
from rl_coach.graph_managers.graph_manager import SimpleSchedule
from rl_coach.memories.memory import MemoryGranularity
from rl_coach.schedules import LinearSchedule
from simple_pg_params import SimplePgAgentParameters


####################
# Graph Scheduling #
####################
schedule_params = SimpleSchedule()


#########
# Agent #
#########
agent_params = ???
agent_params.input_filter = NoInputFilter()
agent_params.output_filter = NoOutputFilter()


###############
# Environment #
###############
env_params = GymEnvironmentParameters(level='short_corridor_env:ShortCorridorEnv')

#################
# Graph Manager #
#################
graph_manager = BasicRLGraphManager(agent_params=agent_params,
                                    env_params=env_params,
                                    schedule_params=schedule_params)



# + [markdown] colab_type="text" id="Ufc_EAe3HZMK"
# ##### ***Run preset of the new agent on the new environment***
#
# **`Complete code`**
#
#
#

# + colab_type="code" id="ad2xpEu8HZML" colab={}
# ???

StarcoderdataJupyterScriptsDedupFiltered