import os
import pickle
from datetime import datetime

import gym
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import tensorflow as tf

from laser_trajectory_control_env import LaserTrajectoryControlEnv
from tango_connection import TangoConnection

# set random seed
random_seed = 111
np.random.seed(random_seed)


conf_file = '/home/niky/FERMI/2020_10_06/configuration/conf_fel2.json'
tango = TangoConnection(conf_file=conf_file)
real_env = LaserTrajectoryControlEnv(tango=tango)

# Todo: Add info to buffer and test ensemble advantage
# We will take the avarage cumulative reward and the entropy of the trajectories as an observable,
# as well as the reward on the measured data during the data collection


# For Niky Hyper parameters:
#------------------------------------------------------
steps_per_env = 15
init_random_steps = 50
# num_epochs = int((300 - init_random_steps) / (steps_per_env)) + 1
num_epochs = 20
# num_epochs = int((125-init_random_steps)/(steps_per_env))
print('Number of epoches: ', num_epochs)

hidden_sizes = [100, 100]
cr_lr = 1e-3
gamma = 0.9999
lam = 0.95

max_training_iterations = 50

delay_before_convergence_check = 1

# number_envs = 1
critic_iter = 40
delta = 0.01
algorithm = 'TRPO'
conj_iters = 15
minibatch_size = 100
simulated_steps = 1000

mb_lr = 1e-3
model_batch_size = 5  # train the dynamics model

num_ensemble_models = 3
model_iter = 25

# How often to check the progress of the network training
# e.g. lambda it, episode: (it + 1) % max(3, (ep+1)*2) == 0
# dynamic_wait_time = lambda it, ep: True  # (it + 1) % 1 == 0  #
dynamic_wait_time = lambda it, ep: (it + 1) % 2 == 0  #

# Set the priors for the anchor method:
#  TODO: How to set these correctly?
init_params = dict(init_stddev_1_w=np.sqrt(1),
                   init_stddev_1_b=np.sqrt(1),
                   init_stddev_2_w=1 / np.sqrt(100))

data_noise = 1e-4  # estimated noise variance

lambda_anchor = data_noise / (np.array([init_params['init_stddev_1_w'],
                                        init_params['init_stddev_1_b'],
                                        init_params['init_stddev_2_w']]) ** 2)

training_steps_model = 4000

# method = 'random'
method = 'pessimistic'

# Create the logging directory:
project_directory = 'Data_logging/AE_TRPO/at_1/'

#------------------------------------------------------
#------------------------------------------------------
#------------------------------------------------------


hyp_str_all = '-nr_steps_' + str(steps_per_env) + '-cr_lr' + str(cr_lr) + '-crit_it_' + str(
    critic_iter) + '-d_' + str(delta) + '-conj_iters_' + str(conj_iters) + '-n_ep_' + str(num_epochs) + \
              '-mini_bs_' + str(minibatch_size) + '-m_bs_' + str(model_batch_size) + \
              '-mb_lr_' + str(mb_lr) + \
              '-sim_steps_' + str(simulated_steps) + \
              '-m_iter_' + str(model_iter) + '-ensnr_' + str(num_ensemble_models) + '-init_' + str(
    init_random_steps) + '/'
project_directory = project_directory + hyp_str_all

# To label the plots:
hyp_str_all = '-nr_steps_' + str(steps_per_env) + '-cr_lr' + str(cr_lr) + '-crit_it_' + str(
    critic_iter) + '-d_' + str(delta) + '-conj_iters_' + str(conj_iters) + '-n_ep_' + str(num_epochs) + \
              '\n-mini_bs_' + str(minibatch_size) + '-m_bs_' + str(model_batch_size) + \
              '-mb_lr_' + str(mb_lr) + \
              '-sim_steps_' + str(simulated_steps) + \
              '-m_iter_' + str(model_iter) + \
              '\n-ensnr_' + str(num_ensemble_models)
if not os.path.isdir(project_directory):
    os.makedirs(project_directory)
    print("created folder : ", project_directory)


# Class for data storage during the tests
class TrajectoryBuffer():
    '''Class for data storage during the tests'''

    def __init__(self, name, directory):
        self.save_frequency = 100000
        self.directory = directory
        self.name = name
        self.rews = []
        self.obss = []
        self.acts = []
        self.dones = []
        self.info = ""
        self.idx = -1

    def new_trajectory(self, obs):
        self.idx += 1
        self.rews.append([])
        self.acts.append([])
        self.obss.append([])
        self.dones.append([])
        self.store_step(obs=obs)

    def store_step(self, obs=None, act=None, rew=None, done=None):
        self.rews[self.idx].append(rew)
        self.obss[self.idx].append(obs)
        self.acts[self.idx].append(act)
        self.dones[self.idx].append(done)

        if self.__len__() % self.save_frequency == 0:
            self.save_buffer()

    def __len__(self):
        assert (len(self.rews) == len(self.obss) == len(self.acts) == len(self.dones))
        return len(self.obss)

    def save_buffer(self, **kwargs):
        if 'info' in kwargs:
            self.info = kwargs.get('info')
        now = datetime.now()
        # clock_time = "{}_{}_{}_{}_".format(now.day, now.hour, now.minute, now.second)
        clock_time = f'{now.month:0>2}_{now.day:0>2}_{now.hour:0>2}_{now.minute:0>2}_{now.second:0>2}_'
        data = dict(obss=self.obss,
                    acts=self.acts,
                    rews=self.rews,
                    dones=self.dones,
                    info=self.info)
        # print('saving...', data)
        out_put_writer = open(self.directory + clock_time + self.name, 'wb')
        pickle.dump(data, out_put_writer, -1)
        # pickle.dump(self.actions, out_put_writer, -1)
        out_put_writer.close()

    def get_data(self):
        return dict(obss=self.obss,
                    acts=self.acts,
                    rews=self.rews,
                    dones=self.dones,
                    info=self.info)


class MonitoringEnv(gym.Wrapper):
    '''
    Gym Wrapper to store information for scaling to correct scpace and for post analysis.
    '''

    def __init__(self, env, **kwargs):
        gym.Wrapper.__init__(self, env)
        self.data_dict = dict()
        self.environment_usage = 'default'
        self.directory = project_directory
        self.data_dict[self.environment_usage] = TrajectoryBuffer(name=self.environment_usage,
                                                                  directory=self.directory)
        self.current_buffer = self.data_dict.get(self.environment_usage)

        self.test_env_flag = False
        if 'test_env' in kwargs:
            self.test_env_flag = True

    def reset(self, **kwargs):
        init_obs = self.env.reset(**kwargs)
        # print('Reset Env: ', (init_obs),10*'-- ')
        self.current_buffer.new_trajectory(init_obs)
        init_obs = self.scale_state_env(init_obs)
        # print('Reset Menv: ', (init_obs))
        return init_obs

    def step(self, action):
        # print('a', action)
        action = self.descale_action_env(action)
        # print('as', action)
        ob, reward, done, info = self.env.step(action)
        # print('Env: ', reward)
        # print('Env: ', ob, 'r:', reward, done)
        self.current_buffer.store_step(obs=ob, act=action, rew=reward, done=done)
        ob = self.scale_state_env(ob)
        reward = self.rew_scale(reward)
        # if done:
        #     print('Menv: ', ob, 'r:', reward, done)
        # print('Menv: ', reward)
        return ob, reward, done, info

    def set_usage(self, usage):
        self.environment_usage = usage
        if usage in self.data_dict:
            self.current_buffer = self.data_dict.get(usage)
        else:
            self.data_dict[self.environment_usage] = TrajectoryBuffer(name=self.environment_usage,
                                                                      directory=self.directory)
            self.current_buffer = self.data_dict.get(usage)

    def close_usage(self, usage):
        # Todo: Implement to save complete data
        self.current_buffer = self.data_dict.get(usage)
        self.current_buffer.save_buffer()

    def scale_state_env(self, ob):
        scale = (self.env.observation_space.high - self.env.observation_space.low)
        return (2 * ob - (self.env.observation_space.high + self.env.observation_space.low)) / scale

    def descale_action_env(self, act):
        scale = (self.env.action_space.high - self.env.action_space.low)
        return (scale * act + self.env.action_space.high + self.env.action_space.low) / 2

    def rew_scale(self, rew):
        # we only scale for the network training:
        if not self.test_env_flag:
            rew = rew * 2 + 1
        # if self.test_env_flag:
        #     '''Rescale reward from [-1,0] to [-1,1] for the training of the network in case of tests'''
        #     rew = rew * 2 + 1
        #     if rew < -1:
        #         print('Hallo was geht: ', rew)
        #     else:
        #         print('Okay...', rew)
        return rew

    def save_current_buffer(self, info=''):
        self.current_buffer = self.data_dict.get(self.environment_usage)
        self.current_buffer.save_buffer(info=info)
        print('Saved current buffer', self.environment_usage)

    def set_directory(self, directory):
        self.directory = directory


env_monitored = MonitoringEnv(env=real_env)


def make_env(**kwargs):
    '''Create the environement'''
    return MonitoringEnv(env=real_env, **kwargs)


def mlp(x, hidden_layers, output_layer, activation, last_activation=None, init_params=dict()):
    '''
    Multi-layer perceptron with init conditions for anchor method
    '''
    if init_params:
        layer_1_w = tf.layers.Dense(hidden_layers[0],
                                    activation=tf.tanh,  # try using tf.nn.relu, tf.erf, tf.nn.tanh etc.
                                    kernel_initializer=tf.random_normal_initializer(mean=0.,
                                                                                    stddev=init_params.get(
                                                                                        'init_stddev_1_w')),
                                    bias_initializer=tf.random_normal_initializer(mean=0.,
                                                                                  stddev=init_params.get(
                                                                                      'init_stddev_1_b')))

        layer_1 = layer_1_w.apply(x)

        layer_2_w = tf.layers.Dense(hidden_layers[1],
                                    activation=tf.tanh,  # try using tf.nn.relu, tf.erf, tf.nn.tanh etc.
                                    kernel_initializer=tf.random_normal_initializer(mean=0.,
                                                                                    stddev=init_params.get(
                                                                                        'init_stddev_1_w')),
                                    bias_initializer=tf.random_normal_initializer(mean=0.,
                                                                                  stddev=init_params.get(
                                                                                      'init_stddev_1_b')))

        layer_2 = layer_2_w.apply(layer_1)

        output_w = tf.layers.Dense(output_layer,
                                   activation=tf.tanh,
                                   use_bias=False,
                                   kernel_initializer=tf.random_normal_initializer(mean=0.,
                                                                                   stddev=init_params.get(
                                                                                       'init_stddev_2_w')))

        return output_w.apply(layer_2)
    else:
        for l in hidden_layers:
            x = tf.layers.dense(x, units=l, activation=activation)
        return tf.layers.dense(x, units=output_layer, activation=last_activation)


def softmax_entropy(logits):
    '''
    Softmax Entropy
    '''
    return -tf.reduce_sum(tf.nn.softmax(logits, axis=-1) * tf.nn.log_softmax(logits, axis=-1), axis=-1)


def gaussian_log_likelihood(ac, mean, log_std):
    '''
    Gaussian Log Likelihood
    '''
    log_p = ((ac - mean) ** 2 / (tf.exp(log_std) ** 2 + 1e-9) + 2 * log_std) + np.log(2 * np.pi)
    return -0.5 * tf.reduce_sum(log_p, axis=-1)


def conjugate_gradient(A, b, x=None, iters=10):
    '''
    Conjugate gradient method: approximate the solution of Ax=b
    It solve Ax=b without forming the full matrix, just compute the matrix-vector product (The Fisher-vector product)
    NB: A is not the full matrix but is a useful matrix-vector product between the averaged Fisher information matrix and arbitrary vectors
    Descibed in Appendix C.1 of the TRPO paper
    '''
    if x is None:
        x = np.zeros_like(b)

    r = A(x) - b
    p = -r
    for _ in range(iters):
        a = np.dot(r, r) / (np.dot(p, A(p)) + 1e-8)
        x += a * p
        r_n = r + a * A(p)
        b = np.dot(r_n, r_n) / (np.dot(r, r) + 1e-8)
        p = -r_n + b * p
        r = r_n
    return x


def gaussian_DKL(mu_q, log_std_q, mu_p, log_std_p):
    '''
    Gaussian KL divergence in case of a diagonal covariance matrix
    '''
    return tf.reduce_mean(tf.reduce_sum(
        0.5 * (log_std_p - log_std_q + tf.exp(log_std_q - log_std_p) + (mu_q - mu_p) ** 2 / tf.exp(log_std_p) - 1),
        axis=1))


def backtracking_line_search(Dkl, delta, old_loss, p=0.8):
    '''
    Backtracking line searc. It look for a coefficient s.t. the constraint on the DKL is satisfied
    It has both to
     - improve the non-linear objective
     - satisfy the constraint
    '''
    ## Explained in Appendix C of the TRPO paper
    a = 1
    it = 0

    new_dkl, new_loss = Dkl(a)
    while (new_dkl > delta) or (new_loss > old_loss):
        a *= p
        it += 1
        new_dkl, new_loss = Dkl(a)

    return a


def GAE(rews, v, v_last, gamma=0.99, lam=0.95):
    '''
    Generalized Advantage Estimation
    '''
    assert len(rews) == len(v)
    vs = np.append(v, v_last)
    d = np.array(rews) + gamma * vs[1:] - vs[:-1]
    gae_advantage = discounted_rewards(d, 0, gamma * lam)
    return gae_advantage


def discounted_rewards(rews, last_sv, gamma):
    '''
    Discounted reward to go
    Parameters:
    ----------
    rews: list of rewards
    last_sv: value of the last state
    gamma: discount value
    '''
    rtg = np.zeros_like(rews, dtype=np.float64)
    rtg[-1] = rews[-1] + gamma * last_sv
    for i in reversed(range(len(rews) - 1)):
        rtg[i] = rews[i] + gamma * rtg[i + 1]
    return rtg


def flatten_list(tensor_list):
    '''
    Flatten a list of tensors
    '''
    return tf.concat([flatten(t) for t in tensor_list], axis=0)


def flatten(tensor):
    '''
    Flatten a tensor
    '''
    return tf.reshape(tensor, shape=(-1,))


def test_agent(env_test, agent_op, num_games=10):
    '''
    Test an agent 'agent_op', 'num_games' times
    Return mean and std
    '''
    games_r = []
    games_length = []
    games_dones = []
    for _ in range(num_games):
        d = False
        game_r = 0
        o = env_test.reset()
        game_length = 0
        while not d:
            a_s, _ = agent_op([o])
            o, r, d, _ = env_test.step(a_s)
            game_r += r
            game_length += 1

        games_r.append(game_r)
        games_length.append(game_length)
        games_dones.append(d)
    return np.mean(games_r), np.std(games_r), np.mean(games_length), np.mean(games_dones)


class Buffer():
    '''
    Class to store the experience from a unique policy
    '''

    def __init__(self, gamma=0.99, lam=0.95):
        self.gamma = gamma
        self.lam = lam
        self.adv = []
        self.ob = []
        self.ac = []
        self.rtg = []

    def store(self, temp_traj, last_sv):
        '''
        Add temp_traj values to the buffers and compute the advantage and reward to go
        Parameters:
        -----------
        temp_traj: list where each element is a list that contains: observation, reward, action, state-value
        last_sv: value of the last state (Used to Bootstrap)
        '''
        # store only if there are temporary trajectories
        if len(temp_traj) > 0:
            self.ob.extend(temp_traj[:, 0])
            rtg = discounted_rewards(temp_traj[:, 1], last_sv, self.gamma)
            self.adv.extend(GAE(temp_traj[:, 1], temp_traj[:, 3], last_sv, self.gamma, self.lam))
            self.rtg.extend(rtg)
            self.ac.extend(temp_traj[:, 2])

    def get_batch(self):
        # standardize the advantage values
        norm_adv = (self.adv - np.mean(self.adv)) / (np.std(self.adv) + 1e-10)
        return np.array(self.ob), np.array(np.expand_dims(self.ac, -1)), np.array(norm_adv), np.array(self.rtg)

    def __len__(self):
        assert (len(self.adv) == len(self.ob) == len(self.ac) == len(self.rtg))
        return len(self.ob)


class FullBuffer():
    def __init__(self):
        self.rew = []
        self.obs = []
        self.act = []
        self.nxt_obs = []
        self.done = []

        self.train_idx = []
        self.valid_idx = []
        self.full_idx = []
        self.idx = 0

    def store(self, obs, act, rew, nxt_obs, done):
        self.rew.append(rew)
        self.obs.append(obs)
        self.act.append(act)
        self.nxt_obs.append(nxt_obs)
        self.done.append(done)

        self.idx += 1

    def generate_random_dataset(self):
        rnd = np.arange(len(self.obs))
        np.random.shuffle(rnd)
        # proportion for training vs validation
        self.valid_idx = rnd[: int(len(self.obs) / 3)]
        self.train_idx = rnd[int(len(self.obs) / 3):]
        print('Train set:', len(self.train_idx), 'Valid set:', len(self.valid_idx))

    def get_training_batch(self):
        return np.array(self.obs)[self.train_idx], np.array(np.expand_dims(self.act, -1))[self.train_idx], \
               np.array(self.rew)[self.train_idx], np.array(self.nxt_obs)[self.train_idx], np.array(self.done)[
                   self.train_idx]

    def get_valid_batch(self):
        return np.array(self.obs)[self.valid_idx], np.array(np.expand_dims(self.act, -1))[self.valid_idx], \
               np.array(self.rew)[self.valid_idx], np.array(self.nxt_obs)[self.valid_idx], np.array(self.done)[
                   self.valid_idx]

    def get_full_batch(self):
        rnd = np.arange(len(self.obs))
        np.random.shuffle(rnd)
        self.full_idx = rnd
        return np.array(self.obs)[self.full_idx], np.array(np.expand_dims(self.act, -1))[self.full_idx], \
               np.array(self.rew)[self.full_idx], np.array(self.nxt_obs)[self.full_idx], np.array(self.done)[
                   self.full_idx]

    def __len__(self):
        assert (len(self.rew) == len(self.obs) == len(self.act) == len(self.nxt_obs) == len(self.done))
        return len(self.obs)


def simulate_environment(env, policy, simulated_steps):
    '''Lists to store rewards and length of the trajectories completed'''
    buffer = Buffer(0.99, 0.95)
    steps = 0
    number_episodes = 0

    while steps < simulated_steps:
        temp_buf = []
        obs = env.reset()
        number_episodes += 1
        done = False

        while not done:
            act, val = policy([obs])

            obs2, rew, done, _ = env.step([act])

            temp_buf.append([obs.copy(), rew, np.squeeze(act), np.squeeze(val)])

            obs = obs2.copy()
            steps += 1

            if done:
                buffer.store(np.array(temp_buf), 0)
                temp_buf = []

            if steps == simulated_steps:
                break

        buffer.store(np.array(temp_buf), np.squeeze(policy([obs])[1]))

    print('Sim ep:', number_episodes, end=' \n')

    return buffer.get_batch(), number_episodes


class NetworkEnv(gym.Wrapper):
    '''
    Wrapper to handle the network interaction
    '''

    def __init__(self, env, model_func=None, done_func=None, number_models=1, uncertainty_method='random'):
        gym.Wrapper.__init__(self, env)

        self.uncertainty_method = uncertainty_method
        if self.uncertainty_method not in ['random', 'pessimistic']:
            self.uncertainty_method = 'random'
            print('switched to', self.uncertainty_method, ' since ', uncertainty_method, ' is unknown.')
        self.model_func = model_func
        self.done_func = done_func
        self.number_models = number_models
        self.len_episode = 0
        # self.threshold = self.env.threshold
        # print('the threshold is: ', self.threshold)
        self.max_steps = env.max_steps

    def reset(self, **kwargs):

        # self.threshold = -0.05 * 2 + 1  # rescaled [-1,1]
        self.len_episode = 0
        self.done = False
        # kwargs['simulation'] = True
        # action = self.env.reset(**kwargs)
        if self.model_func is not None:
            obs = np.random.uniform(-1, 1, self.env.observation_space.shape)
            # print('reset', obs)
            # Todo: remove
            obs = self.env.reset()
            # obs = self.env.reset()
        else:
            # obs = self.env.reset(**kwargs)
            pass
        # Does this work?
        # self.obs = np.clip(obs, -1.0, 1.0)
        self.obs = obs.copy()
        # if self.test_phase:
        #     print('test reset', self.obs)
        # print('Reset : ',self.obs)
        return self.obs

    def step(self, action):
        if self.model_func is not None:
            # predict the next state on a random model

            if self.uncertainty_method == 'pessimistic':
                obss = []
                rews = []
                for model_idx in range(num_ensemble_models):
                    obs, rew = self.model_func(self.obs, [np.squeeze(action)], model_idx)
                    obss.append(obs)
                    rews.append(rew)
                idx = np.argmin(rew)
                # idx = np.argmax(rew)
                obs = obss[idx].copy()
                rew = rews[idx]
            else:
                obs, rew = self.model_func(self.obs, [np.squeeze(action)], np.random.randint(0, self.number_models))
            self.obs = np.clip(obs.copy(), -1.0, 1.0)
            # obs_real, rew_real, _, _ = self.env.step(action)
            # #
            # print('Diff: ', np.linalg.norm(obs - obs_real), np.linalg.norm(rew - rew_real))
            # print('Diff: ', action, obs, obs_real, rew, rew_real)
            # obs += np.random.randn(obs.shape[-1])
            # # Todo: remove
            # self.env.state = self.obs
            # done = rew > self.threshold

            self.len_episode += 1
            # print('threshold at:', self.threshold)
            # For Niky: hardcoded reward
            if rew > 0.9:  # self.threshold: TODO: to be changed
                self.done = True
                # print("Done", rew)
            if self.len_episode >= self.max_steps:
                self.done = True

            # if self.done:
            #     print('NetEnv:', self.obs, rew)
            return self.obs, (rew - 1) / 2, self.done, ""
        else:
            # self.obs, rew, done, _ = real_env.step(action)
            # return self.obs, rew, done, ""
            pass
        # return env.step(action)


class StructEnv(gym.Wrapper):
    '''
    Gym Wrapper to store information like number of steps and total reward of the last espisode.
    '''

    def __init__(self, env):
        gym.Wrapper.__init__(self, env)
        self.n_obs = self.env.reset()
        self.total_rew = 0
        self.len_episode = 0

    def reset(self, **kwargs):
        self.n_obs = self.env.reset(**kwargs)
        self.total_rew = 0
        self.len_episode = 0
        return self.n_obs.copy()

    def step(self, action):
        ob, reward, done, info = self.env.step(action)
        # print('reward in struct', reward)
        self.total_rew += reward
        self.len_episode += 1
        return ob, reward, done, info

    def get_episode_reward(self):
        return self.total_rew

    def get_episode_length(self):
        return self.len_episode


def restore_model(old_model_variables, m_variables):
    # variable used as index for restoring the actor's parameters
    it_v2 = tf.Variable(0, trainable=False)

    restore_m_params = []
    for m_v in m_variables:
        upd_m_rsh = tf.reshape(old_model_variables[it_v2: it_v2 + tf.reduce_prod(m_v.shape)], shape=m_v.shape)
        restore_m_params.append(m_v.assign(upd_m_rsh))
        it_v2 += tf.reduce_prod(m_v.shape)

    return tf.group(*restore_m_params)


def METRPO(env_name, hidden_sizes=[32, 32], cr_lr=5e-3, num_epochs=50, gamma=0.99, lam=0.95, number_envs=1,
           critic_iter=10, steps_per_env=100, delta=0.05, algorithm='TRPO', conj_iters=10, minibatch_size=1000,
           mb_lr_start=0.0001, model_batch_size=512, simulated_steps=1000, num_ensemble_models=2, model_iter=15,
           init_random_steps=steps_per_env, method='random'):
    '''
    Model Ensemble Trust Region Policy Optimization
    The states and actions are provided by the gym environement with the correct boxs.
    The reward has to be between [-1,0].
    Parameters:
    -----------
    env_name: Name of the environment
    hidden_sizes: list of the number of hidden units for each layer
    cr_lr: critic learning rate
    num_epochs: number of training epochs
    gamma: discount factor
    lam: lambda parameter for computing the GAE
    number_envs: number of "parallel" synchronous environments
        # NB: it isn't distributed across multiple CPUs
    critic_iter: Number of SGD iterations on the critic per epoch
    steps_per_env: number of steps per environment
            # NB: the total number of steps per epoch will be: steps_per_env*number_envs
    delta: Maximum KL divergence between two policies. Scalar value
    algorithm: type of algorithm. Either 'TRPO' or 'NPO'
    conj_iters: number of conjugate gradient iterations
    minibatch_size: Batch size used to train the critic
    mb_lr: learning rate of the environment model
    model_batch_size: batch size of the environment model
    simulated_steps: number of simulated steps for each policy update
    num_ensemble_models: number of models
    model_iter: number of iterations without improvement before stopping training the model
    '''
    # TODO: add ME-TRPO hyperparameters

    tf.reset_default_graph()

    # Create a few environments to collect the trajectories

    # envs = [StructEnv(gym.make(env_name)) for _ in range(number_envs)]
    envs = [StructEnv(make_env()) for _ in range(number_envs)]
    env_test = StructEnv(make_env(test_env=True))
    # env_test = gym.make(env_name)
    print('env_test' * 4)

    # env_test = make_env(test=True)
    # env_test = gym.wrappers.Monitor(env_test, "VIDEOS/", force=True, video_callable=lambda x: x%10 == 0)
    # to be changed in real test
    # env_test = FelLocalEnv(tango=tango)
    # env_test.test = True
    # env_test_1 = FelLocalEnv(tango=tango)
    # env_test_1.test = True

    # If the scaling is not perfomed this has to be changed
    low_action_space = -1  # envs[0].action_space.low
    high_action_space = 1  # envs[0].action_space.high

    obs_dim = envs[0].observation_space.shape
    act_dim = envs[0].action_space.shape[0]

    # print(envs[0].action_space, envs[0].observation_space, low_action_space,
    #       high_action_space)

    # Placeholders for model
    act_ph = tf.placeholder(shape=(None, act_dim), dtype=tf.float64, name='act')
    obs_ph = tf.placeholder(shape=(None, obs_dim[0]), dtype=tf.float64, name='obs')
    # NEW
    nobs_ph = tf.placeholder(shape=(None, obs_dim[0]), dtype=tf.float64, name='nobs')
    rew_ph = tf.placeholder(shape=(None, 1), dtype=tf.float64, name='rew')

    ret_ph = tf.placeholder(shape=(None,), dtype=tf.float64, name='ret')
    adv_ph = tf.placeholder(shape=(None,), dtype=tf.float64, name='adv')
    old_p_log_ph = tf.placeholder(shape=(None,), dtype=tf.float64, name='old_p_log')
    old_mu_ph = tf.placeholder(shape=(None, act_dim), dtype=tf.float64, name='old_mu')
    old_log_std_ph = tf.placeholder(shape=(act_dim), dtype=tf.float64, name='old_log_std')
    p_ph = tf.placeholder(shape=(None,), dtype=tf.float64, name='p_ph')

    # Placeholder for learning rate
    mb_lr_ = tf.placeholder("float", None)  # , name='mb_lr')

    # result of the conjugate gradient algorithm
    cg_ph = tf.placeholder(shape=(None,), dtype=tf.float64, name='cg')

    #########################################################
    ######################## POLICY #########################
    #########################################################

    old_model_variables = tf.placeholder(shape=(None,), dtype=tf.float64, name='old_model_variables')

    # Neural network that represent the policy
    with tf.variable_scope('actor_nn'):
        p_means = mlp(obs_ph, hidden_sizes, act_dim, tf.tanh, last_activation=tf.tanh)
        p_means = tf.clip_by_value(p_means, low_action_space, high_action_space)
        log_std = tf.get_variable(name='log_std', initializer=np.ones(act_dim, dtype=np.float64))

    # Neural network that represent the value function
    with tf.variable_scope('critic_nn'):
        s_values = mlp(obs_ph, hidden_sizes, 1, tf.tanh, last_activation=None)
        s_values = tf.squeeze(s_values)

        # Add "noise" to the predicted mean following the Gaussian distribution with standard deviation e^(log_std)
    p_noisy = p_means + tf.random_normal(tf.shape(p_means), 0, 1, dtype=tf.float64) * tf.exp(log_std)
    # Clip the noisy actions
    a_sampl = tf.clip_by_value(p_noisy, low_action_space, high_action_space)

    # Compute the gaussian log likelihood
    p_log = gaussian_log_likelihood(act_ph, p_means, log_std)

    # Measure the divergence
    diverg = tf.reduce_mean(tf.exp(old_p_log_ph - p_log))

    # ratio
    ratio_new_old = tf.exp(p_log - old_p_log_ph)
    # TRPO surrogate loss function
    p_loss = - tf.reduce_mean(ratio_new_old * adv_ph)

    # MSE loss function
    v_loss = tf.reduce_mean((ret_ph - s_values) ** 2)
    # Critic optimization
    v_opt = tf.train.AdamOptimizer(cr_lr).minimize(v_loss)

    def variables_in_scope(scope):
        # get all trainable variables in 'scope'
        return tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, scope)

    # Gather and flatten the actor parameters
    p_variables = variables_in_scope('actor_nn')
    p_var_flatten = flatten_list(p_variables)

    # Gradient of the policy loss with respect to the actor parameters
    p_grads = tf.gradients(p_loss, p_variables)
    p_grads_flatten = flatten_list(p_grads)

    ########### RESTORE ACTOR PARAMETERS ###########
    p_old_variables = tf.placeholder(shape=(None,), dtype=tf.float64, name='p_old_variables')
    # variable used as index for restoring the actor's parameters
    it_v1 = tf.Variable(0, trainable=False)
    restore_params = []

    for p_v in p_variables:
        upd_rsh = tf.reshape(p_old_variables[it_v1: it_v1 + tf.reduce_prod(p_v.shape)], shape=p_v.shape)
        restore_params.append(p_v.assign(upd_rsh))
        it_v1 += tf.reduce_prod(p_v.shape)

    restore_params = tf.group(*restore_params)

    # gaussian KL divergence of the two policies
    dkl_diverg = gaussian_DKL(old_mu_ph, old_log_std_ph, p_means, log_std)

    # Jacobian of the KL divergence (Needed for the Fisher matrix-vector product)
    dkl_diverg_grad = tf.gradients(dkl_diverg, p_variables)

    dkl_matrix_product = tf.reduce_sum(flatten_list(dkl_diverg_grad) * p_ph)
    print('dkl_matrix_product', dkl_matrix_product.shape)
    # Fisher vector product
    # The Fisher-vector product is a way to compute the A matrix without the need of the full A
    Fx = flatten_list(tf.gradients(dkl_matrix_product, p_variables))

    ## Step length
    beta_ph = tf.placeholder(shape=(), dtype=tf.float64, name='beta')
    # NPG update
    npg_update = beta_ph * cg_ph

    ## alpha is found through line search
    alpha = tf.Variable(1., trainable=False, dtype=tf.float64)
    # TRPO update
    trpo_update = alpha * npg_update

    ####################   POLICY UPDATE  ###################
    # variable used as an index
    it_v = tf.Variable(0, trainable=False)
    p_opt = []
    # Apply the updates to the policy
    for p_v in p_variables:
        print(p_v)
        upd_rsh = tf.reshape(trpo_update[it_v: it_v + tf.reduce_prod(p_v.shape)], shape=p_v.shape)
        p_opt.append(p_v.assign_sub(upd_rsh))
        it_v += tf.reduce_prod(p_v.shape)

    p_opt = tf.group(*p_opt)

    #########################################################
    ######################### MODEL #########################
    #########################################################

    # Create a new class for the model:
    # NN class
    class NN:
        def __init__(self, x, y, y_dim, hidden_size, n, learning_rate, init_params):
            self.init_params = init_params

            # set up NN
            with tf.variable_scope('model_' + str(n) + '_nn'):
                self.inputs = x
                self.y_target = y

                self.layer_1_w = tf.layers.Dense(hidden_size,
                                                 activation=tf.tanh,
                                                 kernel_initializer=tf.random_normal_initializer(mean=0.,
                                                                                                 stddev=self.init_params.get(
                                                                                                     'init_stddev_1_w')),
                                                 bias_initializer=tf.random_normal_initializer(mean=0.,
                                                                                               stddev=self.init_params.get(
                                                                                                   'init_stddev_1_b')))

                self.layer_1 = self.layer_1_w.apply(self.inputs)

                self.layer_2_w = tf.layers.Dense(hidden_size,
                                                 activation=tf.tanh,
                                                 kernel_initializer=tf.random_normal_initializer(mean=0.,
                                                                                                 stddev=self.init_params.get(
                                                                                                     'init_stddev_1_w')),
                                                 bias_initializer=tf.random_normal_initializer(mean=0.,
                                                                                               stddev=self.init_params.get(
                                                                                                   'init_stddev_1_b')))

                self.layer_2 = self.layer_2_w.apply(self.layer_1)

                self.output_w = tf.layers.Dense(y_dim,
                                                activation=tf.tanh,
                                                use_bias=False,
                                                kernel_initializer=tf.random_normal_initializer(mean=0.,
                                                                                                stddev=self.init_params.get(
                                                                                                    'init_stddev_2_w')))
                # self.layer_1_w = tf.layers.Dense(hidden_size,
                #                                  activation=tf.tanh
                #                                  )
                #
                # self.layer_1 = self.layer_1_w.apply(self.inputs)
                #
                # self.layer_2_w = tf.layers.Dense(hidden_size,
                #                                  activation=tf.tanh)
                #
                # self.layer_2 = self.layer_2_w.apply(self.layer_1)
                #
                # self.output_w = tf.layers.Dense(y_dim,
                #                                 activation=tf.tanh)
                # #

                self.output = self.output_w.apply(self.layer_2)

                # set up loss and optimiser - we'll modify this later with anchoring regularisation
                self.opt_method = tf.train.AdamOptimizer(learning_rate)
                self.mse_ = 1 / tf.shape(self.inputs, out_type=tf.int64)[0] * \
                            tf.reduce_sum(tf.square(self.y_target - self.output))
                self.loss_ = 1 / tf.shape(self.inputs, out_type=tf.int64)[0] * \
                             tf.reduce_sum(tf.square(self.y_target - self.output))
                self.optimizer = self.opt_method.minimize(self.loss_)
                self.optimizer_mse = self.opt_method.minimize(self.mse_)
                # return self.mse_

        def get_weights(self):
            '''method to return current params'''

            ops = [self.layer_1_w.kernel, self.layer_1_w.bias,
                   self.layer_2_w.kernel, self.layer_2_w.bias,
                   self.output_w.kernel]
            w1, b1, w2, b2, w = sess.run(ops)

            return w1, b1, w2, b2, w

        # def set_weights(self, variables):
        #     model_assign(self.n, initial_variables_models[model_idx])

        def anchor(self, lambda_anchor):
            '''regularise around initialised parameters'''

            w1, b1, w2, b2, w = self.get_weights()

            # get initial params to hold for future trainings
            self.w1_init, self.b1_init, self.w2_init, self.b2_init, self.w_out_init = w1, b1, w2, b2, w

            loss_anchor = lambda_anchor[0] * tf.reduce_sum(tf.square(self.w1_init - self.layer_1_w.kernel))
            loss_anchor += lambda_anchor[1] * tf.reduce_sum(tf.square(self.b1_init - self.layer_1_w.bias))

            loss_anchor = lambda_anchor[0] * tf.reduce_sum(tf.square(self.w2_init - self.layer_2_w.kernel))
            loss_anchor += lambda_anchor[1] * tf.reduce_sum(tf.square(self.b2_init - self.layer_2_w.bias))

            loss_anchor += lambda_anchor[2] * tf.reduce_sum(tf.square(self.w_out_init - self.output_w.kernel))

            # combine with original loss
            # norm_val = 1/tf.shape(self.inputs)[0]
            self.loss_ = self.loss_ + tf.scalar_mul(1 / tf.shape(self.inputs)[0], loss_anchor)
            # self.loss_ = self.loss_ + tf.scalar_mul(1 / 1000, loss_anchor)
            self.optimizer = self.opt_method.minimize(self.loss_)
            return self.optimizer, self.loss_

    m_opts = []
    m_losses = []

    nobs_pred_m = []
    act_obs = tf.concat([obs_ph, act_ph], 1)
    target = tf.concat([nobs_ph, rew_ph], 1)

    # computational graph of N models and the correct losses for the anchor method
    m_classes = []
    for i in range(num_ensemble_models):
        # TODO: Add variable size of network
        hidden_sizes = [100, 100]
        nobs_pred = mlp(x=act_obs, hidden_layers=hidden_sizes, output_layer=obs_dim[0] + 1,
                        activation=tf.tanh, last_activation=tf.tanh)  # , init_params=init_params)
        m_class = NN(x=tf.concat([obs_ph, act_ph], 1), y=tf.concat([nobs_ph, rew_ph], 1), y_dim=obs_dim[0] + 1,
                     learning_rate=1e-3, n=i,
                     hidden_size=100, init_params=init_params)

        nobs_pred = m_class.output

        nobs_pred_m.append(nobs_pred)

        # m_loss = tf.reduce_mean((tf.concat([nobs_ph, rew_ph], 1) - nobs_pred) ** 2)
        # m_opts.append(tf.train.AdamOptimizer(learning_rate=mb_lr_).minimize(m_loss))
        # m_losses.append(m_loss)

        m_classes.append(m_class)
        # m_losses.append(m_class.mse_)
        # m_opts.append(m_class.optimizer_mse)

    ##################### RESTORE MODEL ######################
    initialize_models = []
    models_variables = []
    for i in range(num_ensemble_models):
        m_variables = variables_in_scope('model_' + str(i) + '_nn')
        # print('m_variables', m_variables)
        initialize_models.append(restore_model(old_model_variables, m_variables))
        #  List of weights as numpy
        models_variables.append(flatten_list(m_variables))

    #########################################################
    ##################### END MODEL #########################
    #########################################################
    # Time
    now = datetime.now()
    clock_time = "{}_{}_{}_{}".format(now.day, now.hour, now.minute, now.second)
    print('Time:', clock_time)

    # Set scalars and hisograms for TensorBoard
    tf.summary.scalar('p_loss', p_loss, collections=['train'])
    tf.summary.scalar('v_loss', v_loss, collections=['train'])
    tf.summary.scalar('p_divergence', diverg, collections=['train'])
    tf.summary.scalar('ratio_new_old', tf.reduce_mean(ratio_new_old), collections=['train'])
    tf.summary.scalar('dkl_diverg', dkl_diverg, collections=['train'])
    tf.summary.scalar('alpha', alpha, collections=['train'])
    tf.summary.scalar('beta', beta_ph, collections=['train'])
    tf.summary.scalar('p_std_mn', tf.reduce_mean(tf.exp(log_std)), collections=['train'])
    tf.summary.scalar('s_values_mn', tf.reduce_mean(s_values), collections=['train'])
    tf.summary.histogram('p_log', p_log, collections=['train'])
    tf.summary.histogram('p_means', p_means, collections=['train'])
    tf.summary.histogram('s_values', s_values, collections=['train'])
    tf.summary.histogram('adv_ph', adv_ph, collections=['train'])
    tf.summary.histogram('log_std', log_std, collections=['train'])
    scalar_summary = tf.summary.merge_all('train')

    tf.summary.scalar('old_v_loss', v_loss, collections=['pre_train'])
    tf.summary.scalar('old_p_loss', p_loss, collections=['pre_train'])
    pre_scalar_summary = tf.summary.merge_all('pre_train')

    hyp_str = '-spe_' + str(steps_per_env) + '-envs_' + str(number_envs) + '-cr_lr' + str(cr_lr) + '-crit_it_' + str(
        critic_iter) + '-delta_' + str(delta) + '-conj_iters_' + str(conj_iters)

    file_writer = tf.summary.FileWriter('log_dir/' + env_name + '/' + algorithm + '_' + clock_time + '_' + hyp_str,
                                        tf.get_default_graph())

    #################################################################################################
    # Session start!!!!!!!!
    # create a session
    sess = tf.Session()
    # initialize the variables
    sess.run(tf.global_variables_initializer())

    def action_op(o):
        return sess.run([p_means, s_values], feed_dict={obs_ph: o})

    def action_op_noise(o):
        return sess.run([a_sampl, s_values], feed_dict={obs_ph: o})

    def model_op(o, a, md_idx):
        mo = sess.run(nobs_pred_m[md_idx], feed_dict={obs_ph: [o], act_ph: [a[0]]})
        return np.squeeze(mo[:, :-1]), np.squeeze(mo[:, -1])

    def run_model_loss(model_idx, r_obs, r_act, r_nxt_obs, r_rew):
        # print({'obs_ph': r_obs.shape, 'act_ph': r_act.shape, 'nobs_ph': r_nxt_obs.shape})
        r_act = np.squeeze(r_act, axis=2)
        # print(r_act.shape)
        r_rew = np.reshape(r_rew, (-1, 1))
        # print(r_rew.shape)
        return_val = sess.run(m_loss_anchor[model_idx],
                              feed_dict={obs_ph: r_obs, act_ph: r_act, nobs_ph: r_nxt_obs, rew_ph: r_rew})
        # return_val = sess.run(m_losses[model_idx],
        #                       feed_dict={obs_ph: r_obs, act_ph: r_act, nobs_ph: r_nxt_obs, rew_ph: r_rew})
        return return_val

    def run_model_opt_loss(model_idx, r_obs, r_act, r_nxt_obs, r_rew, mb_lr):
        r_act = np.squeeze(r_act, axis=2)
        r_rew = np.reshape(r_rew, (-1, 1))
        # return sess.run([m_opts[model_idx], m_losses[model_idx]],
        #                 feed_dict={obs_ph: r_obs, act_ph: r_act, nobs_ph: r_nxt_obs, rew_ph: r_rew, mb_lr_: mb_lr})
        return sess.run([m_opts_anchor[model_idx], m_loss_anchor[model_idx]],
                        feed_dict={obs_ph: r_obs, act_ph: r_act, nobs_ph: r_nxt_obs, rew_ph: r_rew, mb_lr_: mb_lr})

    def model_assign(i, model_variables_to_assign):
        '''
        Update the i-th model's parameters
        '''
        return sess.run(initialize_models[i], feed_dict={old_model_variables: model_variables_to_assign})

    # def anchor(model_idx):
    #     m_classes[model_idx].anchor(lambda_anchor=lambda_anchor)

    def policy_update(obs_batch, act_batch, adv_batch, rtg_batch, it):
        # log probabilities, logits and log std of the "old" policy
        # "old" policy refer to the policy to optimize and that has been used to sample from the environment
        act_batch = np.squeeze(act_batch, axis=2)
        old_p_log, old_p_means, old_log_std = sess.run([p_log, p_means, log_std],
                                                       feed_dict={obs_ph: obs_batch, act_ph: act_batch,
                                                                  adv_ph: adv_batch, ret_ph: rtg_batch})
        # get also the "old" parameters
        old_actor_params = sess.run(p_var_flatten)
        if it < 1:
            std_vals = sess.run([log_std], feed_dict={log_std: np.ones(act_dim)})
            # print(std_vals)
        # old_p_loss is later used in the line search
        # run pre_scalar_summary for a summary before the optimization
        old_p_loss, summary = sess.run([p_loss, pre_scalar_summary],
                                       feed_dict={obs_ph: obs_batch, act_ph: act_batch, adv_ph: adv_batch,
                                                  ret_ph: rtg_batch, old_p_log_ph: old_p_log})
        file_writer.add_summary(summary, step_count)

        file_writer.add_summary(summary, step_count)
        file_writer.flush()

        def H_f(p):
            '''
            Run the Fisher-Vector product on 'p' to approximate the Hessian of the DKL
            '''
            return sess.run(Fx,
                            feed_dict={old_mu_ph: old_p_means, old_log_std_ph: old_log_std, p_ph: p, obs_ph: obs_batch,
                                       act_ph: act_batch, adv_ph: adv_batch, ret_ph: rtg_batch})

        g_f = sess.run(p_grads_flatten,
                       feed_dict={old_mu_ph: old_p_means, obs_ph: obs_batch, act_ph: act_batch, adv_ph: adv_batch,
                                  ret_ph: rtg_batch, old_p_log_ph: old_p_log})
        ## Compute the Conjugate Gradient so to obtain an approximation of H^(-1)*g
        # Where H in reality isn't the true Hessian of the KL divergence but an approximation of it computed via Fisher-Vector Product (F)
        conj_grad = conjugate_gradient(H_f, g_f, iters=conj_iters)

        # Compute the step length
        beta_np = np.sqrt(2 * delta / (1e-10 + np.sum(conj_grad * H_f(conj_grad))))

        def DKL(alpha_v):
            '''
            Compute the KL divergence.
            It optimize the function to compute the DKL. Afterwards it restore the old parameters.
            '''
            sess.run(p_opt, feed_dict={beta_ph: beta_np, alpha: alpha_v, cg_ph: conj_grad, obs_ph: obs_batch,
                                       act_ph: act_batch, adv_ph: adv_batch, old_p_log_ph: old_p_log})
            a_res = sess.run([dkl_diverg, p_loss],
                             feed_dict={old_mu_ph: old_p_means, old_log_std_ph: old_log_std, obs_ph: obs_batch,
                                        act_ph: act_batch, adv_ph: adv_batch, ret_ph: rtg_batch,
                                        old_p_log_ph: old_p_log})
            sess.run(restore_params, feed_dict={p_old_variables: old_actor_params})
            return a_res

        # Actor optimization step
        # Different for TRPO or NPG
        # Backtracing line search to find the maximum alpha coefficient s.t. the constraint is valid
        best_alpha = backtracking_line_search(DKL, delta, old_p_loss, p=0.8)
        sess.run(p_opt, feed_dict={beta_ph: beta_np, alpha: best_alpha,
                                   cg_ph: conj_grad, obs_ph: obs_batch, act_ph: act_batch,
                                   adv_ph: adv_batch, old_p_log_ph: old_p_log})

        lb = len(obs_batch)
        shuffled_batch = np.arange(lb)
        np.random.shuffle(shuffled_batch)

        # Value function optimization steps
        for _ in range(critic_iter):
            # shuffle the batch on every iteration
            np.random.shuffle(shuffled_batch)
            for idx in range(0, lb, minibatch_size):
                minib = shuffled_batch[idx:min(idx + minibatch_size, lb)]
                sess.run(v_opt, feed_dict={obs_ph: obs_batch[minib], ret_ph: rtg_batch[minib]})

    def train_model(tr_obs, tr_act, tr_nxt_obs, tr_rew, v_obs, v_act, v_nxt_obs, v_rew,
                    f_obs, f_act, f_nxt_obs, f_rew,
                    step_count, model_idx, mb_lr):

        # Get validation loss on the old model only used for monitoring
        mb_valid_loss1 = run_model_loss(model_idx, v_obs, v_act, v_nxt_obs, v_rew)

        # Restore the initial random weights to have a new, clean neural network
        # initial_variables_models - list stored before already in the code below -
        # important for the anchor method
        model_assign(model_idx, initial_variables_models[model_idx])

        # update the anchor model losses:
        # m_opts_anchor = []
        # m_loss_anchor = []
        # for i in range(num_ensemble_models):
        # This stores the anchor weights to the models for the regularisation:
        # opt, loss = m_classes[model_idx].anchor(lambda_anchor=lambda_anchor)
        # m_opts_anchor[model_idx] = opt
        # m_loss_anchor[model_idx] = loss

        # Get validation loss on the now initialized model
        mb_valid_loss = run_model_loss(model_idx, v_obs, v_act, v_nxt_obs, v_rew)

        acc_m_losses = []
        last_m_losses = []
        md_params = sess.run(models_variables[model_idx])
        best_mb = {'iter': 0, 'loss': mb_valid_loss, 'params': md_params}
        iteration = 0

        # Create mini-batch for training
        lb = len(f_obs)
        shuffled_batch = np.arange(lb)
        np.random.shuffle(shuffled_batch)

        # Run until the number of model_iter has passed from the best val loss at it on...
        for iteration in range(0, training_steps_model):

            # update the model on each mini-batch
            last_m_losses = []
            # for idx in range(0, 20): # repeat gradient step for some time
            # minib = shuffled_batch[idx:min(idx + minibatch_size, lb)]
            # want all data (shuffeld)
            # np.random.shuffle(shuffled_batch)
            minib = shuffled_batch
            if len(minib) != minibatch_size:
                _, ml = run_model_opt_loss(model_idx, f_obs[minib], f_act[minib], f_nxt_obs[minib],
                                           f_rew[minib], mb_lr=mb_lr)
                acc_m_losses.append(ml)
                last_m_losses.append(ml)
            else:
                pass
                # print('Warning!')

            # Check if the loss on the validation set has improved
            mb_valid_loss = run_model_loss(model_idx, v_obs, v_act, v_nxt_obs, v_rew)
            # if mb_valid_loss < best_mb['loss']:
            best_mb['loss'] = mb_valid_loss
            best_mb['iter'] = iteration
            # store the parameters to the array
            best_mb['params'] = sess.run(models_variables[model_idx])

            # it += 1

            # if it>=10000:
            #     break
            # print('iteration: ', it)

        # Restore the model with the lower validation loss
        model_assign(model_idx, best_mb['params'])

        print('Model:{}, iter:{} -- Old Val loss:{:.6f}  '
              'New Val loss:{:.6f} -- New Train loss:{:.6f}'.format(model_idx, iteration,
                                                                    mb_valid_loss1, best_mb['loss'],
                                                                    np.mean(last_m_losses)))
        summary = tf.Summary()
        summary.value.add(tag='supplementary/m_loss', simple_value=np.mean(acc_m_losses))
        summary.value.add(tag='supplementary/iterations', simple_value=iteration)
        file_writer.add_summary(summary, step_count)
        file_writer.flush()

    def plot_results(env_wrapper, label, **kwargs):
        # plotting
        print('now plotting...')
        rewards = env_wrapper.env.current_buffer.get_data()['rews']

        # states = env_wrapper.env.current_buffer.get_data()['obss']
        # pd.DataFrame(states).hist()
        # plt.show()

        # initial_states = env.initial_conditions

        iterations = []
        finals = []
        means = []
        stds = []

        # init_states = pd.read_pickle('/Users/shirlaen/PycharmProjects/DeepLearning/spinningup/Environments/initData')

        for i in range(len(rewards)):
            if (len(rewards[i]) > 1):
                # finals.append(rewards[i][len(rewards[i]) - 1])
                finals.append(rewards[i][-1])
                means.append(np.mean(rewards[i][1:]))
                stds.append(np.std(rewards[i][1:]))
                iterations.append(len(rewards[i]))
        # print(iterations)
        x = range(len(iterations))
        iterations = np.array(iterations)
        finals = np.array(finals)
        means = np.array(means)
        stds = np.array(stds)

        plot_suffix = label  # , Fermi time: {env.TOTAL_COUNTER / 600:.1f} h'

        fig, axs = plt.subplots(2, 1, sharex=True)

        ax = axs[0]
        ax.plot(x, iterations)
        ax.set_ylabel('Iterations (1)')
        ax.set_title(plot_suffix)
        # fig.suptitle(label, fontsize=12)
        if 'data_number' in kwargs:
            ax1 = plt.twinx(ax)
            color = 'lime'
            ax1.set_ylabel('Mean reward', color=color)  # we already handled the x-label with ax1
            ax1.tick_params(axis='y', labelcolor=color)
            ax1.plot(x, kwargs.get('data_number'), color=color)

        ax = axs[1]
        color = 'blue'
        ax.set_ylabel('Final reward', color=color)  # we already handled the x-label with ax1
        ax.tick_params(axis='y', labelcolor=color)
        ax.plot(x, finals, color=color)

        ax.set_title('Final reward per episode')  # + plot_suffix)
        ax.set_xlabel('Episodes (1)')

        ax1 = plt.twinx(ax)
        color = 'lime'
        ax1.set_ylabel('Mean reward', color=color)  # we already handled the x-label with ax1
        ax1.tick_params(axis='y', labelcolor=color)
        ax1.fill_between(x, means - stds, means + stds,
                         alpha=0.5, edgecolor=color, facecolor='#FF9848')
        ax1.plot(x, means, color=color)

        # ax.set_ylim(ax1.get_ylim())
        if 'save_name' in kwargs:
            plt.savefig(kwargs.get('save_name') + '.pdf')
        # fig.tight_layout()
        plt.show()

    def plot_observables(data, label, **kwargs):
        """plot observables during the test"""

        sim_rewards_all = np.array(data.get('sim_rewards_all'))
        step_counts_all = np.array(data.get('step_counts_all'))
        batch_rews_all = np.array(data.get('batch_rews_all'))
        tests_all = np.array(data.get('tests_all'))

        fig, axs = plt.subplots(2, 1, sharex=True)
        x = np.arange(len(batch_rews_all[0]))
        ax = axs[0]
        ax.step(x, batch_rews_all[0])
        ax.fill_between(x, batch_rews_all[0] - batch_rews_all[1], batch_rews_all[0] + batch_rews_all[1],
                        alpha=0.5)
        ax.set_ylabel('rews per batch')

        ax.set_title(label)

        # plt.tw
        ax2 = ax.twinx()

        color = 'lime'
        ax2.set_ylabel('data points', color=color)  # we already handled the x-label with ax1
        ax2.tick_params(axis='y', labelcolor=color)
        ax2.step(x, step_counts_all, color=color)

        ax = axs[1]
        ax.grid(True)
        ax.plot(sim_rewards_all[0], ls=':')
        ax.fill_between(x, sim_rewards_all[0] - sim_rewards_all[1], sim_rewards_all[0] + sim_rewards_all[1],
                        alpha=0.5)

        try:
            ax.plot(tests_all[0])
            ax.fill_between(x, tests_all[0] - tests_all[1], tests_all[0] + tests_all[1],
                            alpha=0.5)
        except:
            pass

        ax.set_ylabel('rewards tests')
        # plt.tw
        ax2 = ax.twinx()

        color = 'lime'
        ax2.set_ylabel('entropy', color=color)  # we already handled the x-label with ax1
        ax2.tick_params(axis='y', labelcolor=color)
        ax2.plot(entropy_all, color=color)

        plt.show()

    def save_data(data, **kwargs):
        '''logging functon'''
        # if 'directory_name' in kwargs:
        #     project_directory = kwargs.get('directory_name')
        now = datetime.now()
        clock_time = f'{now.month:0>2}_{now.day:0>2}_{now.hour:0>2}_{now.minute:0>2}_{now.second:0>2}'
        out_put_writer = open(project_directory + clock_time + '_training_observables', 'wb')
        pickle.dump(data, out_put_writer, -1)
        out_put_writer.close()

    # variable to store the total number of steps
    step_count = 0
    model_buffer = FullBuffer()
    print('Env batch size:', steps_per_env, ' Batch size:', steps_per_env * number_envs)

    # Create a simulated environment
    sim_env = NetworkEnv(make_env(), model_op, None, number_models=num_ensemble_models, uncertainty_method=method)

    # ------------------------------------------------------------------------------------------------------
    # -------------------------------------Try to set correct anchors---------------------------------------
    # Get the initial parameters of each model
    # These are used in later epochs when we aim to re-train the models anew with the new dataset
    initial_variables_models = []
    for model_var in models_variables:
        # num list of parameters
        initial_variables_models.append(sess.run(model_var))

    # update the anchor model losses:
    m_opts_anchor = []
    m_loss_anchor = []
    for i in range(num_ensemble_models):
        # This stores the anchor weights to the models for the regularisation:
        opt, loss = m_classes[i].anchor(lambda_anchor=lambda_anchor)
        m_opts_anchor.append(opt)
        m_loss_anchor.append(loss)

    # ------------------------------------------------------------------------------------------------------
    # -------------------------------------Try to set correct anchors---------------------------------------

    total_iterations = 0

    converged_flag = False
    # save_data = save_data(clock_time)
    sim_rewards_all = []
    sim_rewards_std_all = []
    entropy_all = []
    tests_all = []
    tests_std_all = []
    batch_rews_all = []
    batch_rews_std_all = []
    step_counts_all = []



    for ep in range(num_epochs):
        if (converged_flag):
            print('Converged!!!!')
            break
        # lists to store rewards and length of the trajectories completed
        batch_rew = []
        batch_len = []
        print('============================', ep, '============================')
        # Execute in serial the environment, storing temporarily the trajectories.
        for env in envs:
            # Todo: Test randomization stronger if reward lower...we need a good scheme
            # target_threshold ?????
            init_log_std = np.ones(act_dim) * np.log(np.random.rand() * 1)
            env.reset()

            # iterate over a fixed number of steps
            steps_train = init_random_steps if ep == 0 else steps_per_env
            # steps_train = steps_per_env
            for _ in range(steps_train):
                # found = False
                # while not(found):
                # run the policy

                if ep == 0:
                    # Sample random action during the first epoch
                    act = np.random.uniform(-1, 1, size=env.action_space.shape[-1])

                else:

                    act = sess.run(a_sampl, feed_dict={obs_ph: [env.n_obs], log_std: init_log_std})
                    # act = np.clip(act + np.random.randn(act.shape[0], act.shape[1]) * 0.1, -1, 1)

                act = np.squeeze(act)
                # print('act', act*12)
                # take a step in the environment
                obs2, rew, done, _ = env.step(np.array(act))

                # add the new transition to the temporary buffer
                model_buffer.store(env.n_obs.copy(), act, rew.copy(), obs2.copy(), done)

                env.n_obs = obs2.copy()
                step_count += 1

                if done:
                    batch_rew.append(env.get_episode_reward())
                    batch_len.append(env.get_episode_length())

                    env.reset()
                    init_log_std = np.ones(act_dim) * np.log(np.random.rand() * 1)

                # if ep == 0:
                #     # try:
                #     # Initialize randomly a training and validation set
                #     model_buffer.generate_random_dataset()
                #     # get both datasets
                #     train_obs, train_act, train_rew, train_nxt_obs, _ = model_buffer.get_training_batch()
                #     valid_obs, valid_act, valid_rew, valid_nxt_obs, _ = model_buffer.get_valid_batch()
                #     target_threshold = max(max(valid_rew), max(train_rew))
                #     # print('-- '*38, target_threshold)
                #     found = target_threshold>=-0.1 and step_count>=191
                #     # except:
                #     #     pass

        # save the data for plotting the collected data for the model
        env.save_current_buffer()

        print('Ep:%d Rew:%.2f -- Step:%d' % (ep, np.mean(batch_rew), step_count))

        # env_test.env.set_usage('default')
        # plot_results(env_test, f'Total {total_iterations}, '
        #                        f'data points: {len(model_buffer.train_idx) + len(model_buffer.valid_idx)}, '
        #                        f'modelit: {ep}')
        ############################################################
        ###################### MODEL LEARNING ######################
        ############################################################

        # Initialize randomly a training and validation set
        model_buffer.generate_random_dataset()

        # get both datasets
        train_obs, train_act, train_rew, train_nxt_obs, _ = model_buffer.get_training_batch()
        valid_obs, valid_act, valid_rew, valid_nxt_obs, _ = model_buffer.get_valid_batch()
        full_obs, full_act, full_rew, full_nxt_obs, _ = model_buffer.get_full_batch()

        std_vals = sess.run(log_std)
        print('Log Std policy:', std_vals, np.mean(std_vals))

        target_threshold = max(max(valid_rew), max(train_rew))
        sim_env.threshold = target_threshold  # min(target_threshold, -0.05)
        print('maximum: ', sim_env.threshold)

        # Learning rate as function of ep
        lr = lambda ep: 1e-3 - ep / num_epochs * (1e-3 - 5e-4)
        mb_lr = 1e-3  # if ep < 1 else 1e-3
        # simulated_steps = simulated_steps if ep<10 else 10000
        print('mb_lr: ', mb_lr)
        for i in range(num_ensemble_models):
            # train the dynamic model on the datasets just sampled
            train_model(train_obs, train_act, train_nxt_obs, train_rew, valid_obs, valid_act, valid_nxt_obs, valid_rew,
                        full_obs, full_act, full_nxt_obs, full_rew,
                        step_count, i, mb_lr=mb_lr)

        ############################################################
        ###################### POLICY LEARNING ######################
        ############################################################

        best_sim_test = -1e16 * np.ones(num_ensemble_models)

        # plot_results(env_test, f'Total {total_iterations}, '
        #                        f'data points: {len(model_buffer.train_idx) + len(model_buffer.valid_idx)}, '
        #                        f'modelit: {ep}')

        # # For Niky:  perform default test! -----------------------------
        # if ep == 0:
        #     it = 0
        #     # perform test! -----------------------------
        #     env_test.env.set_usage('test')
        #     mn_test, mn_test_std, mn_length, mn_success = test_agent(env_test, action_op, num_games=50)
        #     # perform test! -----------------------------
        #     label = f'Init {total_iterations}, ' + \
        #             f'data points: {len(model_buffer.train_idx) + len(model_buffer.valid_idx)}, ' + \
        #             f'ep: {ep}, it: {it}\n' + hyp_str_all
        #     # plot results of test -----------------------------
        #     plot_results(env_test, label=label)

        env_test.env.set_usage('default')

        for it in range(max_training_iterations):
            if converged_flag:
                break
            total_iterations += 1
            print('\t Policy it', it, end='..')

            ##################### MODEL SIMLUATION #####################
            # obs_batch, act_batch, adv_batch, rtg_batch = simulate_environment(sim_env, action_op_noise, simulated_steps)
            batch, ep_length = simulate_environment(sim_env, action_op_noise, simulated_steps)
            # verification_simulate_environment(sim_env, env_test, action_op_noise, 50)
            obs_batch, act_batch, adv_batch, rtg_batch = batch

            ################# TRPO UPDATE ################
            policy_update(obs_batch, act_batch, adv_batch, rtg_batch, it)
            std_vals = sess.run(log_std)
            print('Log Std policy inner:', np.mean(std_vals))
            if np.mean(std_vals) < -5:
                converged_flag = True
            # Testing the policy on a real environment
            # mn_test, mn_test_std, mn_length = test_agent(env_test, action_op, num_games=1)
            # plot_results(env_test, 'ME-TRPO')
            # print(' Test score: ', np.round(mn_test, 2), np.round(mn_test_std, 2), np.round(mn_length, 2))
            # mn_test, mn_test_std, mn_length = test_agent(env_test, action_op, num_games=1)
            # summary = tf.Summary()
            # summary.value.add(tag='test/performance', simple_value=mn_test)
            # file_writer.add_summary(summary, step_count)
            # file_writer.flush()

            # Test the policy on simulated environment.
            # dynamic_wait_time_count = dynamic_wait_time(ep)
            if dynamic_wait_time(it, ep):

                print('Iterations: ', total_iterations)
                label = f'Total {total_iterations}, ' + \
                        f'data points: {len(model_buffer.train_idx) + len(model_buffer.valid_idx)}, ' + \
                        f'ep: {ep}, it: {it}\n' + hyp_str_all

                # # For Niky:  perform test! -----------------------------
                # env_test.env.set_usage('test')
                # mn_test, mn_test_std, mn_length, mn_success = test_agent(env_test, action_op, num_games=50)
                # # perform test! -----------------------------
                #
                # # plot results of test -----------------------------
                # plot_results(env_test, label=label)
                #
                # # saving test buffer -----------------------------
                # # env_test.save_current_buffer(info=label)
                # print(' Test score: ', np.round(mn_test, 2), np.round(mn_test_std, 2),
                #       np.round(mn_length, 2), np.round(mn_success, 2))
                #
                # # save the data for plotting the tests
                # tests_all.append(mn_test)
                # tests_std_all.append(mn_test_std)
                #
                # # perform test end! -----------------------------
                #
                #

                env_test.env.set_usage('default')

                # Start of test on the simulation:
                print('Simulated test:', end=' ** ')

                sim_rewards = []
                for i in range(num_ensemble_models):
                    sim_m_env = NetworkEnv(make_env(), model_op, None, i + 1)
                    mn_sim_rew, _, _, _ = test_agent(sim_m_env, action_op, num_games=25)
                    mn_sim_rew = np.round(mn_sim_rew, 2)
                    sim_rewards.append(mn_sim_rew)
                    print(mn_sim_rew, end=' ** ')

                print("")
                # Stop of test on the simulation

                # collecting/summarizing the data for plotting the progress -------------------
                entropy_all.append(np.mean(std_vals))
                step_counts_all.append(step_count)

                sim_rewards = np.array(sim_rewards)
                sim_rewards_all.append(np.mean(sim_rewards))
                sim_rewards_std_all.append(np.std(sim_rewards))

                batch_rews_all.append(np.mean(batch_rew))
                batch_rews_std_all.append(np.std(batch_rew))

                data = dict(sim_rewards_all=[sim_rewards_all, sim_rewards_std_all],
                            entropy_all=entropy_all,
                            step_counts_all=step_counts_all,
                            batch_rews_all=[batch_rews_all, batch_rews_std_all],
                            tests_all=[tests_all, tests_std_all], # data of test with real env
                            info=label)

                # save the data for plotting the progress -------------------
                save_data(data=data)

                # plotting the progress -------------------
                plot_observables(data=data, label=label)

                # stop training if the policy hasn't improved
                if np.sum(best_sim_test >= sim_rewards) > int(num_ensemble_models * 2/3):
                    # or (len(sim_rewards[sim_rewards >= 990]) > int(num_ensemble_models * 0.7)):
                    if it > delay_before_convergence_check and ep < num_epochs - 1:
                        # Test the entropy measure as convergence criterion
                        # if np.diff(entropy_all)[-1] < 0:
                        #     print('break')
                        #     break
                        print('Break...')
                        break
                else:
                    best_sim_test = sim_rewards
    # Final verification:
    # env_final = FelLocalEnv(tango=tango)
    # env_final.test = True
    # env.TOTAL_COUNTER = len(model_buffer.train_idx) + len(model_buffer.valid_idx)
    # mn_test, mn_test_std, mn_length = test_agent(env_final, action_op, num_games=100)
    # plot_results(env_final, 'ME-TRPO', save_name='Fermi')


    # For Niky:  perform final test! -----------------------------
    env_test.env.set_usage('final')
    mn_test, mn_test_std, mn_length, _ = test_agent(env_test, action_op, num_games=50)

    label = f'Verification : total {total_iterations}, ' + \
            f'data points: {len(model_buffer.train_idx) + len(model_buffer.valid_idx)}, ' + \
            f'ep: {ep}, it: {it}\n' + \
            f'rew: {mn_test}, std: {mn_test_std}'
    plot_results(env_test, label=label)

    env_test.save_current_buffer(info=label)

    # For Niky:  end final test! -----------------------------

    # env_test.env.set_usage('default')

    # closing environments..
    for env in envs:
        env.close()
    file_writer.close()


if __name__ == '__main__':
    METRPO('', hidden_sizes=hidden_sizes, cr_lr=cr_lr, gamma=gamma, lam=lam, num_epochs=num_epochs,
           steps_per_env=steps_per_env,
           number_envs=1, critic_iter=critic_iter, delta=delta, algorithm='AETRPO', conj_iters=conj_iters,
           minibatch_size=minibatch_size,
           mb_lr_start=mb_lr, model_batch_size=model_batch_size, simulated_steps=simulated_steps,
           num_ensemble_models=num_ensemble_models, model_iter=model_iter, init_random_steps=init_random_steps,
           method=method)
    # plot the results

# important notes:
# Storage
# Hyperparameters
# Scaling

# Changes:
# No init steps and less step per env 31 instead of 51 and the number of iterations is dynamic

import os
import pickle
from datetime import datetime

import gym
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import tensorflow as tf

from laser_trajectory_control_env import LaserTrajectoryControlEnv
from tango_connection import TangoConnection

# set random seed
random_seed = 111
np.random.seed(random_seed)


conf_file = '/home/niky/FERMI/2020_10_06/configuration/conf_fel2.json'
tango = TangoConnection(conf_file=conf_file)
real_env = LaserTrajectoryControlEnv(tango=tango)

# Todo: Add info to buffer and test ensemble advantage
# We will take the avarage cumulative reward and the entropy of the trajectories as an observable,
# as well as the reward on the measured data during the data collection


# For Niky Hyper parameters:
#------------------------------------------------------
steps_per_env = 15
init_random_steps = 50
# num_epochs = int((300 - init_random_steps) / (steps_per_env)) + 1
num_epochs = 20
# num_epochs = int((125-init_random_steps)/(steps_per_env))
print('Number of epoches: ', num_epochs)

hidden_sizes = [100, 100]
cr_lr = 1e-3
gamma = 0.9999
lam = 0.95

max_training_iterations = 50

delay_before_convergence_check = 1

# number_envs = 1
critic_iter = 40
delta = 0.01
algorithm = 'TRPO'
conj_iters = 15
minibatch_size = 100
simulated_steps = 1000

mb_lr = 1e-3
model_batch_size = 5  # train the dynamics model

num_ensemble_models = 3
model_iter = 25

# How often to check the progress of the network training
# e.g. lambda it, episode: (it + 1) % max(3, (ep+1)*2) == 0
# dynamic_wait_time = lambda it, ep: True  # (it + 1) % 1 == 0  #
dynamic_wait_time = lambda it, ep: (it + 1) % 2 == 0  #

# Set the priors for the anchor method:
#  TODO: How to set these correctly?
init_params = dict(init_stddev_1_w=np.sqrt(1),
                   init_stddev_1_b=np.sqrt(1),
                   init_stddev_2_w=1 / np.sqrt(100))

data_noise = 1e-4  # estimated noise variance

lambda_anchor = data_noise / (np.array([init_params['init_stddev_1_w'],
                                        init_params['init_stddev_1_b'],
                                        init_params['init_stddev_2_w']]) ** 2)

training_steps_model = 4000

# method = 'random'
method = 'pessimistic'

# Create the logging directory:
project_directory = 'Data_logging/AE_TRPO/at_1/'

#------------------------------------------------------
#------------------------------------------------------
#------------------------------------------------------


hyp_str_all = '-nr_steps_' + str(steps_per_env) + '-cr_lr' + str(cr_lr) + '-crit_it_' + str(
    critic_iter) + '-d_' + str(delta) + '-conj_iters_' + str(conj_iters) + '-n_ep_' + str(num_epochs) + \
              '-mini_bs_' + str(minibatch_size) + '-m_bs_' + str(model_batch_size) + \
              '-mb_lr_' + str(mb_lr) + \
              '-sim_steps_' + str(simulated_steps) + \
              '-m_iter_' + str(model_iter) + '-ensnr_' + str(num_ensemble_models) + '-init_' + str(
    init_random_steps) + '/'
project_directory = project_directory + hyp_str_all

# To label the plots:
hyp_str_all = '-nr_steps_' + str(steps_per_env) + '-cr_lr' + str(cr_lr) + '-crit_it_' + str(
    critic_iter) + '-d_' + str(delta) + '-conj_iters_' + str(conj_iters) + '-n_ep_' + str(num_epochs) + \
              '\n-mini_bs_' + str(minibatch_size) + '-m_bs_' + str(model_batch_size) + \
              '-mb_lr_' + str(mb_lr) + \
              '-sim_steps_' + str(simulated_steps) + \
              '-m_iter_' + str(model_iter) + \
              '\n-ensnr_' + str(num_ensemble_models)
if not os.path.isdir(project_directory):
    os.makedirs(project_directory)
    print("created folder : ", project_directory)


# Class for data storage during the tests
class TrajectoryBuffer():
    '''Class for data storage during the tests'''

    def __init__(self, name, directory):
        self.save_frequency = 100000
        self.directory = directory
        self.name = name
        self.rews = []
        self.obss = []
        self.acts = []
        self.dones = []
        self.info = ""
        self.idx = -1

    def new_trajectory(self, obs):
        self.idx += 1
        self.rews.append([])
        self.acts.append([])
        self.obss.append([])
        self.dones.append([])
        self.store_step(obs=obs)

    def store_step(self, obs=None, act=None, rew=None, done=None):
        self.rews[self.idx].append(rew)
        self.obss[self.idx].append(obs)
        self.acts[self.idx].append(act)
        self.dones[self.idx].append(done)

        if self.__len__() % self.save_frequency == 0:
            self.save_buffer()

    def __len__(self):
        assert (len(self.rews) == len(self.obss) == len(self.acts) == len(self.dones))
        return len(self.obss)

    def save_buffer(self, **kwargs):
        if 'info' in kwargs:
            self.info = kwargs.get('info')
        now = datetime.now()
        # clock_time = "{}_{}_{}_{}_".format(now.day, now.hour, now.minute, now.second)
        clock_time = f'{now.month:0>2}_{now.day:0>2}_{now.hour:0>2}_{now.minute:0>2}_{now.second:0>2}_'
        data = dict(obss=self.obss,
                    acts=self.acts,
                    rews=self.rews,
                    dones=self.dones,
                    info=self.info)
        # print('saving...', data)
        out_put_writer = open(self.directory + clock_time + self.name, 'wb')
        pickle.dump(data, out_put_writer, -1)
        # pickle.dump(self.actions, out_put_writer, -1)
        out_put_writer.close()

    def get_data(self):
        return dict(obss=self.obss,
                    acts=self.acts,
                    rews=self.rews,
                    dones=self.dones,
                    info=self.info)


class MonitoringEnv(gym.Wrapper):
    '''
    Gym Wrapper to store information for scaling to correct scpace and for post analysis.
    '''

    def __init__(self, env, **kwargs):
        gym.Wrapper.__init__(self, env)
        self.data_dict = dict()
        self.environment_usage = 'default'
        self.directory = project_directory
        self.data_dict[self.environment_usage] = TrajectoryBuffer(name=self.environment_usage,
                                                                  directory=self.directory)
        self.current_buffer = self.data_dict.get(self.environment_usage)

        self.test_env_flag = False
        if 'test_env' in kwargs:
            self.test_env_flag = True

    def reset(self, **kwargs):
        init_obs = self.env.reset(**kwargs)
        # print('Reset Env: ', (init_obs),10*'-- ')
        self.current_buffer.new_trajectory(init_obs)
        init_obs = self.scale_state_env(init_obs)
        # print('Reset Menv: ', (init_obs))
        return init_obs

    def step(self, action):
        # print('a', action)
        action = self.descale_action_env(action)
        # print('as', action)
        ob, reward, done, info = self.env.step(action)
        # print('Env: ', reward)
        # print('Env: ', ob, 'r:', reward, done)
        self.current_buffer.store_step(obs=ob, act=action, rew=reward, done=done)
        ob = self.scale_state_env(ob)
        reward = self.rew_scale(reward)
        # if done:
        #     print('Menv: ', ob, 'r:', reward, done)
        # print('Menv: ', reward)
        return ob, reward, done, info

    def set_usage(self, usage):
        self.environment_usage = usage
        if usage in self.data_dict:
            self.current_buffer = self.data_dict.get(usage)
        else:
            self.data_dict[self.environment_usage] = TrajectoryBuffer(name=self.environment_usage,
                                                                      directory=self.directory)
            self.current_buffer = self.data_dict.get(usage)

    def close_usage(self, usage):
        # Todo: Implement to save complete data
        self.current_buffer = self.data_dict.get(usage)
        self.current_buffer.save_buffer()

    def scale_state_env(self, ob):
        scale = (self.env.observation_space.high - self.env.observation_space.low)
        return (2 * ob - (self.env.observation_space.high + self.env.observation_space.low)) / scale

    def descale_action_env(self, act):
        scale = (self.env.action_space.high - self.env.action_space.low)
        return (scale * act + self.env.action_space.high + self.env.action_space.low) / 2

    def rew_scale(self, rew):
        # we only scale for the network training:
        if not self.test_env_flag:
            rew = rew * 2 + 1
        # if self.test_env_flag:
        #     '''Rescale reward from [-1,0] to [-1,1] for the training of the network in case of tests'''
        #     rew = rew * 2 + 1
        #     if rew < -1:
        #         print('Hallo was geht: ', rew)
        #     else:
        #         print('Okay...', rew)
        return rew

    def save_current_buffer(self, info=''):
        self.current_buffer = self.data_dict.get(self.environment_usage)
        self.current_buffer.save_buffer(info=info)
        print('Saved current buffer', self.environment_usage)

    def set_directory(self, directory):
        self.directory = directory


env_monitored = MonitoringEnv(env=real_env)


def make_env(**kwargs):
    '''Create the environement'''
    return MonitoringEnv(env=real_env, **kwargs)


def mlp(x, hidden_layers, output_layer, activation, last_activation=None, init_params=dict()):
    '''
    Multi-layer perceptron with init conditions for anchor method
    '''
    if init_params:
        layer_1_w = tf.layers.Dense(hidden_layers[0],
                                    activation=tf.tanh,  # try using tf.nn.relu, tf.erf, tf.nn.tanh etc.
                                    kernel_initializer=tf.random_normal_initializer(mean=0.,
                                                                                    stddev=init_params.get(
                                                                                        'init_stddev_1_w')),
                                    bias_initializer=tf.random_normal_initializer(mean=0.,
                                                                                  stddev=init_params.get(
                                                                                      'init_stddev_1_b')))

        layer_1 = layer_1_w.apply(x)

        layer_2_w = tf.layers.Dense(hidden_layers[1],
                                    activation=tf.tanh,  # try using tf.nn.relu, tf.erf, tf.nn.tanh etc.
                                    kernel_initializer=tf.random_normal_initializer(mean=0.,
                                                                                    stddev=init_params.get(
                                                                                        'init_stddev_1_w')),
                                    bias_initializer=tf.random_normal_initializer(mean=0.,
                                                                                  stddev=init_params.get(
                                                                                      'init_stddev_1_b')))

        layer_2 = layer_2_w.apply(layer_1)

        output_w = tf.layers.Dense(output_layer,
                                   activation=tf.tanh,
                                   use_bias=False,
                                   kernel_initializer=tf.random_normal_initializer(mean=0.,
                                                                                   stddev=init_params.get(
                                                                                       'init_stddev_2_w')))

        return output_w.apply(layer_2)
    else:
        for l in hidden_layers:
            x = tf.layers.dense(x, units=l, activation=activation)
        return tf.layers.dense(x, units=output_layer, activation=last_activation)


def softmax_entropy(logits):
    '''
    Softmax Entropy
    '''
    return -tf.reduce_sum(tf.nn.softmax(logits, axis=-1) * tf.nn.log_softmax(logits, axis=-1), axis=-1)


def gaussian_log_likelihood(ac, mean, log_std):
    '''
    Gaussian Log Likelihood
    '''
    log_p = ((ac - mean) ** 2 / (tf.exp(log_std) ** 2 + 1e-9) + 2 * log_std) + np.log(2 * np.pi)
    return -0.5 * tf.reduce_sum(log_p, axis=-1)


def conjugate_gradient(A, b, x=None, iters=10):
    '''
    Conjugate gradient method: approximate the solution of Ax=b
    It solve Ax=b without forming the full matrix, just compute the matrix-vector product (The Fisher-vector product)
    NB: A is not the full matrix but is a useful matrix-vector product between the averaged Fisher information matrix and arbitrary vectors
    Descibed in Appendix C.1 of the TRPO paper
    '''
    if x is None:
        x = np.zeros_like(b)

    r = A(x) - b
    p = -r
    for _ in range(iters):
        a = np.dot(r, r) / (np.dot(p, A(p)) + 1e-8)
        x += a * p
        r_n = r + a * A(p)
        b = np.dot(r_n, r_n) / (np.dot(r, r) + 1e-8)
        p = -r_n + b * p
        r = r_n
    return x


def gaussian_DKL(mu_q, log_std_q, mu_p, log_std_p):
    '''
    Gaussian KL divergence in case of a diagonal covariance matrix
    '''
    return tf.reduce_mean(tf.reduce_sum(
        0.5 * (log_std_p - log_std_q + tf.exp(log_std_q - log_std_p) + (mu_q - mu_p) ** 2 / tf.exp(log_std_p) - 1),
        axis=1))


def backtracking_line_search(Dkl, delta, old_loss, p=0.8):
    '''
    Backtracking line searc. It look for a coefficient s.t. the constraint on the DKL is satisfied
    It has both to
     - improve the non-linear objective
     - satisfy the constraint
    '''
    ## Explained in Appendix C of the TRPO paper
    a = 1
    it = 0

    new_dkl, new_loss = Dkl(a)
    while (new_dkl > delta) or (new_loss > old_loss):
        a *= p
        it += 1
        new_dkl, new_loss = Dkl(a)

    return a


def GAE(rews, v, v_last, gamma=0.99, lam=0.95):
    '''
    Generalized Advantage Estimation
    '''
    assert len(rews) == len(v)
    vs = np.append(v, v_last)
    d = np.array(rews) + gamma * vs[1:] - vs[:-1]
    gae_advantage = discounted_rewards(d, 0, gamma * lam)
    return gae_advantage


def discounted_rewards(rews, last_sv, gamma):
    '''
    Discounted reward to go
    Parameters:
    ----------
    rews: list of rewards
    last_sv: value of the last state
    gamma: discount value
    '''
    rtg = np.zeros_like(rews, dtype=np.float64)
    rtg[-1] = rews[-1] + gamma * last_sv
    for i in reversed(range(len(rews) - 1)):
        rtg[i] = rews[i] + gamma * rtg[i + 1]
    return rtg


def flatten_list(tensor_list):
    '''
    Flatten a list of tensors
    '''
    return tf.concat([flatten(t) for t in tensor_list], axis=0)


def flatten(tensor):
    '''
    Flatten a tensor
    '''
    return tf.reshape(tensor, shape=(-1,))


def test_agent(env_test, agent_op, num_games=10):
    '''
    Test an agent 'agent_op', 'num_games' times
    Return mean and std
    '''
    games_r = []
    games_length = []
    games_dones = []
    for _ in range(num_games):
        d = False
        game_r = 0
        o = env_test.reset()
        game_length = 0
        while not d:
            a_s, _ = agent_op([o])
            o, r, d, _ = env_test.step(a_s)
            game_r += r
            game_length += 1

        games_r.append(game_r)
        games_length.append(game_length)
        games_dones.append(d)
    return np.mean(games_r), np.std(games_r), np.mean(games_length), np.mean(games_dones)


class Buffer():
    '''
    Class to store the experience from a unique policy
    '''

    def __init__(self, gamma=0.99, lam=0.95):
        self.gamma = gamma
        self.lam = lam
        self.adv = []
        self.ob = []
        self.ac = []
        self.rtg = []

    def store(self, temp_traj, last_sv):
        '''
        Add temp_traj values to the buffers and compute the advantage and reward to go
        Parameters:
        -----------
        temp_traj: list where each element is a list that contains: observation, reward, action, state-value
        last_sv: value of the last state (Used to Bootstrap)
        '''
        # store only if there are temporary trajectories
        if len(temp_traj) > 0:
            self.ob.extend(temp_traj[:, 0])
            rtg = discounted_rewards(temp_traj[:, 1], last_sv, self.gamma)
            self.adv.extend(GAE(temp_traj[:, 1], temp_traj[:, 3], last_sv, self.gamma, self.lam))
            self.rtg.extend(rtg)
            self.ac.extend(temp_traj[:, 2])

    def get_batch(self):
        # standardize the advantage values
        norm_adv = (self.adv - np.mean(self.adv)) / (np.std(self.adv) + 1e-10)
        return np.array(self.ob), np.array(np.expand_dims(self.ac, -1)), np.array(norm_adv), np.array(self.rtg)

    def __len__(self):
        assert (len(self.adv) == len(self.ob) == len(self.ac) == len(self.rtg))
        return len(self.ob)


class FullBuffer():
    def __init__(self):
        self.rew = []
        self.obs = []
        self.act = []
        self.nxt_obs = []
        self.done = []

        self.train_idx = []
        self.valid_idx = []
        self.full_idx = []
        self.idx = 0

    def store(self, obs, act, rew, nxt_obs, done):
        self.rew.append(rew)
        self.obs.append(obs)
        self.act.append(act)
        self.nxt_obs.append(nxt_obs)
        self.done.append(done)

        self.idx += 1

    def generate_random_dataset(self):
        rnd = np.arange(len(self.obs))
        np.random.shuffle(rnd)
        # proportion for training vs validation
        self.valid_idx = rnd[: int(len(self.obs) / 3)]
        self.train_idx = rnd[int(len(self.obs) / 3):]
        print('Train set:', len(self.train_idx), 'Valid set:', len(self.valid_idx))

    def get_training_batch(self):
        return np.array(self.obs)[self.train_idx], np.array(np.expand_dims(self.act, -1))[self.train_idx], \
               np.array(self.rew)[self.train_idx], np.array(self.nxt_obs)[self.train_idx], np.array(self.done)[
                   self.train_idx]

    def get_valid_batch(self):
        return np.array(self.obs)[self.valid_idx], np.array(np.expand_dims(self.act, -1))[self.valid_idx], \
               np.array(self.rew)[self.valid_idx], np.array(self.nxt_obs)[self.valid_idx], np.array(self.done)[
                   self.valid_idx]

    def get_full_batch(self):
        rnd = np.arange(len(self.obs))
        np.random.shuffle(rnd)
        self.full_idx = rnd
        return np.array(self.obs)[self.full_idx], np.array(np.expand_dims(self.act, -1))[self.full_idx], \
               np.array(self.rew)[self.full_idx], np.array(self.nxt_obs)[self.full_idx], np.array(self.done)[
                   self.full_idx]

    def __len__(self):
        assert (len(self.rew) == len(self.obs) == len(self.act) == len(self.nxt_obs) == len(self.done))
        return len(self.obs)


def simulate_environment(env, policy, simulated_steps):
    '''Lists to store rewards and length of the trajectories completed'''
    buffer = Buffer(0.99, 0.95)
    steps = 0
    number_episodes = 0

    while steps < simulated_steps:
        temp_buf = []
        obs = env.reset()
        number_episodes += 1
        done = False

        while not done:
            act, val = policy([obs])

            obs2, rew, done, _ = env.step([act])

            temp_buf.append([obs.copy(), rew, np.squeeze(act), np.squeeze(val)])

            obs = obs2.copy()
            steps += 1

            if done:
                buffer.store(np.array(temp_buf), 0)
                temp_buf = []

            if steps == simulated_steps:
                break

        buffer.store(np.array(temp_buf), np.squeeze(policy([obs])[1]))

    print('Sim ep:', number_episodes, end=' \n')

    return buffer.get_batch(), number_episodes


class NetworkEnv(gym.Wrapper):
    '''
    Wrapper to handle the network interaction
    '''

    def __init__(self, env, model_func=None, done_func=None, number_models=1, uncertainty_method='random'):
        gym.Wrapper.__init__(self, env)

        self.uncertainty_method = uncertainty_method
        if self.uncertainty_method not in ['random', 'pessimistic']:
            self.uncertainty_method = 'random'
            print('switched to', self.uncertainty_method, ' since ', uncertainty_method, ' is unknown.')
        self.model_func = model_func
        self.done_func = done_func
        self.number_models = number_models
        self.len_episode = 0
        # self.threshold = self.env.threshold
        # print('the threshold is: ', self.threshold)
        self.max_steps = env.max_steps

    def reset(self, **kwargs):

        # self.threshold = -0.05 * 2 + 1  # rescaled [-1,1]
        self.len_episode = 0
        self.done = False
        # kwargs['simulation'] = True
        # action = self.env.reset(**kwargs)
        if self.model_func is not None:
            obs = np.random.uniform(-1, 1, self.env.observation_space.shape)
            # print('reset', obs)
            # Todo: remove
            obs = self.env.reset()
            # obs = self.env.reset()
        else:
            # obs = self.env.reset(**kwargs)
            pass
        # Does this work?
        # self.obs = np.clip(obs, -1.0, 1.0)
        self.obs = obs.copy()
        # if self.test_phase:
        #     print('test reset', self.obs)
        # print('Reset : ',self.obs)
        return self.obs

    def step(self, action):
        if self.model_func is not None:
            # predict the next state on a random model

            if self.uncertainty_method == 'pessimistic':
                obss = []
                rews = []
                for model_idx in range(num_ensemble_models):
                    obs, rew = self.model_func(self.obs, [np.squeeze(action)], model_idx)
                    obss.append(obs)
                    rews.append(rew)
                idx = np.argmin(rew)
                # idx = np.argmax(rew)
                obs = obss[idx].copy()
                rew = rews[idx]
            else:
                obs, rew = self.model_func(self.obs, [np.squeeze(action)], np.random.randint(0, self.number_models))
            self.obs = np.clip(obs.copy(), -1.0, 1.0)
            # obs_real, rew_real, _, _ = self.env.step(action)
            # #
            # print('Diff: ', np.linalg.norm(obs - obs_real), np.linalg.norm(rew - rew_real))
            # print('Diff: ', action, obs, obs_real, rew, rew_real)
            # obs += np.random.randn(obs.shape[-1])
            # # Todo: remove
            # self.env.state = self.obs
            # done = rew > self.threshold

            self.len_episode += 1
            # print('threshold at:', self.threshold)
            # For Niky: hardcoded reward
            if rew > 0.9:  # self.threshold: TODO: to be changed
                self.done = True
                # print("Done", rew)
            if self.len_episode >= self.max_steps:
                self.done = True

            # if self.done:
            #     print('NetEnv:', self.obs, rew)
            return self.obs, (rew - 1) / 2, self.done, ""
        else:
            # self.obs, rew, done, _ = real_env.step(action)
            # return self.obs, rew, done, ""
            pass
        # return env.step(action)


class StructEnv(gym.Wrapper):
    '''
    Gym Wrapper to store information like number of steps and total reward of the last espisode.
    '''

    def __init__(self, env):
        gym.Wrapper.__init__(self, env)
        self.n_obs = self.env.reset()
        self.total_rew = 0
        self.len_episode = 0

    def reset(self, **kwargs):
        self.n_obs = self.env.reset(**kwargs)
        self.total_rew = 0
        self.len_episode = 0
        return self.n_obs.copy()

    def step(self, action):
        ob, reward, done, info = self.env.step(action)
        # print('reward in struct', reward)
        self.total_rew += reward
        self.len_episode += 1
        return ob, reward, done, info

    def get_episode_reward(self):
        return self.total_rew

    def get_episode_length(self):
        return self.len_episode


def restore_model(old_model_variables, m_variables):
    # variable used as index for restoring the actor's parameters
    it_v2 = tf.Variable(0, trainable=False)

    restore_m_params = []
    for m_v in m_variables:
        upd_m_rsh = tf.reshape(old_model_variables[it_v2: it_v2 + tf.reduce_prod(m_v.shape)], shape=m_v.shape)
        restore_m_params.append(m_v.assign(upd_m_rsh))
        it_v2 += tf.reduce_prod(m_v.shape)

    return tf.group(*restore_m_params)


def METRPO(env_name, hidden_sizes=[32, 32], cr_lr=5e-3, num_epochs=50, gamma=0.99, lam=0.95, number_envs=1,
           critic_iter=10, steps_per_env=100, delta=0.05, algorithm='TRPO', conj_iters=10, minibatch_size=1000,
           mb_lr_start=0.0001, model_batch_size=512, simulated_steps=1000, num_ensemble_models=2, model_iter=15,
           init_random_steps=steps_per_env, method='random'):
    '''
    Model Ensemble Trust Region Policy Optimization
    The states and actions are provided by the gym environement with the correct boxs.
    The reward has to be between [-1,0].
    Parameters:
    -----------
    env_name: Name of the environment
    hidden_sizes: list of the number of hidden units for each layer
    cr_lr: critic learning rate
    num_epochs: number of training epochs
    gamma: discount factor
    lam: lambda parameter for computing the GAE
    number_envs: number of "parallel" synchronous environments
        # NB: it isn't distributed across multiple CPUs
    critic_iter: Number of SGD iterations on the critic per epoch
    steps_per_env: number of steps per environment
            # NB: the total number of steps per epoch will be: steps_per_env*number_envs
    delta: Maximum KL divergence between two policies. Scalar value
    algorithm: type of algorithm. Either 'TRPO' or 'NPO'
    conj_iters: number of conjugate gradient iterations
    minibatch_size: Batch size used to train the critic
    mb_lr: learning rate of the environment model
    model_batch_size: batch size of the environment model
    simulated_steps: number of simulated steps for each policy update
    num_ensemble_models: number of models
    model_iter: number of iterations without improvement before stopping training the model
    '''
    # TODO: add ME-TRPO hyperparameters

    tf.reset_default_graph()

    # Create a few environments to collect the trajectories

    # envs = [StructEnv(gym.make(env_name)) for _ in range(number_envs)]
    envs = [StructEnv(make_env()) for _ in range(number_envs)]
    env_test = StructEnv(make_env(test_env=True))
    # env_test = gym.make(env_name)
    print('env_test' * 4)

    # env_test = make_env(test=True)
    # env_test = gym.wrappers.Monitor(env_test, "VIDEOS/", force=True, video_callable=lambda x: x%10 == 0)
    # to be changed in real test
    # env_test = FelLocalEnv(tango=tango)
    # env_test.test = True
    # env_test_1 = FelLocalEnv(tango=tango)
    # env_test_1.test = True

    # If the scaling is not perfomed this has to be changed
    low_action_space = -1  # envs[0].action_space.low
    high_action_space = 1  # envs[0].action_space.high

    obs_dim = envs[0].observation_space.shape
    act_dim = envs[0].action_space.shape[0]

    # print(envs[0].action_space, envs[0].observation_space, low_action_space,
    #       high_action_space)

    # Placeholders for model
    act_ph = tf.placeholder(shape=(None, act_dim), dtype=tf.float64, name='act')
    obs_ph = tf.placeholder(shape=(None, obs_dim[0]), dtype=tf.float64, name='obs')
    # NEW
    nobs_ph = tf.placeholder(shape=(None, obs_dim[0]), dtype=tf.float64, name='nobs')
    rew_ph = tf.placeholder(shape=(None, 1), dtype=tf.float64, name='rew')

    ret_ph = tf.placeholder(shape=(None,), dtype=tf.float64, name='ret')
    adv_ph = tf.placeholder(shape=(None,), dtype=tf.float64, name='adv')
    old_p_log_ph = tf.placeholder(shape=(None,), dtype=tf.float64, name='old_p_log')
    old_mu_ph = tf.placeholder(shape=(None, act_dim), dtype=tf.float64, name='old_mu')
    old_log_std_ph = tf.placeholder(shape=(act_dim), dtype=tf.float64, name='old_log_std')
    p_ph = tf.placeholder(shape=(None,), dtype=tf.float64, name='p_ph')

    # Placeholder for learning rate
    mb_lr_ = tf.placeholder("float", None)  # , name='mb_lr')

    # result of the conjugate gradient algorithm
    cg_ph = tf.placeholder(shape=(None,), dtype=tf.float64, name='cg')

    #########################################################
    ######################## POLICY #########################
    #########################################################

    old_model_variables = tf.placeholder(shape=(None,), dtype=tf.float64, name='old_model_variables')

    # Neural network that represent the policy
    with tf.variable_scope('actor_nn'):
        p_means = mlp(obs_ph, hidden_sizes, act_dim, tf.tanh, last_activation=tf.tanh)
        p_means = tf.clip_by_value(p_means, low_action_space, high_action_space)
        log_std = tf.get_variable(name='log_std', initializer=np.ones(act_dim, dtype=np.float64))

    # Neural network that represent the value function
    with tf.variable_scope('critic_nn'):
        s_values = mlp(obs_ph, hidden_sizes, 1, tf.tanh, last_activation=None)
        s_values = tf.squeeze(s_values)

        # Add "noise" to the predicted mean following the Gaussian distribution with standard deviation e^(log_std)
    p_noisy = p_means + tf.random_normal(tf.shape(p_means), 0, 1, dtype=tf.float64) * tf.exp(log_std)
    # Clip the noisy actions
    a_sampl = tf.clip_by_value(p_noisy, low_action_space, high_action_space)

    # Compute the gaussian log likelihood
    p_log = gaussian_log_likelihood(act_ph, p_means, log_std)

    # Measure the divergence
    diverg = tf.reduce_mean(tf.exp(old_p_log_ph - p_log))

    # ratio
    ratio_new_old = tf.exp(p_log - old_p_log_ph)
    # TRPO surrogate loss function
    p_loss = - tf.reduce_mean(ratio_new_old * adv_ph)

    # MSE loss function
    v_loss = tf.reduce_mean((ret_ph - s_values) ** 2)
    # Critic optimization
    v_opt = tf.train.AdamOptimizer(cr_lr).minimize(v_loss)

    def variables_in_scope(scope):
        # get all trainable variables in 'scope'
        return tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, scope)

    # Gather and flatten the actor parameters
    p_variables = variables_in_scope('actor_nn')
    p_var_flatten = flatten_list(p_variables)

    # Gradient of the policy loss with respect to the actor parameters
    p_grads = tf.gradients(p_loss, p_variables)
    p_grads_flatten = flatten_list(p_grads)

    ########### RESTORE ACTOR PARAMETERS ###########
    p_old_variables = tf.placeholder(shape=(None,), dtype=tf.float64, name='p_old_variables')
    # variable used as index for restoring the actor's parameters
    it_v1 = tf.Variable(0, trainable=False)
    restore_params = []

    for p_v in p_variables:
        upd_rsh = tf.reshape(p_old_variables[it_v1: it_v1 + tf.reduce_prod(p_v.shape)], shape=p_v.shape)
        restore_params.append(p_v.assign(upd_rsh))
        it_v1 += tf.reduce_prod(p_v.shape)

    restore_params = tf.group(*restore_params)

    # gaussian KL divergence of the two policies
    dkl_diverg = gaussian_DKL(old_mu_ph, old_log_std_ph, p_means, log_std)

    # Jacobian of the KL divergence (Needed for the Fisher matrix-vector product)
    dkl_diverg_grad = tf.gradients(dkl_diverg, p_variables)

    dkl_matrix_product = tf.reduce_sum(flatten_list(dkl_diverg_grad) * p_ph)
    print('dkl_matrix_product', dkl_matrix_product.shape)
    # Fisher vector product
    # The Fisher-vector product is a way to compute the A matrix without the need of the full A
    Fx = flatten_list(tf.gradients(dkl_matrix_product, p_variables))

    ## Step length
    beta_ph = tf.placeholder(shape=(), dtype=tf.float64, name='beta')
    # NPG update
    npg_update = beta_ph * cg_ph

    ## alpha is found through line search
    alpha = tf.Variable(1., trainable=False, dtype=tf.float64)
    # TRPO update
    trpo_update = alpha * npg_update

    ####################   POLICY UPDATE  ###################
    # variable used as an index
    it_v = tf.Variable(0, trainable=False)
    p_opt = []
    # Apply the updates to the policy
    for p_v in p_variables:
        print(p_v)
        upd_rsh = tf.reshape(trpo_update[it_v: it_v + tf.reduce_prod(p_v.shape)], shape=p_v.shape)
        p_opt.append(p_v.assign_sub(upd_rsh))
        it_v += tf.reduce_prod(p_v.shape)

    p_opt = tf.group(*p_opt)

    #########################################################
    ######################### MODEL #########################
    #########################################################

    # Create a new class for the model:
    # NN class
    class NN:
        def __init__(self, x, y, y_dim, hidden_size, n, learning_rate, init_params):
            self.init_params = init_params

            # set up NN
            with tf.variable_scope('model_' + str(n) + '_nn'):
                self.inputs = x
                self.y_target = y

                self.layer_1_w = tf.layers.Dense(hidden_size,
                                                 activation=tf.tanh,
                                                 kernel_initializer=tf.random_normal_initializer(mean=0.,
                                                                                                 stddev=self.init_params.get(
                                                                                                     'init_stddev_1_w')),
                                                 bias_initializer=tf.random_normal_initializer(mean=0.,
                                                                                               stddev=self.init_params.get(
                                                                                                   'init_stddev_1_b')))

                self.layer_1 = self.layer_1_w.apply(self.inputs)

                self.layer_2_w = tf.layers.Dense(hidden_size,
                                                 activation=tf.tanh,
                                                 kernel_initializer=tf.random_normal_initializer(mean=0.,
                                                                                                 stddev=self.init_params.get(
                                                                                                     'init_stddev_1_w')),
                                                 bias_initializer=tf.random_normal_initializer(mean=0.,
                                                                                               stddev=self.init_params.get(
                                                                                                   'init_stddev_1_b')))

                self.layer_2 = self.layer_2_w.apply(self.layer_1)

                self.output_w = tf.layers.Dense(y_dim,
                                                activation=tf.tanh,
                                                use_bias=False,
                                                kernel_initializer=tf.random_normal_initializer(mean=0.,
                                                                                                stddev=self.init_params.get(
                                                                                                    'init_stddev_2_w')))
                # self.layer_1_w = tf.layers.Dense(hidden_size,
                #                                  activation=tf.tanh
                #                                  )
                #
                # self.layer_1 = self.layer_1_w.apply(self.inputs)
                #
                # self.layer_2_w = tf.layers.Dense(hidden_size,
                #                                  activation=tf.tanh)
                #
                # self.layer_2 = self.layer_2_w.apply(self.layer_1)
                #
                # self.output_w = tf.layers.Dense(y_dim,
                #                                 activation=tf.tanh)
                # #

                self.output = self.output_w.apply(self.layer_2)

                # set up loss and optimiser - we'll modify this later with anchoring regularisation
                self.opt_method = tf.train.AdamOptimizer(learning_rate)
                self.mse_ = 1 / tf.shape(self.inputs, out_type=tf.int64)[0] * \
                            tf.reduce_sum(tf.square(self.y_target - self.output))
                self.loss_ = 1 / tf.shape(self.inputs, out_type=tf.int64)[0] * \
                             tf.reduce_sum(tf.square(self.y_target - self.output))
                self.optimizer = self.opt_method.minimize(self.loss_)
                self.optimizer_mse = self.opt_method.minimize(self.mse_)
                # return self.mse_

        def get_weights(self):
            '''method to return current params'''

            ops = [self.layer_1_w.kernel, self.layer_1_w.bias,
                   self.layer_2_w.kernel, self.layer_2_w.bias,
                   self.output_w.kernel]
            w1, b1, w2, b2, w = sess.run(ops)

            return w1, b1, w2, b2, w

        # def set_weights(self, variables):
        #     model_assign(self.n, initial_variables_models[model_idx])

        def anchor(self, lambda_anchor):
            '''regularise around initialised parameters'''

            w1, b1, w2, b2, w = self.get_weights()

            # get initial params to hold for future trainings
            self.w1_init, self.b1_init, self.w2_init, self.b2_init, self.w_out_init = w1, b1, w2, b2, w

            loss_anchor = lambda_anchor[0] * tf.reduce_sum(tf.square(self.w1_init - self.layer_1_w.kernel))
            loss_anchor += lambda_anchor[1] * tf.reduce_sum(tf.square(self.b1_init - self.layer_1_w.bias))

            loss_anchor = lambda_anchor[0] * tf.reduce_sum(tf.square(self.w2_init - self.layer_2_w.kernel))
            loss_anchor += lambda_anchor[1] * tf.reduce_sum(tf.square(self.b2_init - self.layer_2_w.bias))

            loss_anchor += lambda_anchor[2] * tf.reduce_sum(tf.square(self.w_out_init - self.output_w.kernel))

            # combine with original loss
            # norm_val = 1/tf.shape(self.inputs)[0]
            self.loss_ = self.loss_ + tf.scalar_mul(1 / tf.shape(self.inputs)[0], loss_anchor)
            # self.loss_ = self.loss_ + tf.scalar_mul(1 / 1000, loss_anchor)
            self.optimizer = self.opt_method.minimize(self.loss_)
            return self.optimizer, self.loss_

    m_opts = []
    m_losses = []

    nobs_pred_m = []
    act_obs = tf.concat([obs_ph, act_ph], 1)
    target = tf.concat([nobs_ph, rew_ph], 1)

    # computational graph of N models and the correct losses for the anchor method
    m_classes = []
    for i in range(num_ensemble_models):
        # TODO: Add variable size of network
        hidden_sizes = [100, 100]
        nobs_pred = mlp(x=act_obs, hidden_layers=hidden_sizes, output_layer=obs_dim[0] + 1,
                        activation=tf.tanh, last_activation=tf.tanh)  # , init_params=init_params)
        m_class = NN(x=tf.concat([obs_ph, act_ph], 1), y=tf.concat([nobs_ph, rew_ph], 1), y_dim=obs_dim[0] + 1,
                     learning_rate=1e-3, n=i,
                     hidden_size=100, init_params=init_params)

        nobs_pred = m_class.output

        nobs_pred_m.append(nobs_pred)

        # m_loss = tf.reduce_mean((tf.concat([nobs_ph, rew_ph], 1) - nobs_pred) ** 2)
        # m_opts.append(tf.train.AdamOptimizer(learning_rate=mb_lr_).minimize(m_loss))
        # m_losses.append(m_loss)

        m_classes.append(m_class)
        # m_losses.append(m_class.mse_)
        # m_opts.append(m_class.optimizer_mse)

    ##################### RESTORE MODEL ######################
    initialize_models = []
    models_variables = []
    for i in range(num_ensemble_models):
        m_variables = variables_in_scope('model_' + str(i) + '_nn')
        # print('m_variables', m_variables)
        initialize_models.append(restore_model(old_model_variables, m_variables))
        #  List of weights as numpy
        models_variables.append(flatten_list(m_variables))

    #########################################################
    ##################### END MODEL #########################
    #########################################################
    # Time
    now = datetime.now()
    clock_time = "{}_{}_{}_{}".format(now.day, now.hour, now.minute, now.second)
    print('Time:', clock_time)

    # Set scalars and hisograms for TensorBoard
    tf.summary.scalar('p_loss', p_loss, collections=['train'])
    tf.summary.scalar('v_loss', v_loss, collections=['train'])
    tf.summary.scalar('p_divergence', diverg, collections=['train'])
    tf.summary.scalar('ratio_new_old', tf.reduce_mean(ratio_new_old), collections=['train'])
    tf.summary.scalar('dkl_diverg', dkl_diverg, collections=['train'])
    tf.summary.scalar('alpha', alpha, collections=['train'])
    tf.summary.scalar('beta', beta_ph, collections=['train'])
    tf.summary.scalar('p_std_mn', tf.reduce_mean(tf.exp(log_std)), collections=['train'])
    tf.summary.scalar('s_values_mn', tf.reduce_mean(s_values), collections=['train'])
    tf.summary.histogram('p_log', p_log, collections=['train'])
    tf.summary.histogram('p_means', p_means, collections=['train'])
    tf.summary.histogram('s_values', s_values, collections=['train'])
    tf.summary.histogram('adv_ph', adv_ph, collections=['train'])
    tf.summary.histogram('log_std', log_std, collections=['train'])
    scalar_summary = tf.summary.merge_all('train')

    tf.summary.scalar('old_v_loss', v_loss, collections=['pre_train'])
    tf.summary.scalar('old_p_loss', p_loss, collections=['pre_train'])
    pre_scalar_summary = tf.summary.merge_all('pre_train')

    hyp_str = '-spe_' + str(steps_per_env) + '-envs_' + str(number_envs) + '-cr_lr' + str(cr_lr) + '-crit_it_' + str(
        critic_iter) + '-delta_' + str(delta) + '-conj_iters_' + str(conj_iters)

    file_writer = tf.summary.FileWriter('log_dir/' + env_name + '/' + algorithm + '_' + clock_time + '_' + hyp_str,
                                        tf.get_default_graph())

    #################################################################################################
    # Session start!!!!!!!!
    # create a session
    sess = tf.Session()
    # initialize the variables
    sess.run(tf.global_variables_initializer())

    def action_op(o):
        return sess.run([p_means, s_values], feed_dict={obs_ph: o})

    def action_op_noise(o):
        return sess.run([a_sampl, s_values], feed_dict={obs_ph: o})

    def model_op(o, a, md_idx):
        mo = sess.run(nobs_pred_m[md_idx], feed_dict={obs_ph: [o], act_ph: [a[0]]})
        return np.squeeze(mo[:, :-1]), np.squeeze(mo[:, -1])

    def run_model_loss(model_idx, r_obs, r_act, r_nxt_obs, r_rew):
        # print({'obs_ph': r_obs.shape, 'act_ph': r_act.shape, 'nobs_ph': r_nxt_obs.shape})
        r_act = np.squeeze(r_act, axis=2)
        # print(r_act.shape)
        r_rew = np.reshape(r_rew, (-1, 1))
        # print(r_rew.shape)
        return_val = sess.run(m_loss_anchor[model_idx],
                              feed_dict={obs_ph: r_obs, act_ph: r_act, nobs_ph: r_nxt_obs, rew_ph: r_rew})
        # return_val = sess.run(m_losses[model_idx],
        #                       feed_dict={obs_ph: r_obs, act_ph: r_act, nobs_ph: r_nxt_obs, rew_ph: r_rew})
        return return_val

    def run_model_opt_loss(model_idx, r_obs, r_act, r_nxt_obs, r_rew, mb_lr):
        r_act = np.squeeze(r_act, axis=2)
        r_rew = np.reshape(r_rew, (-1, 1))
        # return sess.run([m_opts[model_idx], m_losses[model_idx]],
        #                 feed_dict={obs_ph: r_obs, act_ph: r_act, nobs_ph: r_nxt_obs, rew_ph: r_rew, mb_lr_: mb_lr})
        return sess.run([m_opts_anchor[model_idx], m_loss_anchor[model_idx]],
                        feed_dict={obs_ph: r_obs, act_ph: r_act, nobs_ph: r_nxt_obs, rew_ph: r_rew, mb_lr_: mb_lr})

    def model_assign(i, model_variables_to_assign):
        '''
        Update the i-th model's parameters
        '''
        return sess.run(initialize_models[i], feed_dict={old_model_variables: model_variables_to_assign})

    # def anchor(model_idx):
    #     m_classes[model_idx].anchor(lambda_anchor=lambda_anchor)

    def policy_update(obs_batch, act_batch, adv_batch, rtg_batch, it):
        # log probabilities, logits and log std of the "old" policy
        # "old" policy refer to the policy to optimize and that has been used to sample from the environment
        act_batch = np.squeeze(act_batch, axis=2)
        old_p_log, old_p_means, old_log_std = sess.run([p_log, p_means, log_std],
                                                       feed_dict={obs_ph: obs_batch, act_ph: act_batch,
                                                                  adv_ph: adv_batch, ret_ph: rtg_batch})
        # get also the "old" parameters
        old_actor_params = sess.run(p_var_flatten)
        if it < 1:
            std_vals = sess.run([log_std], feed_dict={log_std: np.ones(act_dim)})
            # print(std_vals)
        # old_p_loss is later used in the line search
        # run pre_scalar_summary for a summary before the optimization
        old_p_loss, summary = sess.run([p_loss, pre_scalar_summary],
                                       feed_dict={obs_ph: obs_batch, act_ph: act_batch, adv_ph: adv_batch,
                                                  ret_ph: rtg_batch, old_p_log_ph: old_p_log})
        file_writer.add_summary(summary, step_count)

        file_writer.add_summary(summary, step_count)
        file_writer.flush()

        def H_f(p):
            '''
            Run the Fisher-Vector product on 'p' to approximate the Hessian of the DKL
            '''
            return sess.run(Fx,
                            feed_dict={old_mu_ph: old_p_means, old_log_std_ph: old_log_std, p_ph: p, obs_ph: obs_batch,
                                       act_ph: act_batch, adv_ph: adv_batch, ret_ph: rtg_batch})

        g_f = sess.run(p_grads_flatten,
                       feed_dict={old_mu_ph: old_p_means, obs_ph: obs_batch, act_ph: act_batch, adv_ph: adv_batch,
                                  ret_ph: rtg_batch, old_p_log_ph: old_p_log})
        ## Compute the Conjugate Gradient so to obtain an approximation of H^(-1)*g
        # Where H in reality isn't the true Hessian of the KL divergence but an approximation of it computed via Fisher-Vector Product (F)
        conj_grad = conjugate_gradient(H_f, g_f, iters=conj_iters)

        # Compute the step length
        beta_np = np.sqrt(2 * delta / (1e-10 + np.sum(conj_grad * H_f(conj_grad))))

        def DKL(alpha_v):
            '''
            Compute the KL divergence.
            It optimize the function to compute the DKL. Afterwards it restore the old parameters.
            '''
            sess.run(p_opt, feed_dict={beta_ph: beta_np, alpha: alpha_v, cg_ph: conj_grad, obs_ph: obs_batch,
                                       act_ph: act_batch, adv_ph: adv_batch, old_p_log_ph: old_p_log})
            a_res = sess.run([dkl_diverg, p_loss],
                             feed_dict={old_mu_ph: old_p_means, old_log_std_ph: old_log_std, obs_ph: obs_batch,
                                        act_ph: act_batch, adv_ph: adv_batch, ret_ph: rtg_batch,
                                        old_p_log_ph: old_p_log})
            sess.run(restore_params, feed_dict={p_old_variables: old_actor_params})
            return a_res

        # Actor optimization step
        # Different for TRPO or NPG
        # Backtracing line search to find the maximum alpha coefficient s.t. the constraint is valid
        best_alpha = backtracking_line_search(DKL, delta, old_p_loss, p=0.8)
        sess.run(p_opt, feed_dict={beta_ph: beta_np, alpha: best_alpha,
                                   cg_ph: conj_grad, obs_ph: obs_batch, act_ph: act_batch,
                                   adv_ph: adv_batch, old_p_log_ph: old_p_log})

        lb = len(obs_batch)
        shuffled_batch = np.arange(lb)
        np.random.shuffle(shuffled_batch)

        # Value function optimization steps
        for _ in range(critic_iter):
            # shuffle the batch on every iteration
            np.random.shuffle(shuffled_batch)
            for idx in range(0, lb, minibatch_size):
                minib = shuffled_batch[idx:min(idx + minibatch_size, lb)]
                sess.run(v_opt, feed_dict={obs_ph: obs_batch[minib], ret_ph: rtg_batch[minib]})

    def train_model(tr_obs, tr_act, tr_nxt_obs, tr_rew, v_obs, v_act, v_nxt_obs, v_rew,
                    f_obs, f_act, f_nxt_obs, f_rew,
                    step_count, model_idx, mb_lr):

        # Get validation loss on the old model only used for monitoring
        mb_valid_loss1 = run_model_loss(model_idx, v_obs, v_act, v_nxt_obs, v_rew)

        # Restore the initial random weights to have a new, clean neural network
        # initial_variables_models - list stored before already in the code below -
        # important for the anchor method
        model_assign(model_idx, initial_variables_models[model_idx])

        # update the anchor model losses:
        # m_opts_anchor = []
        # m_loss_anchor = []
        # for i in range(num_ensemble_models):
        # This stores the anchor weights to the models for the regularisation:
        # opt, loss = m_classes[model_idx].anchor(lambda_anchor=lambda_anchor)
        # m_opts_anchor[model_idx] = opt
        # m_loss_anchor[model_idx] = loss

        # Get validation loss on the now initialized model
        mb_valid_loss = run_model_loss(model_idx, v_obs, v_act, v_nxt_obs, v_rew)

        acc_m_losses = []
        last_m_losses = []
        md_params = sess.run(models_variables[model_idx])
        best_mb = {'iter': 0, 'loss': mb_valid_loss, 'params': md_params}
        iteration = 0

        # Create mini-batch for training
        lb = len(f_obs)
        shuffled_batch = np.arange(lb)
        np.random.shuffle(shuffled_batch)

        # Run until the number of model_iter has passed from the best val loss at it on...
        for iteration in range(0, training_steps_model):

            # update the model on each mini-batch
            last_m_losses = []
            # for idx in range(0, 20): # repeat gradient step for some time
            # minib = shuffled_batch[idx:min(idx + minibatch_size, lb)]
            # want all data (shuffeld)
            # np.random.shuffle(shuffled_batch)
            minib = shuffled_batch
            if len(minib) != minibatch_size:
                _, ml = run_model_opt_loss(model_idx, f_obs[minib], f_act[minib], f_nxt_obs[minib],
                                           f_rew[minib], mb_lr=mb_lr)
                acc_m_losses.append(ml)
                last_m_losses.append(ml)
            else:
                pass
                # print('Warning!')

            # Check if the loss on the validation set has improved
            mb_valid_loss = run_model_loss(model_idx, v_obs, v_act, v_nxt_obs, v_rew)
            # if mb_valid_loss < best_mb['loss']:
            best_mb['loss'] = mb_valid_loss
            best_mb['iter'] = iteration
            # store the parameters to the array
            best_mb['params'] = sess.run(models_variables[model_idx])

            # it += 1

            # if it>=10000:
            #     break
            # print('iteration: ', it)

        # Restore the model with the lower validation loss
        model_assign(model_idx, best_mb['params'])

        print('Model:{}, iter:{} -- Old Val loss:{:.6f}  '
              'New Val loss:{:.6f} -- New Train loss:{:.6f}'.format(model_idx, iteration,
                                                                    mb_valid_loss1, best_mb['loss'],
                                                                    np.mean(last_m_losses)))
        summary = tf.Summary()
        summary.value.add(tag='supplementary/m_loss', simple_value=np.mean(acc_m_losses))
        summary.value.add(tag='supplementary/iterations', simple_value=iteration)
        file_writer.add_summary(summary, step_count)
        file_writer.flush()

    def plot_results(env_wrapper, label, **kwargs):
        # plotting
        print('now plotting...')
        rewards = env_wrapper.env.current_buffer.get_data()['rews']

        # states = env_wrapper.env.current_buffer.get_data()['obss']
        # pd.DataFrame(states).hist()
        # plt.show()

        # initial_states = env.initial_conditions

        iterations = []
        finals = []
        means = []
        stds = []

        # init_states = pd.read_pickle('/Users/shirlaen/PycharmProjects/DeepLearning/spinningup/Environments/initData')

        for i in range(len(rewards)):
            if (len(rewards[i]) > 1):
                # finals.append(rewards[i][len(rewards[i]) - 1])
                finals.append(rewards[i][-1])
                means.append(np.mean(rewards[i][1:]))
                stds.append(np.std(rewards[i][1:]))
                iterations.append(len(rewards[i]))
        # print(iterations)
        x = range(len(iterations))
        iterations = np.array(iterations)
        finals = np.array(finals)
        means = np.array(means)
        stds = np.array(stds)

        plot_suffix = label  # , Fermi time: {env.TOTAL_COUNTER / 600:.1f} h'

        fig, axs = plt.subplots(2, 1, sharex=True)

        ax = axs[0]
        ax.plot(x, iterations)
        ax.set_ylabel('Iterations (1)')
        ax.set_title(plot_suffix)
        # fig.suptitle(label, fontsize=12)
        if 'data_number' in kwargs:
            ax1 = plt.twinx(ax)
            color = 'lime'
            ax1.set_ylabel('Mean reward', color=color)  # we already handled the x-label with ax1
            ax1.tick_params(axis='y', labelcolor=color)
            ax1.plot(x, kwargs.get('data_number'), color=color)

        ax = axs[1]
        color = 'blue'
        ax.set_ylabel('Final reward', color=color)  # we already handled the x-label with ax1
        ax.tick_params(axis='y', labelcolor=color)
        ax.plot(x, finals, color=color)

        ax.set_title('Final reward per episode')  # + plot_suffix)
        ax.set_xlabel('Episodes (1)')

        ax1 = plt.twinx(ax)
        color = 'lime'
        ax1.set_ylabel('Mean reward', color=color)  # we already handled the x-label with ax1
        ax1.tick_params(axis='y', labelcolor=color)
        ax1.fill_between(x, means - stds, means + stds,
                         alpha=0.5, edgecolor=color, facecolor='#FF9848')
        ax1.plot(x, means, color=color)

        # ax.set_ylim(ax1.get_ylim())
        if 'save_name' in kwargs:
            plt.savefig(kwargs.get('save_name') + '.pdf')
        # fig.tight_layout()
        plt.show()

    def plot_observables(data, label, **kwargs):
        """plot observables during the test"""

        sim_rewards_all = np.array(data.get('sim_rewards_all'))
        step_counts_all = np.array(data.get('step_counts_all'))
        batch_rews_all = np.array(data.get('batch_rews_all'))
        tests_all = np.array(data.get('tests_all'))

        fig, axs = plt.subplots(2, 1, sharex=True)
        x = np.arange(len(batch_rews_all[0]))
        ax = axs[0]
        ax.step(x, batch_rews_all[0])
        ax.fill_between(x, batch_rews_all[0] - batch_rews_all[1], batch_rews_all[0] + batch_rews_all[1],
                        alpha=0.5)
        ax.set_ylabel('rews per batch')

        ax.set_title(label)

        # plt.tw
        ax2 = ax.twinx()

        color = 'lime'
        ax2.set_ylabel('data points', color=color)  # we already handled the x-label with ax1
        ax2.tick_params(axis='y', labelcolor=color)
        ax2.step(x, step_counts_all, color=color)

        ax = axs[1]
        ax.grid(True)
        ax.plot(sim_rewards_all[0], ls=':')
        ax.fill_between(x, sim_rewards_all[0] - sim_rewards_all[1], sim_rewards_all[0] + sim_rewards_all[1],
                        alpha=0.5)

        try:
            ax.plot(tests_all[0])
            ax.fill_between(x, tests_all[0] - tests_all[1], tests_all[0] + tests_all[1],
                            alpha=0.5)
        except:
            pass

        ax.set_ylabel('rewards tests')
        # plt.tw
        ax2 = ax.twinx()

        color = 'lime'
        ax2.set_ylabel('entropy', color=color)  # we already handled the x-label with ax1
        ax2.tick_params(axis='y', labelcolor=color)
        ax2.plot(entropy_all, color=color)

        plt.show()

    def save_data(data, **kwargs):
        '''logging functon'''
        # if 'directory_name' in kwargs:
        #     project_directory = kwargs.get('directory_name')
        now = datetime.now()
        clock_time = f'{now.month:0>2}_{now.day:0>2}_{now.hour:0>2}_{now.minute:0>2}_{now.second:0>2}'
        out_put_writer = open(project_directory + clock_time + '_training_observables', 'wb')
        pickle.dump(data, out_put_writer, -1)
        out_put_writer.close()

    # variable to store the total number of steps
    step_count = 0
    model_buffer = FullBuffer()
    print('Env batch size:', steps_per_env, ' Batch size:', steps_per_env * number_envs)

    # Create a simulated environment
    sim_env = NetworkEnv(make_env(), model_op, None, number_models=num_ensemble_models, uncertainty_method=method)

    # ------------------------------------------------------------------------------------------------------
    # -------------------------------------Try to set correct anchors---------------------------------------
    # Get the initial parameters of each model
    # These are used in later epochs when we aim to re-train the models anew with the new dataset
    initial_variables_models = []
    for model_var in models_variables:
        # num list of parameters
        initial_variables_models.append(sess.run(model_var))

    # update the anchor model losses:
    m_opts_anchor = []
    m_loss_anchor = []
    for i in range(num_ensemble_models):
        # This stores the anchor weights to the models for the regularisation:
        opt, loss = m_classes[i].anchor(lambda_anchor=lambda_anchor)
        m_opts_anchor.append(opt)
        m_loss_anchor.append(loss)

    # ------------------------------------------------------------------------------------------------------
    # -------------------------------------Try to set correct anchors---------------------------------------

    total_iterations = 0

    converged_flag = False
    # save_data = save_data(clock_time)
    sim_rewards_all = []
    sim_rewards_std_all = []
    entropy_all = []
    tests_all = []
    tests_std_all = []
    batch_rews_all = []
    batch_rews_std_all = []
    step_counts_all = []



    for ep in range(num_epochs):
        if (converged_flag):
            print('Converged!!!!')
            break
        # lists to store rewards and length of the trajectories completed
        batch_rew = []
        batch_len = []
        print('============================', ep, '============================')
        # Execute in serial the environment, storing temporarily the trajectories.
        for env in envs:
            # Todo: Test randomization stronger if reward lower...we need a good scheme
            # target_threshold ?????
            init_log_std = np.ones(act_dim) * np.log(np.random.rand() * 1)
            env.reset()

            # iterate over a fixed number of steps
            steps_train = init_random_steps if ep == 0 else steps_per_env
            # steps_train = steps_per_env
            for _ in range(steps_train):
                # found = False
                # while not(found):
                # run the policy

                if ep == 0:
                    # Sample random action during the first epoch
                    act = np.random.uniform(-1, 1, size=env.action_space.shape[-1])

                else:

                    act = sess.run(a_sampl, feed_dict={obs_ph: [env.n_obs], log_std: init_log_std})
                    # act = np.clip(act + np.random.randn(act.shape[0], act.shape[1]) * 0.1, -1, 1)

                act = np.squeeze(act)
                # print('act', act*12)
                # take a step in the environment
                obs2, rew, done, _ = env.step(np.array(act))

                # add the new transition to the temporary buffer
                model_buffer.store(env.n_obs.copy(), act, rew.copy(), obs2.copy(), done)

                env.n_obs = obs2.copy()
                step_count += 1

                if done:
                    batch_rew.append(env.get_episode_reward())
                    batch_len.append(env.get_episode_length())

                    env.reset()
                    init_log_std = np.ones(act_dim) * np.log(np.random.rand() * 1)

                # if ep == 0:
                #     # try:
                #     # Initialize randomly a training and validation set
                #     model_buffer.generate_random_dataset()
                #     # get both datasets
                #     train_obs, train_act, train_rew, train_nxt_obs, _ = model_buffer.get_training_batch()
                #     valid_obs, valid_act, valid_rew, valid_nxt_obs, _ = model_buffer.get_valid_batch()
                #     target_threshold = max(max(valid_rew), max(train_rew))
                #     # print('-- '*38, target_threshold)
                #     found = target_threshold>=-0.1 and step_count>=191
                #     # except:
                #     #     pass

        # save the data for plotting the collected data for the model
        env.save_current_buffer()

        print('Ep:%d Rew:%.2f -- Step:%d' % (ep, np.mean(batch_rew), step_count))

        # env_test.env.set_usage('default')
        # plot_results(env_test, f'Total {total_iterations}, '
        #                        f'data points: {len(model_buffer.train_idx) + len(model_buffer.valid_idx)}, '
        #                        f'modelit: {ep}')
        ############################################################
        ###################### MODEL LEARNING ######################
        ############################################################

        # Initialize randomly a training and validation set
        model_buffer.generate_random_dataset()

        # get both datasets
        train_obs, train_act, train_rew, train_nxt_obs, _ = model_buffer.get_training_batch()
        valid_obs, valid_act, valid_rew, valid_nxt_obs, _ = model_buffer.get_valid_batch()
        full_obs, full_act, full_rew, full_nxt_obs, _ = model_buffer.get_full_batch()

        std_vals = sess.run(log_std)
        print('Log Std policy:', std_vals, np.mean(std_vals))

        target_threshold = max(max(valid_rew), max(train_rew))
        sim_env.threshold = target_threshold  # min(target_threshold, -0.05)
        print('maximum: ', sim_env.threshold)

        # Learning rate as function of ep
        lr = lambda ep: 1e-3 - ep / num_epochs * (1e-3 - 5e-4)
        mb_lr = 1e-3  # if ep < 1 else 1e-3
        # simulated_steps = simulated_steps if ep<10 else 10000
        print('mb_lr: ', mb_lr)
        for i in range(num_ensemble_models):
            # train the dynamic model on the datasets just sampled
            train_model(train_obs, train_act, train_nxt_obs, train_rew, valid_obs, valid_act, valid_nxt_obs, valid_rew,
                        full_obs, full_act, full_nxt_obs, full_rew,
                        step_count, i, mb_lr=mb_lr)

        ############################################################
        ###################### POLICY LEARNING ######################
        ############################################################

        best_sim_test = -1e16 * np.ones(num_ensemble_models)

        # plot_results(env_test, f'Total {total_iterations}, '
        #                        f'data points: {len(model_buffer.train_idx) + len(model_buffer.valid_idx)}, '
        #                        f'modelit: {ep}')

        # # For Niky:  perform default test! -----------------------------
        # if ep == 0:
        #     it = 0
        #     # perform test! -----------------------------
        #     env_test.env.set_usage('test')
        #     mn_test, mn_test_std, mn_length, mn_success = test_agent(env_test, action_op, num_games=50)
        #     # perform test! -----------------------------
        #     label = f'Init {total_iterations}, ' + \
        #             f'data points: {len(model_buffer.train_idx) + len(model_buffer.valid_idx)}, ' + \
        #             f'ep: {ep}, it: {it}\n' + hyp_str_all
        #     # plot results of test -----------------------------
        #     plot_results(env_test, label=label)

        env_test.env.set_usage('default')

        for it in range(max_training_iterations):
            if converged_flag:
                break
            total_iterations += 1
            print('\t Policy it', it, end='..')

            ##################### MODEL SIMLUATION #####################
            # obs_batch, act_batch, adv_batch, rtg_batch = simulate_environment(sim_env, action_op_noise, simulated_steps)
            batch, ep_length = simulate_environment(sim_env, action_op_noise, simulated_steps)
            # verification_simulate_environment(sim_env, env_test, action_op_noise, 50)
            obs_batch, act_batch, adv_batch, rtg_batch = batch

            ################# TRPO UPDATE ################
            policy_update(obs_batch, act_batch, adv_batch, rtg_batch, it)
            std_vals = sess.run(log_std)
            print('Log Std policy inner:', np.mean(std_vals))
            if np.mean(std_vals) < -5:
                converged_flag = True
            # Testing the policy on a real environment
            # mn_test, mn_test_std, mn_length = test_agent(env_test, action_op, num_games=1)
            # plot_results(env_test, 'ME-TRPO')
            # print(' Test score: ', np.round(mn_test, 2), np.round(mn_test_std, 2), np.round(mn_length, 2))
            # mn_test, mn_test_std, mn_length = test_agent(env_test, action_op, num_games=1)
            # summary = tf.Summary()
            # summary.value.add(tag='test/performance', simple_value=mn_test)
            # file_writer.add_summary(summary, step_count)
            # file_writer.flush()

            # Test the policy on simulated environment.
            # dynamic_wait_time_count = dynamic_wait_time(ep)
            if dynamic_wait_time(it, ep):

                print('Iterations: ', total_iterations)
                label = f'Total {total_iterations}, ' + \
                        f'data points: {len(model_buffer.train_idx) + len(model_buffer.valid_idx)}, ' + \
                        f'ep: {ep}, it: {it}\n' + hyp_str_all

                # # For Niky:  perform test! -----------------------------
                # env_test.env.set_usage('test')
                # mn_test, mn_test_std, mn_length, mn_success = test_agent(env_test, action_op, num_games=50)
                # # perform test! -----------------------------
                #
                # # plot results of test -----------------------------
                # plot_results(env_test, label=label)
                #
                # # saving test buffer -----------------------------
                # # env_test.save_current_buffer(info=label)
                # print(' Test score: ', np.round(mn_test, 2), np.round(mn_test_std, 2),
                #       np.round(mn_length, 2), np.round(mn_success, 2))
                #
                # # save the data for plotting the tests
                # tests_all.append(mn_test)
                # tests_std_all.append(mn_test_std)
                #
                # # perform test end! -----------------------------
                #
                #

                env_test.env.set_usage('default')

                # Start of test on the simulation:
                print('Simulated test:', end=' ** ')

                sim_rewards = []
                for i in range(num_ensemble_models):
                    sim_m_env = NetworkEnv(make_env(), model_op, None, i + 1)
                    mn_sim_rew, _, _, _ = test_agent(sim_m_env, action_op, num_games=25)
                    mn_sim_rew = np.round(mn_sim_rew, 2)
                    sim_rewards.append(mn_sim_rew)
                    print(mn_sim_rew, end=' ** ')

                print("")
                # Stop of test on the simulation

                # collecting/summarizing the data for plotting the progress -------------------
                entropy_all.append(np.mean(std_vals))
                step_counts_all.append(step_count)

                sim_rewards = np.array(sim_rewards)
                sim_rewards_all.append(np.mean(sim_rewards))
                sim_rewards_std_all.append(np.std(sim_rewards))

                batch_rews_all.append(np.mean(batch_rew))
                batch_rews_std_all.append(np.std(batch_rew))

                data = dict(sim_rewards_all=[sim_rewards_all, sim_rewards_std_all],
                            entropy_all=entropy_all,
                            step_counts_all=step_counts_all,
                            batch_rews_all=[batch_rews_all, batch_rews_std_all],
                            tests_all=[tests_all, tests_std_all], # data of test with real env
                            info=label)

                # save the data for plotting the progress -------------------
                save_data(data=data)

                # plotting the progress -------------------
                plot_observables(data=data, label=label)

                # stop training if the policy hasn't improved
                if np.sum(best_sim_test >= sim_rewards) > int(num_ensemble_models * 2/3):
                    # or (len(sim_rewards[sim_rewards >= 990]) > int(num_ensemble_models * 0.7)):
                    if it > delay_before_convergence_check and ep < num_epochs - 1:
                        # Test the entropy measure as convergence criterion
                        # if np.diff(entropy_all)[-1] < 0:
                        #     print('break')
                        #     break
                        print('Break...')
                        break
                else:
                    best_sim_test = sim_rewards
    # Final verification:
    # env_final = FelLocalEnv(tango=tango)
    # env_final.test = True
    # env.TOTAL_COUNTER = len(model_buffer.train_idx) + len(model_buffer.valid_idx)
    # mn_test, mn_test_std, mn_length = test_agent(env_final, action_op, num_games=100)
    # plot_results(env_final, 'ME-TRPO', save_name='Fermi')


    # For Niky:  perform final test! -----------------------------
    env_test.env.set_usage('final')
    mn_test, mn_test_std, mn_length, _ = test_agent(env_test, action_op, num_games=50)

    label = f'Verification : total {total_iterations}, ' + \
            f'data points: {len(model_buffer.train_idx) + len(model_buffer.valid_idx)}, ' + \
            f'ep: {ep}, it: {it}\n' + \
            f'rew: {mn_test}, std: {mn_test_std}'
    plot_results(env_test, label=label)

    env_test.save_current_buffer(info=label)

    # For Niky:  end final test! -----------------------------

    # env_test.env.set_usage('default')

    # closing environments..
    for env in envs:
        env.close()
    file_writer.close()


if __name__ == '__main__':
    METRPO('', hidden_sizes=hidden_sizes, cr_lr=cr_lr, gamma=gamma, lam=lam, num_epochs=num_epochs,
           steps_per_env=steps_per_env,
           number_envs=1, critic_iter=critic_iter, delta=delta, algorithm='AETRPO', conj_iters=conj_iters,
           minibatch_size=minibatch_size,
           mb_lr_start=mb_lr, model_batch_size=model_batch_size, simulated_steps=simulated_steps,
           num_ensemble_models=num_ensemble_models, model_iter=model_iter, init_random_steps=init_random_steps,
           method=method)
    # plot the results

# important notes:
# Storage
# Hyperparameters
# Scaling

# Changes:
# No init steps and less step per env 31 instead of 51 and the number of iterations is dynamic

# set random seed # Todo: Add info to buffer and test ensemble advantage # We will take the avarage cumulative reward and the entropy of the trajectories as an observable, # as well as the reward on the measured data during the data collection # For Niky Hyper parameters: #------------------------------------------------------ # num_epochs = int((300 - init_random_steps) / (steps_per_env)) + 1 # num_epochs = int((125-init_random_steps)/(steps_per_env)) # number_envs = 1 # train the dynamics model # How often to check the progress of the network training # e.g. lambda it, episode: (it + 1) % max(3, (ep+1)*2) == 0 # dynamic_wait_time = lambda it, ep: True  # (it + 1) % 1 == 0  # # # Set the priors for the anchor method: #  TODO: How to set these correctly? # estimated noise variance # method = 'random' # Create the logging directory: #------------------------------------------------------ #------------------------------------------------------ #------------------------------------------------------ # To label the plots: # Class for data storage during the tests Class for data storage during the tests # clock_time = "{}_{}_{}_{}_".format(now.day, now.hour, now.minute, now.second) # print('saving...', data) # pickle.dump(self.actions, out_put_writer, -1) Gym Wrapper to store information for scaling to correct scpace and for post analysis. # print('Reset Env: ', (init_obs),10*'-- ') # print('Reset Menv: ', (init_obs)) # print('a', action) # print('as', action) # print('Env: ', reward) # print('Env: ', ob, 'r:', reward, done) # if done: #     print('Menv: ', ob, 'r:', reward, done) # print('Menv: ', reward) # Todo: Implement to save complete data # we only scale for the network training: # if self.test_env_flag: #     '''Rescale reward from [-1,0] to [-1,1] for the training of the network in case of tests''' #     rew = rew * 2 + 1 #     if rew < -1: #         print('Hallo was geht: ', rew) #     else: #         print('Okay...', rew) Create the environement Multi-layer perceptron with init conditions for anchor method # try using tf.nn.relu, tf.erf, tf.nn.tanh etc. # try using tf.nn.relu, tf.erf, tf.nn.tanh etc. Softmax Entropy Gaussian Log Likelihood Conjugate gradient method: approximate the solution of Ax=b     It solve Ax=b without forming the full matrix, just compute the matrix-vector product (The Fisher-vector product)     NB: A is not the full matrix but is a useful matrix-vector product between the averaged Fisher information matrix and arbitrary vectors     Descibed in Appendix C.1 of the TRPO paper Gaussian KL divergence in case of a diagonal covariance matrix Backtracking line searc. It look for a coefficient s.t. the constraint on the DKL is satisfied     It has both to      - improve the non-linear objective      - satisfy the constraint ## Explained in Appendix C of the TRPO paper Generalized Advantage Estimation Discounted reward to go     Parameters:     ----------     rews: list of rewards     last_sv: value of the last state     gamma: discount value Flatten a list of tensors Flatten a tensor Test an agent 'agent_op', 'num_games' times     Return mean and std Class to store the experience from a unique policy Add temp_traj values to the buffers and compute the advantage and reward to go         Parameters:         -----------         temp_traj: list where each element is a list that contains: observation, reward, action, state-value         last_sv: value of the last state (Used to Bootstrap) # store only if there are temporary trajectories # standardize the advantage values # proportion for training vs validation Lists to store rewards and length of the trajectories completed Wrapper to handle the network interaction # self.threshold = self.env.threshold # print('the threshold is: ', self.threshold) # self.threshold = -0.05 * 2 + 1  # rescaled [-1,1] # kwargs['simulation'] = True # action = self.env.reset(**kwargs) # print('reset', obs) # Todo: remove # obs = self.env.reset() # obs = self.env.reset(**kwargs) # Does this work? # self.obs = np.clip(obs, -1.0, 1.0) # if self.test_phase: #     print('test reset', self.obs) # print('Reset : ',self.obs) # predict the next state on a random model # idx = np.argmax(rew) # obs_real, rew_real, _, _ = self.env.step(action) # # # print('Diff: ', np.linalg.norm(obs - obs_real), np.linalg.norm(rew - rew_real)) # print('Diff: ', action, obs, obs_real, rew, rew_real) # obs += np.random.randn(obs.shape[-1]) # # Todo: remove # self.env.state = self.obs # done = rew > self.threshold # print('threshold at:', self.threshold) # For Niky: hardcoded reward # self.threshold: TODO: to be changed # print("Done", rew) # if self.done: #     print('NetEnv:', self.obs, rew) # self.obs, rew, done, _ = real_env.step(action) # return self.obs, rew, done, "" # return env.step(action) Gym Wrapper to store information like number of steps and total reward of the last espisode. # print('reward in struct', reward) # variable used as index for restoring the actor's parameters Model Ensemble Trust Region Policy Optimization     The states and actions are provided by the gym environement with the correct boxs.     The reward has to be between [-1,0].     Parameters:     -----------     env_name: Name of the environment     hidden_sizes: list of the number of hidden units for each layer     cr_lr: critic learning rate     num_epochs: number of training epochs     gamma: discount factor     lam: lambda parameter for computing the GAE     number_envs: number of "parallel" synchronous environments         # NB: it isn't distributed across multiple CPUs     critic_iter: Number of SGD iterations on the critic per epoch     steps_per_env: number of steps per environment             # NB: the total number of steps per epoch will be: steps_per_env*number_envs     delta: Maximum KL divergence between two policies. Scalar value     algorithm: type of algorithm. Either 'TRPO' or 'NPO'     conj_iters: number of conjugate gradient iterations     minibatch_size: Batch size used to train the critic     mb_lr: learning rate of the environment model     model_batch_size: batch size of the environment model     simulated_steps: number of simulated steps for each policy update     num_ensemble_models: number of models     model_iter: number of iterations without improvement before stopping training the model # TODO: add ME-TRPO hyperparameters # Create a few environments to collect the trajectories # envs = [StructEnv(gym.make(env_name)) for _ in range(number_envs)] # env_test = gym.make(env_name) # env_test = make_env(test=True) # env_test = gym.wrappers.Monitor(env_test, "VIDEOS/", force=True, video_callable=lambda x: x%10 == 0) # to be changed in real test # env_test = FelLocalEnv(tango=tango) # env_test.test = True # env_test_1 = FelLocalEnv(tango=tango) # env_test_1.test = True # If the scaling is not perfomed this has to be changed # envs[0].action_space.low # envs[0].action_space.high # print(envs[0].action_space, envs[0].observation_space, low_action_space, #       high_action_space) # Placeholders for model # NEW # Placeholder for learning rate # , name='mb_lr') # result of the conjugate gradient algorithm ######################################################### ######################## POLICY ######################### ######################################################### # Neural network that represent the policy # Neural network that represent the value function # Add "noise" to the predicted mean following the Gaussian distribution with standard deviation e^(log_std) # Clip the noisy actions # Compute the gaussian log likelihood # Measure the divergence # ratio # TRPO surrogate loss function # MSE loss function # Critic optimization # get all trainable variables in 'scope' # Gather and flatten the actor parameters # Gradient of the policy loss with respect to the actor parameters ########### RESTORE ACTOR PARAMETERS ########### # variable used as index for restoring the actor's parameters # gaussian KL divergence of the two policies # Jacobian of the KL divergence (Needed for the Fisher matrix-vector product) # Fisher vector product # The Fisher-vector product is a way to compute the A matrix without the need of the full A ## Step length # NPG update ## alpha is found through line search # TRPO update ####################   POLICY UPDATE  ################### # variable used as an index # Apply the updates to the policy ######################################################### ######################### MODEL ######################### ######################################################### # Create a new class for the model: # NN class # set up NN # self.layer_1_w = tf.layers.Dense(hidden_size, #                                  activation=tf.tanh #                                  ) # # self.layer_1 = self.layer_1_w.apply(self.inputs) # # self.layer_2_w = tf.layers.Dense(hidden_size, #                                  activation=tf.tanh) # # self.layer_2 = self.layer_2_w.apply(self.layer_1) # # self.output_w = tf.layers.Dense(y_dim, #                                 activation=tf.tanh) # # # set up loss and optimiser - we'll modify this later with anchoring regularisation # return self.mse_ method to return current params # def set_weights(self, variables): #     model_assign(self.n, initial_variables_models[model_idx]) regularise around initialised parameters # get initial params to hold for future trainings # combine with original loss # norm_val = 1/tf.shape(self.inputs)[0] # self.loss_ = self.loss_ + tf.scalar_mul(1 / 1000, loss_anchor) # computational graph of N models and the correct losses for the anchor method # TODO: Add variable size of network # , init_params=init_params) # m_loss = tf.reduce_mean((tf.concat([nobs_ph, rew_ph], 1) - nobs_pred) ** 2) # m_opts.append(tf.train.AdamOptimizer(learning_rate=mb_lr_).minimize(m_loss)) # m_losses.append(m_loss) # m_losses.append(m_class.mse_) # m_opts.append(m_class.optimizer_mse) ##################### RESTORE MODEL ###################### # print('m_variables', m_variables) #  List of weights as numpy ######################################################### ##################### END MODEL ######################### ######################################################### # Time # Set scalars and hisograms for TensorBoard ################################################################################################# # Session start!!!!!!!! # create a session # initialize the variables # print({'obs_ph': r_obs.shape, 'act_ph': r_act.shape, 'nobs_ph': r_nxt_obs.shape}) # print(r_act.shape) # print(r_rew.shape) # return_val = sess.run(m_losses[model_idx], #                       feed_dict={obs_ph: r_obs, act_ph: r_act, nobs_ph: r_nxt_obs, rew_ph: r_rew}) # return sess.run([m_opts[model_idx], m_losses[model_idx]], #                 feed_dict={obs_ph: r_obs, act_ph: r_act, nobs_ph: r_nxt_obs, rew_ph: r_rew, mb_lr_: mb_lr}) Update the i-th model's parameters # def anchor(model_idx): #     m_classes[model_idx].anchor(lambda_anchor=lambda_anchor) # log probabilities, logits and log std of the "old" policy # "old" policy refer to the policy to optimize and that has been used to sample from the environment # get also the "old" parameters # print(std_vals) # old_p_loss is later used in the line search # run pre_scalar_summary for a summary before the optimization Run the Fisher-Vector product on 'p' to approximate the Hessian of the DKL ## Compute the Conjugate Gradient so to obtain an approximation of H^(-1)*g # Where H in reality isn't the true Hessian of the KL divergence but an approximation of it computed via Fisher-Vector Product (F) # Compute the step length Compute the KL divergence.             It optimize the function to compute the DKL. Afterwards it restore the old parameters. # Actor optimization step # Different for TRPO or NPG # Backtracing line search to find the maximum alpha coefficient s.t. the constraint is valid # Value function optimization steps # shuffle the batch on every iteration # Get validation loss on the old model only used for monitoring # Restore the initial random weights to have a new, clean neural network # initial_variables_models - list stored before already in the code below - # important for the anchor method # update the anchor model losses: # m_opts_anchor = [] # m_loss_anchor = [] # for i in range(num_ensemble_models): # This stores the anchor weights to the models for the regularisation: # opt, loss = m_classes[model_idx].anchor(lambda_anchor=lambda_anchor) # m_opts_anchor[model_idx] = opt # m_loss_anchor[model_idx] = loss # Get validation loss on the now initialized model # Create mini-batch for training # Run until the number of model_iter has passed from the best val loss at it on... # update the model on each mini-batch # for idx in range(0, 20): # repeat gradient step for some time # minib = shuffled_batch[idx:min(idx + minibatch_size, lb)] # want all data (shuffeld) # np.random.shuffle(shuffled_batch) # print('Warning!') # Check if the loss on the validation set has improved # if mb_valid_loss < best_mb['loss']: # store the parameters to the array # it += 1 # if it>=10000: #     break # print('iteration: ', it) # Restore the model with the lower validation loss # plotting # states = env_wrapper.env.current_buffer.get_data()['obss'] # pd.DataFrame(states).hist() # plt.show() # initial_states = env.initial_conditions # init_states = pd.read_pickle('/Users/shirlaen/PycharmProjects/DeepLearning/spinningup/Environments/initData') # finals.append(rewards[i][len(rewards[i]) - 1]) # print(iterations) # , Fermi time: {env.TOTAL_COUNTER / 600:.1f} h' # fig.suptitle(label, fontsize=12) # we already handled the x-label with ax1 # we already handled the x-label with ax1 # + plot_suffix) # we already handled the x-label with ax1 # ax.set_ylim(ax1.get_ylim()) # fig.tight_layout() plot observables during the test # plt.tw # we already handled the x-label with ax1 # plt.tw # we already handled the x-label with ax1 logging functon # if 'directory_name' in kwargs: #     project_directory = kwargs.get('directory_name') # variable to store the total number of steps # Create a simulated environment # ------------------------------------------------------------------------------------------------------ # -------------------------------------Try to set correct anchors--------------------------------------- # Get the initial parameters of each model # These are used in later epochs when we aim to re-train the models anew with the new dataset # num list of parameters # update the anchor model losses: # This stores the anchor weights to the models for the regularisation: # ------------------------------------------------------------------------------------------------------ # -------------------------------------Try to set correct anchors--------------------------------------- # save_data = save_data(clock_time) # lists to store rewards and length of the trajectories completed # Execute in serial the environment, storing temporarily the trajectories. # Todo: Test randomization stronger if reward lower...we need a good scheme # target_threshold ????? # iterate over a fixed number of steps # steps_train = steps_per_env # found = False # while not(found): # run the policy # Sample random action during the first epoch # act = np.clip(act + np.random.randn(act.shape[0], act.shape[1]) * 0.1, -1, 1) # print('act', act*12) # take a step in the environment # add the new transition to the temporary buffer # if ep == 0: #     # try: #     # Initialize randomly a training and validation set #     model_buffer.generate_random_dataset() #     # get both datasets #     train_obs, train_act, train_rew, train_nxt_obs, _ = model_buffer.get_training_batch() #     valid_obs, valid_act, valid_rew, valid_nxt_obs, _ = model_buffer.get_valid_batch() #     target_threshold = max(max(valid_rew), max(train_rew)) #     # print('-- '*38, target_threshold) #     found = target_threshold>=-0.1 and step_count>=191 #     # except: #     #     pass # save the data for plotting the collected data for the model # env_test.env.set_usage('default') # plot_results(env_test, f'Total {total_iterations}, ' #                        f'data points: {len(model_buffer.train_idx) + len(model_buffer.valid_idx)}, ' #                        f'modelit: {ep}') ############################################################ ###################### MODEL LEARNING ###################### ############################################################ # Initialize randomly a training and validation set # get both datasets # min(target_threshold, -0.05) # Learning rate as function of ep # if ep < 1 else 1e-3 # simulated_steps = simulated_steps if ep<10 else 10000 # train the dynamic model on the datasets just sampled ############################################################ ###################### POLICY LEARNING ###################### ############################################################ # plot_results(env_test, f'Total {total_iterations}, ' #                        f'data points: {len(model_buffer.train_idx) + len(model_buffer.valid_idx)}, ' #                        f'modelit: {ep}') # # For Niky:  perform default test! ----------------------------- # if ep == 0: #     it = 0 #     # perform test! ----------------------------- #     env_test.env.set_usage('test') #     mn_test, mn_test_std, mn_length, mn_success = test_agent(env_test, action_op, num_games=50) #     # perform test! ----------------------------- #     label = f'Init {total_iterations}, ' + \ #             f'data points: {len(model_buffer.train_idx) + len(model_buffer.valid_idx)}, ' + \ #             f'ep: {ep}, it: {it}\n' + hyp_str_all #     # plot results of test ----------------------------- #     plot_results(env_test, label=label) ##################### MODEL SIMLUATION ##################### # obs_batch, act_batch, adv_batch, rtg_batch = simulate_environment(sim_env, action_op_noise, simulated_steps) # verification_simulate_environment(sim_env, env_test, action_op_noise, 50) ################# TRPO UPDATE ################ # Testing the policy on a real environment # mn_test, mn_test_std, mn_length = test_agent(env_test, action_op, num_games=1) # plot_results(env_test, 'ME-TRPO') # print(' Test score: ', np.round(mn_test, 2), np.round(mn_test_std, 2), np.round(mn_length, 2)) # mn_test, mn_test_std, mn_length = test_agent(env_test, action_op, num_games=1) # summary = tf.Summary() # summary.value.add(tag='test/performance', simple_value=mn_test) # file_writer.add_summary(summary, step_count) # file_writer.flush() # Test the policy on simulated environment. # dynamic_wait_time_count = dynamic_wait_time(ep) # # For Niky:  perform test! ----------------------------- # env_test.env.set_usage('test') # mn_test, mn_test_std, mn_length, mn_success = test_agent(env_test, action_op, num_games=50) # # perform test! ----------------------------- # # # plot results of test ----------------------------- # plot_results(env_test, label=label) # # # saving test buffer ----------------------------- # # env_test.save_current_buffer(info=label) # print(' Test score: ', np.round(mn_test, 2), np.round(mn_test_std, 2), #       np.round(mn_length, 2), np.round(mn_success, 2)) # # # save the data for plotting the tests # tests_all.append(mn_test) # tests_std_all.append(mn_test_std) # # # perform test end! ----------------------------- # # # Start of test on the simulation: # Stop of test on the simulation # collecting/summarizing the data for plotting the progress ------------------- # data of test with real env # save the data for plotting the progress ------------------- # plotting the progress ------------------- # stop training if the policy hasn't improved # or (len(sim_rewards[sim_rewards >= 990]) > int(num_ensemble_models * 0.7)): # Test the entropy measure as convergence criterion # if np.diff(entropy_all)[-1] < 0: #     print('break') #     break # Final verification: # env_final = FelLocalEnv(tango=tango) # env_final.test = True # env.TOTAL_COUNTER = len(model_buffer.train_idx) + len(model_buffer.valid_idx) # mn_test, mn_test_std, mn_length = test_agent(env_final, action_op, num_games=100) # plot_results(env_final, 'ME-TRPO', save_name='Fermi') # For Niky:  perform final test! ----------------------------- # For Niky:  end final test! ----------------------------- # env_test.env.set_usage('default') # closing environments.. # plot the results # important notes: # Storage # Hyperparameters # Scaling # Changes: # No init steps and less step per env 31 instead of 51 and the number of iterations is dynamic