# AUTOGENERATED! DO NOT EDIT! File to edit: 02_pokingbot_training.ipynb (unless otherwise specified).

__all__ = ['PolicyNetwork', 'CriticNetwork', 'Trainer']

# Cell
from typing import Mapping, Sequence
from absl import app
from absl import flags
import acme
from acme import specs
from acme import types
from acme import wrappers
from acme.agents.tf import actors
from acme.agents.tf import d4pg
from acme.tf import networks
from acme.tf import utils as tf2_utils
from acme.tf import utils as tf2_utils
from acme.tf import networks
from acme import types
import dm_env
import gym
import numpy as np
import sonnet as snt
import sonnet as snt
import tensorflow as tf

# Cell
class PolicyNetwork(snt.Module):
    def __init__(self, action_spec):
        super(PolicyNetwork, self).__init__()

        # Get total number of action dimensions from action spec.
        output_dim = np.prod(action_spec.shape, dtype=int)

        self.conv = snt.Sequential([
            snt.Reshape([-1, 4]),
            snt.Conv1D(output_channels=32, kernel_shape=3, stride=1),
            tf.nn.relu,
            snt.Conv1D(output_channels=32, kernel_shape=3, stride=1),
            tf.nn.relu,
            snt.Flatten(),
        ])

        self.hidden = snt.Sequential([
            snt.Linear(512, w_init=tf.initializers.VarianceScaling(distribution='uniform', mode='fan_out', scale=0.333) ),
            tf.nn.elu,
            snt.Linear(512, w_init=tf.initializers.VarianceScaling(distribution='uniform', mode='fan_out', scale=0.333) ),
            tf.nn.elu,
            snt.Linear(256, w_init=tf.initializers.VarianceScaling(distribution='uniform', mode='fan_out', scale=0.333) ),
            tf.nn.elu,
            snt.Linear(256, w_init=tf.initializers.VarianceScaling(distribution='uniform', mode='fan_out', scale=0.333) ),
            tf.nn.elu,
        ])

        self.out = snt.Sequential([
            snt.Linear(
                output_dim, w_init=tf.initializers.VarianceScaling(1e-4)),
            networks.TanhToSpec(action_spec),
        ])

    def __call__(self, x):
        x = tf2_utils.batch_concat(x)

        laser, g= tf.split(x, [4*(241*2), 10*3], axis=1)
        laser = self.conv(laser)

        x = tf.concat([g, laser], 1)
        x = self.hidden(x)
        x = self.out(x)

        return x


class CriticNetwork(snt.Module):
    def __init__(self, vmin, vmax, num_atoms):
        super(CriticNetwork, self).__init__()

        self.conv = snt.Sequential([
            snt.Reshape([-1, 4]),
            snt.Conv1D(output_channels=32, kernel_shape=3, stride=1),
            tf.nn.relu,
            snt.Conv1D(output_channels=32, kernel_shape=3, stride=1),
            tf.nn.relu,
            snt.Flatten(),
        ])

        self.hidden = snt.Sequential([
            snt.Linear(512, w_init=tf.initializers.VarianceScaling(distribution='uniform', mode='fan_out', scale=0.333) ),
            tf.nn.elu,
            snt.Linear(512, w_init=tf.initializers.VarianceScaling(distribution='uniform', mode='fan_out', scale=0.333) ),
            tf.nn.elu,
            snt.Linear(256, w_init=tf.initializers.VarianceScaling(distribution='uniform', mode='fan_out', scale=0.333) ),
            tf.nn.elu,
            snt.Linear(256, w_init=tf.initializers.VarianceScaling(distribution='uniform', mode='fan_out', scale=0.333) ),
            tf.nn.elu,
        ])

        self.head = networks.DiscreteValuedHead(vmin, vmax, num_atoms)

    def __call__(self,
                 observation: types.NestedTensor,
                 action: types.NestedTensor) -> tf.Tensor:

        if hasattr(observation, 'dtype') and hasattr(action, 'dtype'):
            if observation.dtype != action.dtype:
                # Observation and action must be the same type for concat to work
                action = tf.cast(action, observation.dtype)

        x = tf2_utils.batch_concat(observation)
        action = tf2_utils.batch_concat(action)

        laser, g= tf.split(x, [4*(241*2), 10*3], axis=1)
        laser = self.conv(laser)

        x = tf.concat([g, laser, action], 1)
        x = self.hidden(x)
        x = self.head(x)

        return x


# Cell

class Trainer():
	def __init__(self, gym_id, pretrained=None):
		"""
		gym_id (String): The custom OpenAI gym environment have been registered. i.e. 'gym_poking:door-and-ball-v3'

		pretrained (String): While using a pretrained weight, the path of the specific weight, i.e.

								 log_home = '/home/argsubt/pokingbot-rl/acme_logs/'
								 log_id = '355a222e-7e90-11ec-9e19-04d4c401c495/'
								 snapshots_path = 'snapshots/policy'
								 checkpoints_path = 'checkpoints/d4pg_learner/'
								 ckpt_id = 'ckpt-409'
								 pretrained_weight = log_home + log_id + checkpoints_path + ckpt_id

		"""
		self.gym_id = gym_id
		self.specified_path = pretrained
	def __make_environment(self,
							task: str = 'gym_poking:curl-navi-gym-v0') -> dm_env.Environment:
		"""
		Creates an OpenAI Gym environment.
		Args:
			gym_id (String): The custom OpenAI gym environment have been registered.
		Returns:
			environment (env class)
		"""
		environment = gym.make(task)
		environment = wrappers.GymWrapper(environment)
		environment = wrappers.SinglePrecisionWrapper(environment)

		return environment

	def __make_networks(self,
						action_spec: specs.BoundedArray,
						vmin: float = -40.,
						vmax: float = 130.,
						num_atoms: int = 51,
						) -> Mapping[str, types.TensorTransformation]:
		"""
		Creates the networks used by the agent.
		vmin/vmax and num_atoms hyperparameters should be set to
		give the distributional critic a good dynamic range over possible discounted
		returns. Note that this is very different than the scale of immediate rewards.

		The vmin|vmax hyperparameters of the distributional critic may need tuning
		depending on your environment's rewards. A good rule of thumb is to
		set vmax to the discounted sum of the maximum instantaneous rewards
		for the maximum episode length;
		then set vmin to -vmax.

		Args:
			action_spec (specs.BoundedArray)
			vmin (float)
			vmax (float)
			num_atoms (int)
		Returns:
			policy_network (snt.Module)
			critic_network (snt.Module)
			observation_network (tf2_utils.batch_concat)
		"""

		observation_network = tf2_utils.batch_concat
		policy_network = PolicyNetwork(action_spec)
		critic_network = CriticNetwork(vmin, vmax, num_atoms)

		return {
			'policy': policy_network,
			'critic': critic_network,
			'observation': observation_network,
		}
	def main(self):
		"""
		Start running the train loop.
		"""
		environment = self.__make_environment(self.gym_id)
		environment_spec = specs.make_environment_spec(environment)
		agent_networks = self.__make_networks(environment_spec.actions)

		# change batch_size together with samples_per_insert
		# batch_size/samples_per_insert = 8
		# depends on the performance of workstation
		if self.specified_path == None:
			agent = d4pg.D4PG(
				environment_spec=environment_spec,
		        policy_network=agent_networks['policy'],
		        critic_network=agent_networks['critic'],
		        observation_network=agent_networks['observation'],
		        discount=0.99,
		        batch_size=512,
		        prefetch_size=4,
		        target_update_period=100,
		        min_replay_size=10000,
		        max_replay_size=200000,
		        # max_replay_size=1000000,
		        samples_per_insert=64.0,
		        n_step=5,
		        sigma=1.0,
		        clipping=True
			)
		else:
			agent = d4pg.D4PG(
				environment_spec=environment_spec,
		        policy_network=agent_networks['policy'],
		        critic_network=agent_networks['critic'],
		        observation_network=agent_networks['observation'],
		        discount=0.99,
		        batch_size=512,
		        prefetch_size=4,
		        target_update_period=100,
		        min_replay_size=10000,
		        max_replay_size=200000,
		        # max_replay_size=1000000,
		        samples_per_insert=64.0,
		        n_step=5,
		        sigma=1.0,
		        clipping=True,
				specified_path=self.specified_path
			)


		train_loop = acme.EnvironmentLoop(environment, agent, label='train_loop')
		train_loop.run(num_episodes=350000000)