# Path: pyrl/networks.py

# 
# Path: pyrl/sgd.py
# class Adam(object):
#     def __init__(self, trainables, accumulators=None):
#         self.trainables = trainables
# 
#         if accumulators is None:
#             self.means = [theanotools.shared(0*x.get_value()) for x in trainables]
#             self.vars  = [theanotools.shared(0*x.get_value()) for x in trainables]
#             self.time  = theanotools.shared(0)
#         else:
#             self.means = [theanotools.shared(x) for x in accumulators[0]]
#             self.vars  = [theanotools.shared(x) for x in accumulators[1]]
#             self.time  = theanotools.shared(accumulators[2])
# 
#     def get_values(self):
#         means  = [x.get_value() for x in self.means]
#         vars_  = [x.get_value() for x in self.vars]
#         time   = self.time.get_value()
# 
#         return [means, vars_, time]
# 
#     def get_updates(self, loss, lr, max_norm=1, beta1=0.9, beta2=0.999,
#                     epsilon=1e-8, grads=None):
#         # Gradients
#         if grads is None:
#             grads = tensor.grad(loss, self.trainables)
# 
#         # Clipping
#         norm  = tensor.sqrt(sum([tensor.sqr(g).sum() for g in grads]))
#         m     = theanotools.clipping_multiplier(norm, max_norm)
#         grads = [m*g for g in grads]
# 
#         # Safeguard against numerical instability
#         new_cond = tensor.or_(tensor.or_(tensor.isnan(norm), tensor.isinf(norm)),
#                               tensor.or_(norm < 0, norm > 1e10))
#         grads = [tensor.switch(new_cond, np.float32(0), g) for g in grads]
# 
#         # Safeguard against numerical instability
#         #cond  = tensor.or_(norm < 0, tensor.or_(tensor.isnan(norm), tensor.isinf(norm)))
#         #grads = [tensor.switch(cond, np.float32(0), g) for g in grads]
# 
#         # New values
#         t       = self.time + 1
#         lr_t    = lr*tensor.sqrt(1. - beta2**t)/(1. - beta1**t)
#         means_t = [beta1*m + (1. - beta1)*g for g, m in zip(grads, self.means)]
#         vars_t  = [beta2*v + (1. - beta2)*tensor.sqr(g) for g, v in zip(grads, self.vars)]
#         steps   = [lr_t*m_t/(tensor.sqrt(v_t) + epsilon)
#                    for m_t, v_t in zip(means_t, vars_t)]
# 
#         # Updates
#         updates  = [(x, x - step) for x, step in zip(self.trainables, steps)]
#         updates += [(m, m_t) for m, m_t in zip(self.means, means_t)]
#         updates += [(v, v_t) for v, v_t in zip(self.vars, vars_t)]
#         updates += [(self.time, t)]
# 
#         return norm, grads, updates


from   collections import OrderedDict
from   theano import tensor
from .         import nptools, tasktools, theanotools, utils
from .debug    import DEBUG
from .networks import Networks
from .sgd      import Adam
import datetime
import sys
import numpy as np
import theano

    def func_update_baseline(self, use_x0=False, accumulators=None):
        U  = tensor.tensor3('U')
        R  = tensor.matrix('R')
        R_ = R.reshape((R.shape[0], R.shape[1], 1))
        Q  = tensor.tensor3('Q')

        if use_x0:
            x0_ = tensor.matrix('x0_')
        else:
            x0  = self.baseline_net.params['x0']
            x0_ = tensor.alloc(x0, U.shape[1], x0.shape[0])

        z_0   = self.baseline_net.get_outputs_0(x0_)
        r, z  = self.baseline_net.get_outputs(U, Q, x0_)
        z_all = tensor.concatenate([z_0.reshape((1, z_0.shape[0], z_0.shape[1])), z],
                                   axis=0)

        # Learning rate
        lr = tensor.scalar('lr')

        # Reward prediction error
        M    = tensor.matrix('M')
        L2   = tensor.sum((tensor.sqr(z_all[:,:,0] - R))*M)/tensor.sum(M)
        RMSE = tensor.sqrt(L2)

        # Objective function
        obj = L2 + self.baseline_net.get_regs(x0_, r, M)

        # SGD

# Path: pyrl/networks.py

# 
# Path: pyrl/sgd.py
# class Adam(object):
#     def __init__(self, trainables, accumulators=None):
#         self.trainables = trainables
# 
#         if accumulators is None:
#             self.means = [theanotools.shared(0*x.get_value()) for x in trainables]
#             self.vars  = [theanotools.shared(0*x.get_value()) for x in trainables]
#             self.time  = theanotools.shared(0)
#         else:
#             self.means = [theanotools.shared(x) for x in accumulators[0]]
#             self.vars  = [theanotools.shared(x) for x in accumulators[1]]
#             self.time  = theanotools.shared(accumulators[2])
# 
#     def get_values(self):
#         means  = [x.get_value() for x in self.means]
#         vars_  = [x.get_value() for x in self.vars]
#         time   = self.time.get_value()
# 
#         return [means, vars_, time]
# 
#     def get_updates(self, loss, lr, max_norm=1, beta1=0.9, beta2=0.999,
#                     epsilon=1e-8, grads=None):
#         # Gradients
#         if grads is None:
#             grads = tensor.grad(loss, self.trainables)
# 
#         # Clipping
#         norm  = tensor.sqrt(sum([tensor.sqr(g).sum() for g in grads]))
#         m     = theanotools.clipping_multiplier(norm, max_norm)
#         grads = [m*g for g in grads]
# 
#         # Safeguard against numerical instability
#         new_cond = tensor.or_(tensor.or_(tensor.isnan(norm), tensor.isinf(norm)),
#                               tensor.or_(norm < 0, norm > 1e10))
#         grads = [tensor.switch(new_cond, np.float32(0), g) for g in grads]
# 
#         # Safeguard against numerical instability
#         #cond  = tensor.or_(norm < 0, tensor.or_(tensor.isnan(norm), tensor.isinf(norm)))
#         #grads = [tensor.switch(cond, np.float32(0), g) for g in grads]
# 
#         # New values
#         t       = self.time + 1
#         lr_t    = lr*tensor.sqrt(1. - beta2**t)/(1. - beta1**t)
#         means_t = [beta1*m + (1. - beta1)*g for g, m in zip(grads, self.means)]
#         vars_t  = [beta2*v + (1. - beta2)*tensor.sqr(g) for g, v in zip(grads, self.vars)]
#         steps   = [lr_t*m_t/(tensor.sqrt(v_t) + epsilon)
#                    for m_t, v_t in zip(means_t, vars_t)]
# 
#         # Updates
#         updates  = [(x, x - step) for x, step in zip(self.trainables, steps)]
#         updates += [(m, m_t) for m, m_t in zip(self.means, means_t)]
#         updates += [(v, v_t) for v, v_t in zip(self.vars, vars_t)]
#         updates += [(self.time, t)]
# 
#         return norm, grads, updates


# Path: pyrl/policygradient.py
from   collections import OrderedDict
from   theano import tensor
from .         import nptools, tasktools, theanotools, utils
from .debug    import DEBUG
from .networks import Networks
from .sgd      import Adam
import datetime
import sys
import numpy as np
import theano


    def func_update_baseline(self, use_x0=False, accumulators=None):
        U  = tensor.tensor3('U')
        R  = tensor.matrix('R')
        R_ = R.reshape((R.shape[0], R.shape[1], 1))
        Q  = tensor.tensor3('Q')

        if use_x0:
            x0_ = tensor.matrix('x0_')
        else:
            x0  = self.baseline_net.params['x0']
            x0_ = tensor.alloc(x0, U.shape[1], x0.shape[0])

        z_0   = self.baseline_net.get_outputs_0(x0_)
        r, z  = self.baseline_net.get_outputs(U, Q, x0_)
        z_all = tensor.concatenate([z_0.reshape((1, z_0.shape[0], z_0.shape[1])), z],
                                   axis=0)

        # Learning rate
        lr = tensor.scalar('lr')

        # Reward prediction error
        M    = tensor.matrix('M')
        L2   = tensor.sum((tensor.sqr(z_all[:,:,0] - R))*M)/tensor.sum(M)
        RMSE = tensor.sqrt(L2)

        # Objective function
        obj = L2 + self.baseline_net.get_regs(x0_, r, M)

        # SGD

        self.baseline_sgd = Adam(self.baseline_net.trainables, accumulators=accumulators)