<SYSTEM_TASK:>
Given the following text description, write Python code to implement the functionality described below step by step
<END_TASK>
<USER_TASK:>
Description:
Step5:  Contents
1. Bernoulli Bandit
Bonus 1.1. Gittins index (5 points)
HW 1.1. Nonstationary Bernoulli bandit (2 points)
2. Exploration in MDP
2.1 Epsilon-greedy q-learning
2.2 Reward shaping
2.3 Curiosity-driven Exploration
HW 2.1 Random network distillation (3 points)
Part 1. Bernoulli Bandit
We are going to implement several exploration strategies for simplest problem - bernoulli bandit.
The bandit has $K$ actions. Action produce 1.0 reward $r$ with probability $0 \le \theta_k \le 1$ which is unknown to agent, but fixed over time. Agent's objective is to minimize regret over fixed number $T$ of action selections
Step6:  Epsilon-greedy agent
for $t = 1,2,...$ do
&nbsp;&nbsp; for $k = 1,...,K$ do
&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp; $\hat\theta_k \leftarrow \alpha_k / (\alpha_k + \beta_k)$
&nbsp;&nbsp; end for 
&nbsp;&nbsp; $x_t \leftarrow argmax_{k}\hat\theta$ with probability $1 - \epsilon$ or random action with probability $\epsilon$
&nbsp;&nbsp; Apply $x_t$ and observe $r_t$
&nbsp;&nbsp; $(\alpha_{x_t}, \beta_{x_t}) \leftarrow (\alpha_{x_t}, \beta_{x_t}) + (r_t, 1-r_t)$
end for
Implement the algorithm above in the cell below
Step7:  UCB Agent
Epsilon-greedy strategy have no preference for actions. It would be better to select among actions that are uncertain or have potential to be optimal. One can come up with idea of index for each action that represents optimality and uncertainty at the same time. One efficient way to do it is to use UCB1 algorithm
Step8:  Thompson sampling
UCB1 algorithm does not take into account actual distribution of rewards. If we know the distribution - we can do much better by using Thompson sampling
Step10:  Bonus 1.1. Gittins index (5 points).
Bernoulli bandit problem has an optimal solution - Gittins index algorithm. Implement finite horizon version of the algorithm and demonstrate it's performance with experiments. Some articles
Step11:  And a picture how it's reward probabilities change over time
Step12:  Your task is to invent an agent that will have better regret than stationary agents from above.
Step13:  Part 2. Exploration in MDP
The following problem, called "deep see", illustrates importance of exploration in context of mdp's.
<img src="deep_see.png">
The deep sea problem is implemented as an $N×N$ grid with a one-hot encoding for state.
The agent begins each episode in the top left corner of the grid and descends one row
per timestep. Each episode terminates after N steps, when the agent reaches the bottom
row. In each state there is a random but fixed mapping between actions $A = {0,1}$ and
the transitions ‘left’ and ‘right’. At each timestep there is a small cost $r = −0.01/N$ of
moving right, and $r = 0$ for moving left. However, should the agent transition right at every
timestep of the episode it will be rewarded with an additional reward of $+1$.
Question
Step15:  2.1 Epsilon-greedy q-learning
Step16:  Let's try to solve this by q-learning with high epsilon!
Step17:  But if we do bigger env
Step18:  2.2 Reward shaping
Step19:  2.3 Curiosity-driven Exploration
Step20:  2.3.1 Uncertainty with forward dynamics
Step21:  2.3.2 Uncertainty with reverse dynamics
The paper
Step22:  2.3.3 Intrinsic Curiosity Module algorithm
Step23:  HW 2.1
<ASSISTANT_TASK:>
Python Code:
class BernoulliBandit:
    def __init__(self, n_actions=5):
        self._probs = np.random.random(n_actions)
    @property
    def action_count(self):
        return len(self._probs)
    def pull(self, action):
        if np.any(np.random.random() > self._probs[action]):
            return 0.0
        return 1.0
    def optimal_reward(self):
         Used for regret calculation
        
        return np.max(self._probs)
    def step(self):
         Used in nonstationary version
        
        pass
    def reset(self):
         Used in nonstationary version
        
class AbstractAgent(metaclass=ABCMeta):
    def init_actions(self, n_actions):
        self._successes = np.zeros(n_actions)
        self._failures = np.zeros(n_actions)
        self._total_pulls = 0
    @abstractmethod
    def get_action(self):
        
        Get current best action
        :rtype: int
        
        pass
    def update(self, action, reward):
        
        Observe reward from action and update agent's internal parameters
        :type action: int
        :type reward: int
        
        self._total_pulls += 1
        if reward == 1:
            self._successes[action] += 1
        else:
            self._failures[action] += 1
    @property
    def name(self):
        return self.__class__.__name__
class RandomAgent(AbstractAgent):
    def get_action(self):
        return np.random.randint(0, len(self._successes))
Explanation: Contents
1. Bernoulli Bandit
Bonus 1.1. Gittins index (5 points)
HW 1.1. Nonstationary Bernoulli bandit (2 points)
2. Exploration in MDP
2.1 Epsilon-greedy q-learning
2.2 Reward shaping
2.3 Curiosity-driven Exploration
HW 2.1 Random network distillation (3 points)
Part 1. Bernoulli Bandit
We are going to implement several exploration strategies for simplest problem - bernoulli bandit.
The bandit has $K$ actions. Action produce 1.0 reward $r$ with probability $0 \le \theta_k \le 1$ which is unknown to agent, but fixed over time. Agent's objective is to minimize regret over fixed number $T$ of action selections:
$$\rho = T\theta^* - \sum_{t=1}^T r_t$$
Where $\theta^* = \max_k{\theta_k}$
Real-world analogy:
Clinical trials - we have $K$ pills and $T$ ill patient. After taking pill, patient is cured with probability $\theta_k$. Task is to find most efficient pill.
A research on clinical trials - https://arxiv.org/pdf/1507.08025.pdf
End of explanation
class EpsilonGreedyAgent(AbstractAgent):
    def __init__(self, epsilon=0.01):
        self._epsilon = epsilon
    def get_action(self):
        if np.random.random() < self._epsilon:
            return np.random.randint(len(self._successes))
        else:
            return np.argmax(self._successes / (self._successes + self._failures + 0.1))
    @property
    def name(self):
        return self.__class__.__name__ + "(epsilon={})".format(self._epsilon)
Explanation: Epsilon-greedy agent
for $t = 1,2,...$ do
&nbsp;&nbsp; for $k = 1,...,K$ do
&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp; $\hat\theta_k \leftarrow \alpha_k / (\alpha_k + \beta_k)$
&nbsp;&nbsp; end for 
&nbsp;&nbsp; $x_t \leftarrow argmax_{k}\hat\theta$ with probability $1 - \epsilon$ or random action with probability $\epsilon$
&nbsp;&nbsp; Apply $x_t$ and observe $r_t$
&nbsp;&nbsp; $(\alpha_{x_t}, \beta_{x_t}) \leftarrow (\alpha_{x_t}, \beta_{x_t}) + (r_t, 1-r_t)$
end for
Implement the algorithm above in the cell below:
End of explanation
class UCBAgent(AbstractAgent):
    def get_action(self):
        pulls = self._successes + self._failures + 0.1
        return np.argmax(self._successes / pulls + np.sqrt(2 * np.log(self._total_pulls + 0.1) / pulls))
Explanation: UCB Agent
Epsilon-greedy strategy have no preference for actions. It would be better to select among actions that are uncertain or have potential to be optimal. One can come up with idea of index for each action that represents optimality and uncertainty at the same time. One efficient way to do it is to use UCB1 algorithm:
for $t = 1,2,...$ do
&nbsp;&nbsp; for $k = 1,...,K$ do
&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp; $w_k \leftarrow \alpha_k / (\alpha_k + \beta_k) + \sqrt{2\log(t) \ / \ (\alpha_k + \beta_k)}$
&nbsp;&nbsp; end for 
&nbsp;&nbsp; end for 
 $x_t \leftarrow argmax_{k}w$
&nbsp;&nbsp; Apply $x_t$ and observe $r_t$
&nbsp;&nbsp; $(\alpha_{x_t}, \beta_{x_t}) \leftarrow (\alpha_{x_t}, \beta_{x_t}) + (r_t, 1-r_t)$
end for
Note: in practice, one can multiply $\sqrt{2\log(t) \ / \ (\alpha_k + \beta_k)}$ by some tunable parameter to regulate agent's optimism and wilingness to abandon non-promising actions.
More versions and optimality analysis - https://homes.di.unimi.it/~cesabian/Pubblicazioni/ml-02.pdf
End of explanation
class ThompsonSamplingAgent(AbstractAgent):
    def get_action(self):
        return np.argmax(np.random.beta(self._successes + 1, self._failures + 1))
from collections import OrderedDict
def get_regret(env, agents, n_steps=5000, n_trials=50):
    scores = OrderedDict({
        agent.name: [0.0 for step in range(n_steps)] for agent in agents
    })
    for trial in range(n_trials):
        env.reset()
        for a in agents:
            a.init_actions(env.action_count)
        for i in range(n_steps):
            optimal_reward = env.optimal_reward()
            for agent in agents:
                action = agent.get_action()
                reward = env.pull(action)
                agent.update(action, reward)
                scores[agent.name][i] += optimal_reward - reward
            env.step()  # change bandit's state if it is unstationary
    for agent in agents:
        scores[agent.name] = np.cumsum(scores[agent.name]) / n_trials
    return scores
def plot_regret(agents, scores):
    for agent in agents:
        plt.plot(scores[agent.name])
    plt.legend([agent.name for agent in agents])
    plt.ylabel("regret")
    plt.xlabel("steps")
    plt.show()
# Uncomment agents
agents = [
    EpsilonGreedyAgent(),
    UCBAgent(),
    ThompsonSamplingAgent()
]
regret = get_regret(BernoulliBandit(), agents, n_steps=10000, n_trials=10)
plot_regret(agents, regret)
Explanation: Thompson sampling
UCB1 algorithm does not take into account actual distribution of rewards. If we know the distribution - we can do much better by using Thompson sampling:
for $t = 1,2,...$ do
&nbsp;&nbsp; for $k = 1,...,K$ do
&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp; Sample $\hat\theta_k \sim beta(\alpha_k, \beta_k)$
&nbsp;&nbsp; end for 
&nbsp;&nbsp; $x_t \leftarrow argmax_{k}\hat\theta$
&nbsp;&nbsp; Apply $x_t$ and observe $r_t$
&nbsp;&nbsp; $(\alpha_{x_t}, \beta_{x_t}) \leftarrow (\alpha_{x_t}, \beta_{x_t}) + (r_t, 1-r_t)$
end for
More on Thompson Sampling:
https://web.stanford.edu/~bvr/pubs/TS_Tutorial.pdf
End of explanation
class DriftingBandit(BernoulliBandit):
    def __init__(self, n_actions=5, gamma=0.01):
        
        Idea from https://github.com/iosband/ts_tutorial
        
        super().__init__(n_actions)
        self._gamma = gamma
        self._successes = None
        self._failures = None
        self._steps = 0
        self.reset()
    def reset(self):
        self._successes = np.zeros(self.action_count) + 1.0
        self._failures = np.zeros(self.action_count) + 1.0
        self._steps = 0
    def step(self):
        action = np.random.randint(self.action_count)
        reward = self.pull(action)
        self._step(action, reward)
    def _step(self, action, reward):
        self._successes = self._successes * (1 - self._gamma) + self._gamma
        self._failures = self._failures * (1 - self._gamma) + self._gamma
        self._steps += 1
        self._successes[action] += reward
        self._failures[action] += 1.0 - reward
        self._probs = np.random.beta(self._successes, self._failures)
Explanation: Bonus 1.1. Gittins index (5 points).
Bernoulli bandit problem has an optimal solution - Gittins index algorithm. Implement finite horizon version of the algorithm and demonstrate it's performance with experiments. Some articles:
- Wikipedia article - https://en.wikipedia.org/wiki/Gittins_index
- Different algorithms for index computation - http://www.ece.mcgill.ca/~amahaj1/projects/bandits/book/2013-bandit-computations.pdf (see "Bernoulli" section)
HW 1.1. Nonstationary Bernoulli bandit
What if success probabilities change over time? Here is an example of such bandit:
End of explanation
drifting_env = DriftingBandit(n_actions=5)
drifting_probs = []
for i in range(20000):
    drifting_env.step()
    drifting_probs.append(drifting_env._probs)
plt.figure(figsize=(17, 8))
plt.plot(pandas.DataFrame(drifting_probs).rolling(window=20).mean())
plt.xlabel("steps")
plt.ylabel("Success probability")
plt.title("Reward probabilities over time")
plt.legend(["Action {}".format(i) for i in range(drifting_env.action_count)])
plt.show()
Explanation: And a picture how it's reward probabilities change over time
End of explanation
# YOUR AGENT HERE SECTION
drifting_agents = [
    ThompsonSamplingAgent(),
    EpsilonGreedyAgent(),
    UCBAgent(),
#     YourAgent()
]
regret = get_regret(DriftingBandit(), drifting_agents, n_steps=20000, n_trials=10)
plot_regret(drifting_agents, regret)
Explanation: Your task is to invent an agent that will have better regret than stationary agents from above.
End of explanation
import gym
import numpy as np
import matplotlib.pyplot as plt
import torch
from torch import nn
from time import sleep
from tqdm import tqdm
from IPython.display import clear_output
import bsuite
from bsuite import sweep
from bsuite.utils import gym_wrapper
from q_learning_agent import QLearningAgent
from replay_buffer import ReplayBuffer
env_id = sweep.DEEP_SEA[0]
env = bsuite.load_from_id(env_id)
env = gym_wrapper.GymFromDMEnv(env)
sweep.SETTINGS[env_id]
Explanation: Part 2. Exploration in MDP
The following problem, called "deep see", illustrates importance of exploration in context of mdp's.
<img src="deep_see.png">
The deep sea problem is implemented as an $N×N$ grid with a one-hot encoding for state.
The agent begins each episode in the top left corner of the grid and descends one row
per timestep. Each episode terminates after N steps, when the agent reaches the bottom
row. In each state there is a random but fixed mapping between actions $A = {0,1}$ and
the transitions ‘left’ and ‘right’. At each timestep there is a small cost $r = −0.01/N$ of
moving right, and $r = 0$ for moving left. However, should the agent transition right at every
timestep of the episode it will be rewarded with an additional reward of $+1$.
Question: Why is the deep see a challengin exploration problem?
See full paper here
End of explanation
def get_state_number(s):
    return np.argmax(s.flatten())
def test_agent(agent, greedy=True, delay=.5):
    v = get_all_states_value(agent)
    s = env.reset()
    done = False
    while not done:
        fig, ax = plt.subplots(ncols=2)
        ax[0].imshow(env.render())
        ax[0].set_title('State')
        im = ax[1].imshow(v)
        plt.colorbar(im)
        ax[1].set_title('Value function')
        clear_output(True)
        plt.show()
        s = get_state_number(s)
        if greedy:
            a = agent.get_best_action(s)
        else:
            a = agent.get_action(s)
            
        s, r, done, _  = env.step(a)
        sleep(delay)
        
def get_all_states_value(agent):
    s_shape = env.observation_space.shape
    s_shape_flatten = np.prod(s_shape)
    v = np.zeros(s_shape_flatten)
    for i in range(s_shape_flatten):
        v[i] = agent.get_value(i)
    v = v.reshape(s_shape)
    return v
def to_one_hot(x, ndims):
     helper: take an integer vector and convert it to 1-hot matrix. 
    x = x.long().view(-1, 1)
    x = torch.zeros(
        x.size()[0], ndims).scatter_(1, x, 1)
    return x
agent = QLearningAgent(
    epsilon=1, 
    alpha=0.5, 
    discount=1, 
    get_legal_actions=lambda s: range(env.action_space.n)
)
test_agent(agent, greedy=True)
Explanation: 2.1 Epsilon-greedy q-learning
End of explanation
for i in range(5000):
    s = env.reset()
    done = False
    while not done:
        i_s = get_state_number(s)
        a = agent.get_action(i_s)
        s_next, r, done, _  = env.step(a)
        i_s_next = get_state_number(s_next)
        agent.update(i_s, a, r, i_s_next, done)
        s = s_next
test_agent(agent)
Explanation: Let's try to solve this by q-learning with high epsilon!
End of explanation
env_id = sweep.DEEP_SEA[1]
env = bsuite.load_from_id(env_id)
env = gym_wrapper.GymFromDMEnv(env)
sweep.SETTINGS[env_id]
agent = QLearningAgent(
    epsilon=1, 
    alpha=0.5, 
    discount=1, 
    get_legal_actions=lambda s: range(env.action_space.n)
)
for i in range(5000):
    s = env.reset()
    done = False
    while not done:
        i_s = get_state_number(s)
        a = agent.get_action(i_s)
        s_next, r, done, _  = env.step(a)
        i_s_next = get_state_number(s_next)
        agent.update(i_s, a, r, i_s_next, done)
        s = s_next
test_agent(agent)
Explanation: But if we do bigger env:
End of explanation
class BaseIntrinsicRewardModule(nn.Module):
    def __init__(self):
        super().__init__()
    def get_intrinsic_reward(self, state, action, next_state):
        return 0.0
    def get_loss(self, state_batch, action_batch, next_state_batch):
        pass
def train_with_reward(env, agent, reward_module, n_episodes=100, update_reward_period=100, batch_size=100, n_iter=10):
    buffer = ReplayBuffer(size=int(1e6))
    
    if list(reward_module.parameters()):
        optimizer = torch.optim.Adam(reward_module.parameters())
    else:
        optimizer = None
    losses = []
    s = env.reset()
    for i in range(n_episodes):
        done = False
        
        while not done:
            i_s = get_state_number(s)
            a = agent.get_action(i_s)
            s_next, r, done, _  = env.step(a)
            i_s_next = get_state_number(s_next)
            
            state_t = torch.tensor(s).float().view(1, -1)
            action_t = torch.tensor(a).float().view(1, -1)
            next_state_t = torch.tensor(s_next).float().view(1, -1)
            r_intr = reward_module.get_intrinsic_reward(state_t, action_t, next_state_t)
            r += r_intr
            agent.update(i_s, a, r, i_s_next, done)
            buffer.add(s, a, r, s_next, done)
            s = s_next
        if (i + 1) % update_reward_period == 0 and optimizer is not None:
        
            for _ in range(n_iter):
                optimizer.zero_grad()
                state_batch, action_batch, _, next_state_batch, _ = buffer.sample(batch_size)
                
                state_tensor = torch.tensor(state_batch).float().flatten(1, 2)
                action_tensor = torch.tensor(action_batch).float().view(-1, 1)
                next_state_tensor = torch.tensor(next_state_batch).float().flatten(1, 2)
                
                loss = reward_module.get_loss(state_tensor, action_tensor, next_state_tensor)
                loss.backward()
                optimizer.step()
                losses.append(loss.item())
            
            fig, ax = plt.subplots(ncols=2)
            ax[0].set_title('Value function after iter: %d' % i)
            im = ax[0].imshow(get_all_states_value(agent))
            ax[1].plot(losses)
            ax[1].set_title('Random network distillation loss')
            clear_output(True)
            plt.show()
class GoRightReward(BaseIntrinsicRewardModule):
    def __init__(self):
        super().__init__()
    def get_intrinsic_reward(self, state, action, next_state):
        # <YOUR CODE>
agent = QLearningAgent(
    epsilon=.1, 
    alpha=0.5, 
    discount=.9, 
    get_legal_actions=lambda s: range(env.action_space.n)
)
go_right = GoRightReward()
train_with_reward(env, agent, go_right, n_episodes=500)
test_agent(agent, greedy=True)
Explanation: 2.2 Reward shaping
End of explanation
class MLP(nn.Module):
    def __init__(self, input_size, hidden_size, output_size):
        super().__init__()
        self.layers = nn.Sequential(
            nn.Linear(input_size, hidden_size),
            nn.ReLU(),
            nn.Linear(hidden_size, output_size)
        )
        
        def init_weights(tensor):
            if isinstance(tensor, nn.Linear):
                nn.init.xavier_uniform_(tensor.weight)
        
        self.layers.apply(init_weights)
            
    
    def forward(self, x):
        return self.layers(x)
Explanation: 2.3 Curiosity-driven Exploration
End of explanation
class ForwardDynamics(BaseIntrinsicRewardModule):
    def __init__(self, states_size, actions_size, hidden_size, alpha=.1):
        super().__init__()
        self.module = MLP(
            actions_size + states_size,
            hidden_size,
            states_size
        )
        self.alpha = alpha
        self.mean_reward = 0
    
    def forward(self, s, a):
        sa = torch.cat([s, a], dim=-1)
        return s + self.module(sa)
    def get_intrinsic_reward(self, state, action, next_state):
        with torch.no_grad():
            r = # <YOUR CODE>
            r_centered = r - self.mean_reward
            self.mean_reward = self.alpha * (r) + (1 - self.alpha) * self.mean_reward
            return r_centered
    def get_loss(self, state_batch, action_batch, next_state_batch):
        # <YOUR CODE>
agent = QLearningAgent(
    epsilon=.1, 
    alpha=0.5, 
    discount=.9, 
    get_legal_actions=lambda s: range(env.action_space.n)
)
forward_dynamics = ForwardDynamics(
    np.prod(env.observation_space.shape), 
    1, 
    16
)
train_with_reward(env, agent, forward_dynamics, n_episodes=2000, update_reward_period=100, batch_size=100, n_iter=25)
test_agent(agent)
Explanation: 2.3.1 Uncertainty with forward dynamics
End of explanation
class InverseDynamics(BaseIntrinsicRewardModule):
    def __init__(self, states_size, n_actions, hidden_size, alpha=0.1):
        super().__init__()
        self.module = MLP(
            2 * states_size,
            hidden_size,
            n_actions
        )
        self.alpha = alpha
        self.mean_reward = 0
        self.n_actions = n_actions
    
    def forward(self, s, s_next):
        # <YOUR CODE>
    
    
    def get_intrinsic_reward(self, state, action, next_state):
        with torch.no_grad():
            r = # <YOUR CODE>
            
            r_centered = r - self.mean_reward
            self.mean_reward = self.alpha * (r) + (1 - self.alpha) * self.mean_reward
            return r_centered
    def get_loss(self, state_batch, action_batch, next_state_batch):        
        a_pred_proba = self.forward(state_batch, next_state_batch)
        a_one_hot = to_one_hot(action_batch, self.n_actions)
        return -(a_pred_proba * a_one_hot).sum(dim=-1).mean()
agent = QLearningAgent(
    epsilon=.1, 
    alpha=0.5, 
    discount=.9, 
    get_legal_actions=lambda s: range(env.action_space.n)
)
inverse_dynamics = InverseDynamics(
    np.prod(env.observation_space.shape), 
    env.action_space.n, 
    16
)
train_with_reward(env, agent, inverse_dynamics, n_episodes=3000, 
                  update_reward_period=100, batch_size=100, n_iter=25)
test_agent(agent)
Explanation: 2.3.2 Uncertainty with reverse dynamics
The paper
End of explanation
class Embedder(nn.Module):
    def __init__(self, states_size, embedding_size, hidden_size):
        super().__init__()
        self.module = MLP(
            states_size,
            hidden_size,
            embedding_size
        )
        
    def forward(self, s):
        return self.module(s)
        
class ICMModule(BaseIntrinsicRewardModule):
    def __init__(self, states_size, n_actions, hidden_size, embedding_size):
        super().__init__()
        # <YOUR CODE>
    
    def get_intrinsic_reward(self, state, action, next_state):
        with torch.no_grad():            
            # <YOUR CODE>
    def get_loss(self, state_batch, action_batch, next_state_batch):
        # <YOUR CODE>
agent = QLearningAgent(
    epsilon=.1, 
    alpha=0.5, 
    discount=1, 
    get_legal_actions=lambda s: range(env.action_space.n)
)
icm = ICMModule(
    states_size=np.prod(env.observation_space.shape), 
    n_actions=env.action_space.n, 
    hidden_size=16, embedding_size=10
)
train_with_reward(env, agent, icm, n_episodes=3000, update_reward_period=100, batch_size=100, n_iter=200)
test_agent(agent)
Explanation: 2.3.3 Intrinsic Curiosity Module algorithm
End of explanation
class RandomNetworkDistilationModule(BaseIntrinsicRewardModule):
    # <YOUR CODE>
agent = QLearningAgent(
    epsilon=.1, 
    alpha=0.5, 
    discount=.9, 
    get_legal_actions=lambda s: range(env.action_space.n)
)
rnd = RandomNetworkDistilationModule(
    np.prod(env.observation_space.shape), 
    np.prod(env.observation_space.shape), 
    16
)
train_with_reward(env, agent, rnd, n_episodes=2000, update_reward_period=100, batch_size=100, n_iter=25)
test_agent(agent)
Explanation: HW 2.1: Random network distillation
Implement algorithm from this paper
End of explanation
<END_TASK>