<kaggle_start><code># This is a two part tutorial series on TF-Agents
# Part 1 : https://www.kaggle.com/usharengaraju/tfagents-environment-policy-driver
# Part 2 : https://www.kaggle.com/usharengaraju/tf-agents-replay-buffer-network-checkpointer
# Credit : The article series has been adapted from the official tensorflow documentation.
# ## Reinforcement Learning
# Reinforcement learning (RL) is a machine learning framework where agents takes action in an environment in order to maximize the cumulative reward .
# ![](https://drive.google.com/uc?id=1ayosPHB3DIbmgnuu0kkN-dpA1TfnurNo)
# Pic Credit : ai.googleblog.com
# ### TF-Agents :
# A reliable, scalable and easy to use TensorFlow library for Contextual Bandits and Reinforcement Learning.
# ![](https://drive.google.com/uc?id=1zjzi0OY7G-QusVpYppjdVvbAP-pX0-RZ)
# Pic Credit : Tensorflow Dev Summit
# ### Advantages of TF-Agents :
# 🎯 Great Resources to get started with Reinforcement Learning like colab , documentation and resources
# 🎯 Well suited to handle complex RL problems
# 🎯 Helps in developing RL algorithms quickly
# 🎯 Can be configured easily with gin-config
# ![](https://drive.google.com/uc?id=1QBCM4TbPlaoftfpXZVoQ7uUScvDtqga0)
# Pic Credit : Tensorflow Dev Summit
# In this notebook , we will be discussing three concepts - Environment , Policy and Driver
# ### Environment
# The agent receives an observation and chooses an action which it applies on the environment and gets back reward
# and a observation from the environment. The goal of the agent is to train a policy to choose actions which will
# maximize the cumulative rewards.
# TF-Agents has both Python and TensorFlow implentation of environments .
#
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import abc
import tensorflow as tf
import numpy as np
from tf_agents.environments import py_environment
from tf_agents.environments import tf_environment
from tf_agents.environments import tf_py_environment
from tf_agents.environments import utils
from tf_agents.specs import array_spec
from tf_agents.environments import wrappers
from tf_agents.environments import suite_gym
from tf_agents.trajectories import time_step as ts
from tf_agents.specs import array_spec
from tf_agents.specs import tensor_spec
from tf_agents.networks import network
from tf_agents.policies import py_policy
from tf_agents.policies import random_py_policy
from tf_agents.policies import scripted_py_policy
from tf_agents.policies import tf_policy
from tf_agents.policies import random_tf_policy
from tf_agents.policies import actor_policy
from tf_agents.policies import q_policy
from tf_agents.policies import greedy_policy
from tf_agents.environments import suite_gym
from tf_agents.metrics import py_metrics
from tf_agents.metrics import tf_metrics
from tf_agents.drivers import py_driver
from tf_agents.drivers import dynamic_episode_driver

tf.compat.v1.enable_v2_behavior()
from tf_agents.trajectories import time_step as ts

tf.compat.v1.enable_v2_behavior()
tf.compat.v1.enable_v2_behavior()


# ### Python Environments :
# The interface that all python environments must implement is in environments/py_environment.PyEnvironment.
# After interaction with the environment the following information about the next step are available
# 📌 **observation:** This is the part of the environment state that the agent can observe to choose its actions at the next step.
# 📌 **reward:** The agent is learning to maximize the sum of these rewards across multiple steps.
# 📌 **step_type:** Interactions with the environment are usually part of a sequence/episode. e.g. multiple moves in a game of chess. step_type can be either FIRST, MID or LAST to indicate whether this time step is the first, intermediate or last step in a sequence.
# 📌 **discount:** This is a float representing how much to weight the reward at the next time step relative to the reward at the current time step.
#
class PyEnvironment(object):
    def reset(self):
        """Return initial_time_step."""
        self._current_time_step = self._reset()
        return self._current_time_step

    def step(self, action):
        """Apply action and return new time_step."""
        if self._current_time_step is None:
            return self.reset()
        self._current_time_step = self._step(action)
        return self._current_time_step

    def current_time_step(self):
        return self._current_time_step

    def time_step_spec(self):
        """Return time_step_spec."""

    @abc.abstractmethod
    def observation_spec(self):
        """Return observation_spec."""

    @abc.abstractmethod
    def action_spec(self):
        """Return action_spec."""

    @abc.abstractmethod
    def _reset(self):
        """Return initial_time_step."""

    @abc.abstractmethod
    def _step(self, action):
        """Apply action and return new time_step."""
        self._current_time_step = self._step(action)
        return self._current_time_step


# ### Standard Environments :
# TF Agents has built-in wrappers for many standard environments like the OpenAI Gym, DeepMind-control and Atari, which can be easily loaded using our environment suites. The code below loads the CartPole environment from the OpenAI gym
#
environment = suite_gym.load("CartPole-v0")
print("action_spec:", environment.action_spec())
print("time_step_spec.observation:", environment.time_step_spec().observation)
print("time_step_spec.step_type:", environment.time_step_spec().step_type)
print("time_step_spec.discount:", environment.time_step_spec().discount)
print("time_step_spec.reward:", environment.time_step_spec().reward)


# ## Policy
# Policies map an observation from the environment to an action or a distribution over actions.Most policies have a neural network to compute actions and/or distributions over actions from TimeSteps.Policies can be saved/restored, and can be used indepedently of the agent for data collection, evaluation etc.
# Policies contain the following information
# 📌 **action:** The action to be applied to the environment.
# 📌 **state:** The state of the policy (e.g. RNN state) to be fed into the next call to action.
# 📌 **info:** Optional side information such as action log probabilities.
#
class Base(object):
    @abc.abstractmethod
    def __init__(self, time_step_spec, action_spec, policy_state_spec=()):
        self._time_step_spec = time_step_spec
        self._action_spec = action_spec
        self._policy_state_spec = policy_state_spec

    @abc.abstractmethod
    def reset(self, policy_state=()):
        # return initial_policy_state.
        pass

    @abc.abstractmethod
    def action(self, time_step, policy_state=()):
        # return a PolicyStep(action, state, info) named tuple.
        pass

    @abc.abstractmethod
    def distribution(self, time_step, policy_state=()):
        # Not implemented in python, only for TF policies.
        pass

    @abc.abstractmethod
    def update(self, policy):
        # update self to be similar to the input `policy`.
        pass

    @property
    def time_step_spec(self):
        return self._time_step_spec

    @property
    def action_spec(self):
        return self._action_spec

    @property
    def policy_state_spec(self):
        return self._policy_state_spec


# ## Random Python Policy
# RandomPyPolicy generates random actions for the discrete/continuous given action_spec.
action_spec = array_spec.BoundedArraySpec((2,), np.int32, -10, 10)
my_random_py_policy = random_py_policy.RandomPyPolicy(
    time_step_spec=None, action_spec=action_spec
)
time_step = None
action_step = my_random_py_policy.action(time_step)
print(action_step)
action_step = my_random_py_policy.action(time_step)
print(action_step)


# ## Drivers
# Drivers are abstraction for the process of executing a policy in an environment for a specified number of steps  during data collection, evaluation and generating a video of the agent.The data encountered by the driver at each step like observation , action , reward , current and next step is saved in Trajectory and broadcast to a set of observers such as replay buffers and metrics.
# Implementations for drivers are available both in Python and TensorFlow
# **Python Drivers :**
# The PyDriver class takes a python environment, a python policy and a list of observers to update at each step.
class PyDriver(object):
    def __init__(self, env, policy, observers, max_steps=1, max_episodes=1):
        self._env = env
        self._policy = policy
        self._observers = observers or []
        self._max_steps = max_steps or np.inf
        self._max_episodes = max_episodes or np.inf

    def run(self, time_step, policy_state=()):
        num_steps = 0
        num_episodes = 0
        while num_steps < self._max_steps and num_episodes < self._max_episodes:
            # Compute an action using the policy for the given time_step
            action_step = self._policy.action(time_step, policy_state)
            # Apply the action to the environment and get the next step
            next_time_step = self._env.step(action_step.action)
            # Package information into a trajectory
            traj = trajectory.Trajectory(
                time_step.step_type,
                time_step.observation,
                action_step.action,
                action_step.info,
                next_time_step.step_type,
                next_time_step.reward,
                next_time_step.discount,
            )
            for observer in self._observers:
                observer(traj)
            # Update statistics to check termination
            num_episodes += np.sum(traj.is_last())
            num_steps += np.sum(~traj.is_boundary())
            time_step = next_time_step
            policy_state = action_step.state
        return time_step, policy_state


# The code below runs a random policy on the CartPole environment, saving the results to a **replay buffer**.
#
env = suite_gym.load("CartPole-v0")
policy = random_py_policy.RandomPyPolicy(
    time_step_spec=env.time_step_spec(), action_spec=env.action_spec()
)
replay_buffer = []
metric = py_metrics.AverageReturnMetric()
observers = [replay_buffer.append, metric]
driver = py_driver.PyDriver(env, policy, observers, max_steps=20, max_episodes=1)
initial_time_step = env.reset()
final_time_step, _ = driver.run(initial_time_step)
print("Replay Buffer:")
for traj in replay_buffer:
    print(traj)
print("Average Return: ", metric.result())
<|endoftext|>
[{"Id": 48305609, "ScriptId": 7934748, "ParentScriptVersionId": NaN, "ScriptLanguageId": 9, "AuthorUserId": 373959, "CreationDate": "12/01/2020 20:53:25", "VersionNumber": 43.0, "Title": "TFAgents-Environment , Policy , Driver", "EvaluationDate": "12/01/2020", "IsChange": false, "TotalLines": 320.0, "LinesInsertedFromPrevious": 0.0, "LinesChangedFromPrevious": 0.0, "LinesUnchangedFromPrevious": 320.0, "LinesInsertedFromFork": NaN, "LinesDeletedFromFork": NaN, "LinesChangedFromFork": NaN, "LinesUnchangedFromFork": NaN, "TotalVotes": 26}]
# This is a two part tutorial series on TF-Agents
# Part 1 : https://www.kaggle.com/usharengaraju/tfagents-environment-policy-driver
# Part 2 : https://www.kaggle.com/usharengaraju/tf-agents-replay-buffer-network-checkpointer
# Credit : The article series has been adapted from the official tensorflow documentation.
# ## Reinforcement Learning
# Reinforcement learning (RL) is a machine learning framework where agents takes action in an environment in order to maximize the cumulative reward .
# ![](https://drive.google.com/uc?id=1ayosPHB3DIbmgnuu0kkN-dpA1TfnurNo)
# Pic Credit : ai.googleblog.com
# ### TF-Agents :
# A reliable, scalable and easy to use TensorFlow library for Contextual Bandits and Reinforcement Learning.
# ![](https://drive.google.com/uc?id=1zjzi0OY7G-QusVpYppjdVvbAP-pX0-RZ)
# Pic Credit : Tensorflow Dev Summit
# ### Advantages of TF-Agents :
# 🎯 Great Resources to get started with Reinforcement Learning like colab , documentation and resources
# 🎯 Well suited to handle complex RL problems
# 🎯 Helps in developing RL algorithms quickly
# 🎯 Can be configured easily with gin-config
# ![](https://drive.google.com/uc?id=1QBCM4TbPlaoftfpXZVoQ7uUScvDtqga0)
# Pic Credit : Tensorflow Dev Summit
# In this notebook , we will be discussing three concepts - Environment , Policy and Driver
# ### Environment
# The agent receives an observation and chooses an action which it applies on the environment and gets back reward
# and a observation from the environment. The goal of the agent is to train a policy to choose actions which will
# maximize the cumulative rewards.
# TF-Agents has both Python and TensorFlow implentation of environments .
#
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import abc
import tensorflow as tf
import numpy as np
from tf_agents.environments import py_environment
from tf_agents.environments import tf_environment
from tf_agents.environments import tf_py_environment
from tf_agents.environments import utils
from tf_agents.specs import array_spec
from tf_agents.environments import wrappers
from tf_agents.environments import suite_gym
from tf_agents.trajectories import time_step as ts
from tf_agents.specs import array_spec
from tf_agents.specs import tensor_spec
from tf_agents.networks import network
from tf_agents.policies import py_policy
from tf_agents.policies import random_py_policy
from tf_agents.policies import scripted_py_policy
from tf_agents.policies import tf_policy
from tf_agents.policies import random_tf_policy
from tf_agents.policies import actor_policy
from tf_agents.policies import q_policy
from tf_agents.policies import greedy_policy
from tf_agents.environments import suite_gym
from tf_agents.metrics import py_metrics
from tf_agents.metrics import tf_metrics
from tf_agents.drivers import py_driver
from tf_agents.drivers import dynamic_episode_driver

tf.compat.v1.enable_v2_behavior()
from tf_agents.trajectories import time_step as ts

tf.compat.v1.enable_v2_behavior()
tf.compat.v1.enable_v2_behavior()


# ### Python Environments :
# The interface that all python environments must implement is in environments/py_environment.PyEnvironment.
# After interaction with the environment the following information about the next step are available
# 📌 **observation:** This is the part of the environment state that the agent can observe to choose its actions at the next step.
# 📌 **reward:** The agent is learning to maximize the sum of these rewards across multiple steps.
# 📌 **step_type:** Interactions with the environment are usually part of a sequence/episode. e.g. multiple moves in a game of chess. step_type can be either FIRST, MID or LAST to indicate whether this time step is the first, intermediate or last step in a sequence.
# 📌 **discount:** This is a float representing how much to weight the reward at the next time step relative to the reward at the current time step.
#
class PyEnvironment(object):
    def reset(self):
        """Return initial_time_step."""
        self._current_time_step = self._reset()
        return self._current_time_step

    def step(self, action):
        """Apply action and return new time_step."""
        if self._current_time_step is None:
            return self.reset()
        self._current_time_step = self._step(action)
        return self._current_time_step

    def current_time_step(self):
        return self._current_time_step

    def time_step_spec(self):
        """Return time_step_spec."""

    @abc.abstractmethod
    def observation_spec(self):
        """Return observation_spec."""

    @abc.abstractmethod
    def action_spec(self):
        """Return action_spec."""

    @abc.abstractmethod
    def _reset(self):
        """Return initial_time_step."""

    @abc.abstractmethod
    def _step(self, action):
        """Apply action and return new time_step."""
        self._current_time_step = self._step(action)
        return self._current_time_step


# ### Standard Environments :
# TF Agents has built-in wrappers for many standard environments like the OpenAI Gym, DeepMind-control and Atari, which can be easily loaded using our environment suites. The code below loads the CartPole environment from the OpenAI gym
#
environment = suite_gym.load("CartPole-v0")
print("action_spec:", environment.action_spec())
print("time_step_spec.observation:", environment.time_step_spec().observation)
print("time_step_spec.step_type:", environment.time_step_spec().step_type)
print("time_step_spec.discount:", environment.time_step_spec().discount)
print("time_step_spec.reward:", environment.time_step_spec().reward)


# ## Policy
# Policies map an observation from the environment to an action or a distribution over actions.Most policies have a neural network to compute actions and/or distributions over actions from TimeSteps.Policies can be saved/restored, and can be used indepedently of the agent for data collection, evaluation etc.
# Policies contain the following information
# 📌 **action:** The action to be applied to the environment.
# 📌 **state:** The state of the policy (e.g. RNN state) to be fed into the next call to action.
# 📌 **info:** Optional side information such as action log probabilities.
#
class Base(object):
    @abc.abstractmethod
    def __init__(self, time_step_spec, action_spec, policy_state_spec=()):
        self._time_step_spec = time_step_spec
        self._action_spec = action_spec
        self._policy_state_spec = policy_state_spec

    @abc.abstractmethod
    def reset(self, policy_state=()):
        # return initial_policy_state.
        pass

    @abc.abstractmethod
    def action(self, time_step, policy_state=()):
        # return a PolicyStep(action, state, info) named tuple.
        pass

    @abc.abstractmethod
    def distribution(self, time_step, policy_state=()):
        # Not implemented in python, only for TF policies.
        pass

    @abc.abstractmethod
    def update(self, policy):
        # update self to be similar to the input `policy`.
        pass

    @property
    def time_step_spec(self):
        return self._time_step_spec

    @property
    def action_spec(self):
        return self._action_spec

    @property
    def policy_state_spec(self):
        return self._policy_state_spec


# ## Random Python Policy
# RandomPyPolicy generates random actions for the discrete/continuous given action_spec.
action_spec = array_spec.BoundedArraySpec((2,), np.int32, -10, 10)
my_random_py_policy = random_py_policy.RandomPyPolicy(
    time_step_spec=None, action_spec=action_spec
)
time_step = None
action_step = my_random_py_policy.action(time_step)
print(action_step)
action_step = my_random_py_policy.action(time_step)
print(action_step)


# ## Drivers
# Drivers are abstraction for the process of executing a policy in an environment for a specified number of steps  during data collection, evaluation and generating a video of the agent.The data encountered by the driver at each step like observation , action , reward , current and next step is saved in Trajectory and broadcast to a set of observers such as replay buffers and metrics.
# Implementations for drivers are available both in Python and TensorFlow
# **Python Drivers :**
# The PyDriver class takes a python environment, a python policy and a list of observers to update at each step.
class PyDriver(object):
    def __init__(self, env, policy, observers, max_steps=1, max_episodes=1):
        self._env = env
        self._policy = policy
        self._observers = observers or []
        self._max_steps = max_steps or np.inf
        self._max_episodes = max_episodes or np.inf

    def run(self, time_step, policy_state=()):
        num_steps = 0
        num_episodes = 0
        while num_steps < self._max_steps and num_episodes < self._max_episodes:
            # Compute an action using the policy for the given time_step
            action_step = self._policy.action(time_step, policy_state)
            # Apply the action to the environment and get the next step
            next_time_step = self._env.step(action_step.action)
            # Package information into a trajectory
            traj = trajectory.Trajectory(
                time_step.step_type,
                time_step.observation,
                action_step.action,
                action_step.info,
                next_time_step.step_type,
                next_time_step.reward,
                next_time_step.discount,
            )
            for observer in self._observers:
                observer(traj)
            # Update statistics to check termination
            num_episodes += np.sum(traj.is_last())
            num_steps += np.sum(~traj.is_boundary())
            time_step = next_time_step
            policy_state = action_step.state
        return time_step, policy_state


# The code below runs a random policy on the CartPole environment, saving the results to a **replay buffer**.
#
env = suite_gym.load("CartPole-v0")
policy = random_py_policy.RandomPyPolicy(
    time_step_spec=env.time_step_spec(), action_spec=env.action_spec()
)
replay_buffer = []
metric = py_metrics.AverageReturnMetric()
observers = [replay_buffer.append, metric]
driver = py_driver.PyDriver(env, policy, observers, max_steps=20, max_episodes=1)
initial_time_step = env.reset()
final_time_step, _ = driver.run(initial_time_step)
print("Replay Buffer:")
for traj in replay_buffer:
    print(traj)
print("Average Return: ", metric.result())
