def createComputationalGraph(self):
        """
            Create the computational graph in tensorflow
        """

        ######################################
        # Variables that are needed
        self.u = tf.Variable(np.zeros(self.N), dtype=self.dtype)
        self.rLowPass = tf.Variable(np.zeros(self.N), dtype=self.dtype)
        uNoise = tf.Variable(np.zeros(self.N), dtype=self.dtype)
        self.uNoiseLowPass = tf.Variable(np.zeros(self.N), dtype=self.dtype)
        self.uDotOld = tf.Variable(np.zeros(self.N), dtype=self.dtype)
        self.eligNow = tf.Variable(np.zeros((self.N, self.N)),
                                   dtype=self.dtype)
        self.eligibility = tf.Variable(np.zeros((self.N, self.N)),
                                       dtype=self.dtype)
        self.regEligibility = tf.Variable(np.zeros((self.N, self.N)),
                                          dtype=self.dtype)
        self.wTfNoWta = tf.Variable(self.WnoWta, dtype=self.dtype)
        self.wTfOnlyWta = tf.Variable(self.onlyWta, dtype=self.dtype)
        inputMask = self.input.getInput(0.)[2]
        self.inputMaskTf = tf.Variable(inputMask, dtype=self.dtype)
        outputMask = self.target.getTarget(self.T)[2]
        self.outputMaskTf = tf.Variable(outputMask, dtype=self.dtype)
        self.biasTf = tf.Variable(self.biasVector, dtype=self.dtype)
        # set up a mask for the learned weights in self.wTfNoWta
        # note that W.mask must omit the WTA network
        self.wNoWtaMask = tf.Variable(self.Wplastic.astype(float),
                                      dtype=self.dtype)

        #####################################
        # Variables for debugging
        self.error = tf.Variable(np.zeros(self.N), dtype=self.dtype)

        #####################################
        # Placeholders
        # The only datatransfer between the GPU and the CPU should be the
        # input to the input layer and the modulatory signal
        self.inputTf = tf.placeholder(dtype=self.dtype, shape=(self.N))
        self.inputPrimeTf = tf.placeholder(dtype=self.dtype, shape=(self.N))
        self.modulator = tf.placeholder(dtype=self.dtype, shape=())

        ####################################
        # Aux variables for the calculations
        nInput = len(np.where(inputMask == 1)[0])
        nOutput = len(np.where(outputMask == 1)[0])
        nFull = len(inputMask)

        #####################################################
        # Start the actual calculations for the comp graph  #
        #####################################################

        ####################################
        # Calculate the activations functions using the updated values
        self.rho = self.actFunc(self.u)
        rhoPrime = self.actFuncPrime(self.u)
        rhoPrimePrime = self.actFuncPrimePrime(self.u)
        self.rhoOutput = self.actFunc(self.u)

        ###################################
        # Update the exploration noise on the output neurons
        uNoiseOut = tf.slice(uNoise, [nFull - nOutput], [-1])
        duOutNoise = self.noiseTheta * (self.noiseMean - uNoiseOut) * self.timeStep + self.noiseSigma * \
            np.sqrt(self.timeStep) * \
            tf.random_normal([nOutput], mean=0., stddev=1.0, dtype=self.dtype)
        updateNoise = tf.scatter_update(uNoise,
                                        np.arange(nFull - nOutput, nFull),
                                        uNoiseOut + duOutNoise)
        # Update the low-pass noise
        with tf.control_dependencies([updateNoise]):
            self.uNoiseLowPass = self.uNoiseLowPass + (
                self.timeStep / self.tau) * (uNoise - self.uNoiseLowPass)

        ####################################
        # Calculate the updates for the membrane potential and for the
        # eligibility trace

        with tf.control_dependencies(
            [self.uNoiseLowPass, updateNoise, rhoPrime, rhoPrimePrime]):

            # frequently used tensors are claculated early on
            wNoWtaT = tf.transpose(self.wTfNoWta)
            wNoWtaRho = tfTools.tf_mat_vec_dot(self.wTfNoWta, self.rho)
            c = tfTools.tf_mat_vec_dot(
                wNoWtaT, self.u - wNoWtaRho - self.biasTf - self.inputTf)

            # get the matrix side of the equation
            A1 = tf.matmul(self.wTfNoWta, tf.diag(rhoPrime))
            A2 = tf.matmul(tf.diag(c), tf.diag(rhoPrimePrime))
            A3 = tf.matmul(tf.diag(rhoPrime), wNoWtaT)
            A4 = tf.matmul(tf.matmul(tf.diag(rhoPrime), wNoWtaT),
                           tf.matmul(self.wTfNoWta, tf.diag(rhoPrime)))
            #A5 = self.beta * self.alphaWna * tf.matmul(self.wTfOnlyWta, tf.diag(rhoPrime))
            A = self.tau * (tf.eye(self.N) - A1 - A2 - A3 + A4)

            # get the vector side of the equation
            y1 = wNoWtaRho + self.biasTf + self.inputTf + self.beta * uNoise - self.u
            y2 = self.tau * self.inputPrimeTf
            y3 = rhoPrime * c
            y4 = self.tau * rhoPrime * tfTools.tf_mat_vec_dot(
                wNoWtaT, self.inputPrimeTf)
            #y5 = self.beta * self.alphaWna * tfTools.tf_mat_vec_dot(
            #                            self.wTfOnlyWta, self.rho)
            #y6 = rhoPrime * tfTools.tf_mat_vec_dot(wNoWtaT, uNoise - self.uNoiseLowPass)
            y = y1 + y2 + y3 - y4

            # Solve the equation for uDot
            #self.uDiff = (1. / self.tau) * tf.linalg.solve(A, y)
            uDiff = tf.linalg.solve(A, tf.expand_dims(y, 1))[:, 0]
            #chol = tf.cholesky(A)
            #uDiff = tf.cholesky_solve(tf.cholesky(A), tf.expand_dims(y, 1))[:, 0]
            """
            # The regular component with lookahead
            reg = tfTools.tf_mat_vec_dot(
                self.wTfNoWta, self.rho + rhoPrime * self.uDotOld * self.tau) - self.u + self.biasTf

            # Error term from the vanilla lagrange
            eVfirst = rhoPrime * c
            eVsecond = (rhoPrimePrime * self.uDotOld) * c
            eVthird = rhoPrime * \
                tfTools.tf_mat_vec_dot(
                    wNoWtaT,
                    self.uDotOld - tfTools.tf_mat_vec_dot(
                        self.wTfNoWta,
                        rhoPrime * self.uDotOld)
                )
            eV = eVfirst + self.tau * (eVsecond + eVthird)

            # terms from the winner nudges all circuit
            regWna = self.beta * self.alphaWna * tfTools.tf_mat_vec_dot(
                self.wTfOnlyWta, self.rho + self.tau * rhoPrime * self.uDotOld)

            # Terms from the exploration noise term
            #eNoise = self.alphaNoise * self.beta * \
            #    ((uNoise) - (uOut + self.tau * uDotOut))
            eNoise = self.alphaNoise * self.beta * uNoise
            """

        #uDiff = (1. / self.tau) * (reg + eV + regWna + eNoise)
        saveOldUDot = self.uDotOld.assign(uDiff)
        updateLowPassActivity = self.rLowPass.assign(
            (self.rLowPass + self.timeStep / self.tauEligibility * self.rho) *
            tf.exp(-1. * self.timeStep / self.tauEligibility))

        self.eligNowUpdate = self.eligNow.assign(
            tfTools.tf_outer_product(
                self.u - tfTools.tf_mat_vec_dot(self.wTfNoWta, self.rho) -
                self.biasTf, self.rho))
        errorUpdate = self.error.assign(
            self.u - tfTools.tf_mat_vec_dot(self.wTfNoWta, self.rho) -
            self.biasTf - self.inputTf)

        with tf.control_dependencies([
                saveOldUDot, updateLowPassActivity, self.eligNowUpdate,
                errorUpdate
        ]):

            self.updateEligiblity = self.eligibility.assign(
                (self.eligibility + self.timeStep * tfTools.tf_outer_product(
                    self.u - tfTools.tf_mat_vec_dot(self.wTfNoWta, self.rho) -
                    self.biasTf - self.inputTf - self.uNoiseLowPass, self.rho))
                * tf.exp(-1. * self.timeStep / self.tauEligibility))

            self.updateRegEligibility = self.regEligibility.assign(
                (self.regEligibility + self.timeStep *
                 tfTools.tf_outer_product(-1.0 *
                                          (self.u - self.uHigh), self.rho)) *
                tf.exp(-1. * self.timeStep / self.tauEligibility))

            #self.applyMembranePot = tf.scatter_update(self.u, np.arange(
            #    nInput, nFull), tf.slice(self.u, [nInput], [-1]) + self.timeStep * tf.slice(uDiff, [nInput], [-1]))

        with tf.control_dependencies([
                saveOldUDot, updateLowPassActivity, self.eligNowUpdate,
                errorUpdate, self.updateEligiblity, self.updateRegEligibility
        ]):
            self.applyMembranePot = self.u.assign(self.u +
                                                  self.timeStep * uDiff)

        ###############################################
        ## Node to update the weights of the network ##
        ###############################################

        self.updateW = self.wTfNoWta.assign(
            self.wTfNoWta + (self.learningRate / self.tauEligibility) *
            (self.modulator * self.eligibility * self.Wplastic +
             tf.abs(self.modulator) * self.kappaDecay * self.regEligibility *
             self.noWnaMask))
Beispiel #2
0
    def createComputationalGraph(self):
        """
            Create the computational graph in tensorflow
        """
        # track the dependencies
        dependencies = []

        # Set up the variables which will be then tracked
        self.u = tf.Variable(np.zeros(self.N), dtype=self.dtype)
        self.uDotOld = tf.Variable(np.zeros(self.N), dtype=self.dtype)
        self.eligibility = tf.Variable(np.zeros((self.N, self.N)),
                                       dtype=self.dtype)
        self.eligibilityDiff = tf.Variable(np.zeros((self.N, self.N)),
                                           dtype=self.dtype)
        self.rho = tf.Variable(np.zeros(self.N), dtype=self.dtype)
        self.rhoPrime = tf.Variable(np.zeros(self.N), dtype=self.dtype)
        self.rhoPrimePrime = tf.Variable(np.zeros(self.N), dtype=self.dtype)
        self.regTerm = tf.Variable(np.zeros((self.N, self.N)),
                                   dtype=self.dtype)
        self.regTermDiff = tf.Variable(np.zeros((self.N, self.N)),
                                       dtype=self.dtype)

        # We need an identity tensor with float values
        identity = tf.Variable(np.eye(self.N), dtype=self.dtype)
        one = tf.Variable(np.eye(1), dtype=self.dtype)

        # Set up the placeholdeers which can be then modified
        self.tfW = tf.placeholder(shape=self.W.shape, dtype=self.dtype)
        self.tfWnoWta = tf.placeholder(shape=self.W.shape, dtype=self.dtype)
        self.inputTf = tf.placeholder(dtype=self.dtype, shape=(self.N))
        self.inputPrimeTf = tf.placeholder(dtype=self.dtype, shape=(self.N))
        self.inputMaskTf = tf.placeholder(dtype=self.dtype, shape=(self.N))
        self.targetTf = tf.placeholder(dtype=self.dtype, shape=(self.N))
        self.targetPrimeTf = tf.placeholder(dtype=self.dtype, shape=(self.N))
        self.targetMaskTf = tf.placeholder(dtype=self.dtype, shape=(self.N))

        # an example input mask is needed to build the comp graph
        inputMask = self.input.getInput(0.)[2]
        nInput = len(np.where(inputMask == 1)[0])
        nFull = len(inputMask)

        # Start the calculations
        # Calculate the activation functions
        dependencies.append(tf.assign(self.rho, self.actFunc(self.u)))
        dependencies.append(tf.assign(self.rhoPrime,
                                      self.actFuncPrime(self.u)))
        dependencies.append(
            tf.assign(self.rhoPrimePrime, self.actFuncPrimePrime(self.u)))

        # set the membrane potential on the input neurons
        dependencies.append(
            tf.scatter_update(self.u, np.arange(nInput),
                              tf.slice(self.inputTf, [0], [nInput])))
        dependencies.append(
            tf.scatter_update(self.uDotOld, np.arange(nInput),
                              tf.slice(self.inputPrimeTf, [0], [nInput])))

        with tf.control_dependencies(dependencies):

            # frequently used ones
            Wt = tf.transpose(self.tfW)
            Wrho = tfTools.tf_mat_vec_dot(self.tfW, self.rho)
            c = tfTools.tf_mat_vec_dot(Wt, self.u - Wrho)

            # calculate the update
            term1 = Wrho + self.tau * \
                tfTools.tf_mat_vec_dot(
                    self.tfW, self.rhoPrime * self.uDotOld) - self.u
            term2 = c * (self.rhoPrime +
                         self.tau * self.rhoPrimePrime * self.uDotOld)
            term3 = self.tau * self.rhoPrime * (tfTools.tf_mat_vec_dot(
                Wt, self.uDotOld) - tfTools.tf_mat_vec_dot(
                    tf.matmul(Wt, self.tfW), self.rhoPrime * self.uDotOld))
            term4 = self.beta * self.targetMaskTf * \
                (self.targetTf - self.u + self.tau *
                 (self.targetPrimeTf - self.uDotOld))

            uDiffDummy = term1 + term2 + term3 + term4

        self.uDiff = (1. / self.tau) * uDiffDummy
        dependencies.append(self.uDiff)
        dependencies.append(self.uDotOld.assign(self.uDiff))

        # Calculate the step in the eligibility trace
        dependencies.append(
            self.eligibilityDiff.assign(
                self.learningRate * tfTools.tf_outer_product(
                    self.u - tfTools.tf_mat_vec_dot(self.tfWnoWta, self.rho),
                    self.rho)))
        dependencies.append(
            self.regTermDiff.assign(
                self.learningRate * tfTools.tf_outer_product(
                    tf.nn.relu(self.uLow - self.u) -
                    tf.nn.relu(self.u - self.uHigh), self.rho)))

        self.logger.warn('It is the approximativ lagrange')

        with tf.control_dependencies(dependencies):

            # Apply membrane potentials
            self.applyMembranePot = tf.scatter_update(
                self.u, np.arange(nInput, nFull),
                tf.slice(self.u, [nInput], [-1]) +
                self.timeStep * tf.slice(self.uDiff, [nInput], [-1]))

            # Apply eligibility trace
            dependencies.append(
                self.eligibility.assign(self.eligibility +
                                        self.timeStep * self.eligibilityDiff))

            dependencies.append(
                self.regTerm.assign(self.regTerm +
                                    self.timeStep * self.regTermDiff))

        with tf.control_dependencies(dependencies):
            # Apply decay to the elifibility trace
            self.applyEligibility = self.eligibility.assign(
                self.eligibility *
                tf.exp(-1. * one * self.timeStep / self.tauEligibility))
            self.applyRegTerm = self.regTerm.assign(
                self.regTerm *
                tf.exp(-1. * one * self.timeStep / self.tauEligibility))
    def createComputationalGraph(self):
        """
            Create the computational graph in tensorflow
        """

        ######################################
        # Variables that are needed
        self.u = tf.Variable(np.zeros(self.N), dtype=self.dtype)
        self.rLowPass = tf.Variable(np.zeros(self.N), dtype=self.dtype)
        uNoise = tf.Variable(np.zeros(self.N), dtype=self.dtype)
        self.uDotOld = tf.Variable(np.zeros(self.N), dtype=self.dtype)
        self.eligibility = tf.Variable(
            np.zeros((self.N, self.N)), dtype=self.dtype)
        self.regEligibility = tf.Variable(
            np.zeros((self.N, self.N)), dtype=self.dtype)
        self.wTfNoWta = tf.Variable(self.WnoWta, dtype=self.dtype)
        self.wTfOnlyWta = tf.Variable(self.onlyWta, dtype=self.dtype)
        inputMask = self.input.getInput(0.)[2]
        self.inputMaskTf = tf.Variable(inputMask,
                                       dtype=self.dtype)
        outputMask = self.target.getTarget(self.T)[2]
        self.outputMaskTf = tf.Variable(outputMask,
                                        dtype=self.dtype)
        # set up a mask for the learned weights in self.wTfNoWta
        # note that W.mask must omit the WTA network
        self.wNoWtaMask = tf.Variable(self.Wplastic.astype(float), dtype=self.dtype)

        #####################################
        # Placeholders
        # The only datatransfer between the GPU and the CPU should be the
        # input to the input layer and the modulatory signal
        self.inputTf = tf.placeholder(dtype=self.dtype,
                                      shape=(self.N))
        self.inputPrimeTf = tf.placeholder(dtype=self.dtype,
                                           shape=(self.N))
        self.modulator = tf.placeholder(dtype=self.dtype,
                                        shape=())

        ####################################
        # Aux variables for the calculations
        nInput = len(np.where(inputMask == 1)[0])
        nOutput = len(np.where(outputMask == 1)[0])
        nFull = len(inputMask)

        #####################################################
        # Start the actual calculations for the comp graph  #
        #####################################################

        ####################################
        # Update the values of u and uDotOld according to the input
        applyInputU = tf.scatter_update(self.u,
                                        np.arange(nInput),
                                        tf.slice(self.inputTf,
                                                 [0],
                                                 [nInput]
                                                 )
                                        )
        applyInputUDot = tf.scatter_update(self.uDotOld,
                                           np.arange(nInput),
                                           tf.slice(self.inputPrimeTf,
                                                    [0],
                                                    [nInput]
                                                    )
                                           )

        ####################################
        # Calculate the activations functions using the updated values
        with tf.control_dependencies([applyInputU, applyInputUDot]):
            self.rho = self.actFunc(self.u)
            rhoPrime = self.actFuncPrime(self.u)
            rhoPrimePrime = self.actFuncPrimePrime(self.u)
        self.rhoOutput = self.actFunc(self.u)

        ###################################
        # Update the exploration noise on the output neurons
        uNoiseOut = tf.slice(uNoise, [nFull - nOutput], [-1])
        duOutNoise = self.noiseTheta * (self.noiseMean - uNoiseOut) * self.timeStep + self.noiseSigma * \
            np.sqrt(self.timeStep) * \
            tf.random_normal([nOutput], mean=0., stddev=1.0, dtype=self.dtype)
        updateNoise = tf.scatter_update(uNoise,
                                        np.arange(nFull - nOutput, nFull),
                                        uNoiseOut + duOutNoise)

        ####################################
        # Calculate the updates for the membrane potential and for the
        # eligibility trace

        with tf.control_dependencies([updateNoise,
                                      applyInputU,
                                      applyInputUDot]):

            # frequently used tensors are claculated early on
            wNoWtaT = tf.transpose(self.wTfNoWta)
            wNoWtaRho = tfTools.tf_mat_vec_dot(self.wTfNoWta, self.rho)
            c = tfTools.tf_mat_vec_dot(wNoWtaT, self.u - wNoWtaRho)
            uOut = self.u * self.outputMaskTf
            uDotOut = self.uDotOld * self.outputMaskTf
            wOnlyWtaT = tf.transpose(self.wTfOnlyWta)
            wOnlyWtaRho = tfTools.tf_mat_vec_dot(self.wTfOnlyWta, self.rho)
            cOnlyWta = tfTools.tf_mat_vec_dot(wOnlyWtaT, uOut - wOnlyWtaRho)

            # The regular component with lookahead
            reg = tfTools.tf_mat_vec_dot(
                self.wTfNoWta, self.rho + rhoPrime * self.uDotOld * self.tau) - self.u

            # Error term from the vanilla lagrange
            eVfirst = rhoPrime * c
            eVsecond = (rhoPrimePrime * self.uDotOld) * c
            eVthird = rhoPrime * \
                tfTools.tf_mat_vec_dot(
                    wNoWtaT,
                    self.uDotOld - tfTools.tf_mat_vec_dot(
                        self.wTfNoWta,
                        rhoPrime * self.uDotOld)
                )
            eV = eVfirst + self.tau * (eVsecond + eVthird)

            # terms from the winner nudges all circuit
            eWnaFirst = tfTools.tf_mat_vec_dot(
                self.wTfOnlyWta, self.rho + rhoPrime * uDotOut) - (uOut + uDotOut * self.tau)
            eWnaSecond = self.outputMaskTf * rhoPrime * cOnlyWta
            eWnaThird = (self.outputMaskTf *
                         rhoPrimePrime * uDotOut) * cOnlyWta
            eWnaFourth = rhoPrime * \
                tfTools.tf_mat_vec_dot(
                    wOnlyWtaT,
                    uDotOut - tfTools.tf_mat_vec_dot(
                        self.wTfOnlyWta,
                        rhoPrime * uDotOut)
                )
            eWna = self.alphaWna * self.beta * \
                (eWnaFirst + eWnaSecond + self.tau * (eWnaThird + eWnaFourth))

            # Terms from the exploration noise term
            eNoise = self.alphaNoise * self.beta * \
                ((uNoise) - (uOut + self.tau * uDotOut))

        uDiff = (1. / self.tau) * (reg + eV + eWna + eNoise)
        saveOldUDot = self.uDotOld.assign(uDiff)
        updateLowPassActivity = self.rLowPass.assign((self.rLowPass + self.timeStep / self.tauEligibility * self.rho) * tf.exp(-1. * self.timeStep / self.tauEligibility))

        with tf.control_dependencies([saveOldUDot, updateLowPassActivity]):

            self.updateEligiblity = self.eligibility.assign(
                (self.eligibility + self.timeStep * tfTools.tf_outer_product(
                    self.u - tfTools.tf_mat_vec_dot(self.wTfNoWta, self.rho), self.rho)) * tf.exp(-1. * self.timeStep / self.tauEligibility)
            )
            
            self.updateRegEligibility = self.regEligibility.assign(
                (self.regEligibility + self.timeStep * tfTools.tf_outer_product(
                    tf.nn.relu(self.uLow - self.u) -
                    tf.nn.relu(self.u - self.uHigh),
                    self.rho)) * tf.exp(-1. * self.timeStep / self.tauEligibility)
            )

            #self.applyMembranePot = tf.scatter_update(self.u, np.arange(
            #    nInput, nFull), tf.slice(self.u, [nInput], [-1]) + self.timeStep * tf.slice(uDiff, [nInput], [-1]))

            self.applyMembranePot = self.u.assign(self.u + self.timeStep * uDiff)

        ###############################################
        ## Node to update the weights of the network ##
        ###############################################

        self.updateW = self.wTfNoWta.assign(self.wTfNoWta + (self.learningRate / self.tauEligibility) * (
            self.modulator * self.eligibility + self.kappaDecay * self.regEligibility) * self.Wplastic)
    def createComputationalGraph(self):
        """
            Create the computational graph in tensorflow
        """
        # track the dependencies
        dependencies = []

        # Set up the variables which will be then tracked
        self.u = tf.Variable(np.zeros(self.N), dtype=self.dtype)
        self.uDotOld = tf.Variable(np.zeros(self.N), dtype=self.dtype)
        self.tfOnlyWta = tf.Variable(self.onlyWta, dtype=self.dtype)
        self.eligibility = tf.Variable(np.zeros((self.N, self.N)),
                                       dtype=self.dtype)
        self.eligibilityDiff = tf.Variable(np.zeros((self.N, self.N)),
                                           dtype=self.dtype)
        self.rho = tf.Variable(np.zeros(self.N), dtype=self.dtype)
        self.rhoPrime = tf.Variable(np.zeros(self.N), dtype=self.dtype)
        self.rhoPrimePrime = tf.Variable(np.zeros(self.N), dtype=self.dtype)
        self.regEligibility = tf.Variable(np.zeros((self.N, self.N)),
                                          dtype=self.dtype)
        self.regEligibilityDiff = tf.Variable(np.zeros((self.N, self.N)),
                                              dtype=self.dtype)

        # We need an identity tensor with float values
        identity = tf.Variable(np.eye(self.N), dtype=self.dtype)
        one = tf.Variable(np.eye(1), dtype=self.dtype)

        # Set up the placeholdeers which can be then modified
        self.tfW = tf.placeholder(shape=self.W.shape, dtype=self.dtype)
        self.tfWnoWta = tf.placeholder(shape=self.W.shape, dtype=self.dtype)
        self.inputTf = tf.placeholder(dtype=self.dtype, shape=(self.N))
        self.inputPrimeTf = tf.placeholder(dtype=self.dtype, shape=(self.N))
        self.inputMaskTf = tf.placeholder(dtype=self.dtype, shape=(self.N))
        self.targetTf = tf.placeholder(dtype=self.dtype, shape=(self.N))
        self.targetPrimeTf = tf.placeholder(dtype=self.dtype, shape=(self.N))
        self.targetMaskTf = tf.placeholder(dtype=self.dtype, shape=(self.N))

        # an example input mask is needed to build the comp graph
        inputMask = self.input.getInput(0.)[2]
        nInput = len(np.where(inputMask == 1)[0])
        nFull = len(inputMask)

        # Start the calculations
        # Calculate the activation functions
        dependencies.append(tf.assign(self.rho, self.actFunc(self.u)))
        dependencies.append(tf.assign(self.rhoPrime,
                                      self.actFuncPrime(self.u)))
        dependencies.append(
            tf.assign(self.rhoPrimePrime, self.actFuncPrimePrime(self.u)))

        # set the membrane potential on the input neurons
        dependencies.append(
            tf.scatter_update(self.u, np.arange(nInput),
                              tf.slice(self.inputTf, [0], [nInput])))
        dependencies.append(
            tf.scatter_update(self.uDotOld, np.arange(nInput),
                              tf.slice(self.inputPrimeTf, [0], [nInput])))

        with tf.control_dependencies(dependencies):

            # frequently used tensors are claculated early on
            wNoWtaT = tf.transpose(self.tfWnoWta)
            wNoWtaRho = tfTools.tf_mat_vec_dot(self.tfWnoWta, self.rho)
            c = tfTools.tf_mat_vec_dot(wNoWtaT, self.u - wNoWtaRho)
            uOut = self.u * self.targetMaskTf
            uDotOut = self.uDotOld * self.targetMaskTf
            wOnlyWtaT = tf.transpose(self.tfOnlyWta)
            wOnlyWtaRho = tfTools.tf_mat_vec_dot(self.tfOnlyWta, self.rho)
            cOnlyWta = tfTools.tf_mat_vec_dot(wOnlyWtaT, uOut - wOnlyWtaRho)

            # The regular component with lookahead
            reg = tfTools.tf_mat_vec_dot(
                self.tfWnoWta,
                self.rho + self.rhoPrime * self.uDotOld) - self.u

            # Error term from the vanilla lagrange
            eVfirst = self.rhoPrime * c
            eVsecond = (self.rhoPrimePrime * self.uDotOld) * c
            eVthird = self.rhoPrime * \
                tfTools.tf_mat_vec_dot(
                wNoWtaT,
                self.uDotOld - tfTools.tf_mat_vec_dot(
                    self.tfWnoWta,
                    self.rhoPrime * self.uDotOld)
            )
            eV = eVfirst + self.tau * (eVsecond + eVthird)

            # terms from the winner nudges all circuit
            eWnaFirst = tfTools.tf_mat_vec_dot(
                self.tfOnlyWta, self.rho +
                self.rhoPrime * uDotOut) - (uOut + uDotOut * self.tau)
            eWnaSecond = self.targetMaskTf * self.rhoPrime * cOnlyWta
            eWnaThird = (self.targetMaskTf * self.rhoPrimePrime *
                         uDotOut) * cOnlyWta
            eWnaFourth = self.rhoPrime * \
                tfTools.tf_mat_vec_dot(
                wOnlyWtaT,
                uDotOut - tfTools.tf_mat_vec_dot(
                    self.tfOnlyWta,
                    self.rhoPrime * uDotOut)
            )
            eWna = self.alphaWna * self.beta * \
                (eWnaFirst + eWnaSecond + self.tau * (eWnaThird + eWnaFourth))

            # Terms from the exploration noise term
            noise = self.targetTf * self.targetMaskTf
            noiseDot = self.targetPrimeTf * self.targetMaskTf
            eNoise = self.alphaNoise * self.beta * \
                ((noise + self.tau * noiseDot) - (uOut + self.tau * uDotOut))

        self.uDiff = (1. / self.tau) * (reg + eV + eWna + eNoise)
        dependencies.append(self.uDiff)
        dependencies.append(self.uDotOld.assign(self.uDiff))

        # Calculate the step in the eligibility trace
        dependencies.append(
            self.eligibilityDiff.assign(
                tfTools.tf_outer_product(
                    self.u - tfTools.tf_mat_vec_dot(self.tfWnoWta, self.rho),
                    self.rho)))
        dependencies.append(
            self.regEligibilityDiff.assign(
                tfTools.tf_outer_product(
                    tf.nn.relu(self.uLow - self.u) -
                    tf.nn.relu(self.u - self.uHigh), self.rho)))

        self.logger.warn(
            'You are using the directly derived version of the WNA and the noise term'
        )

        with tf.control_dependencies(dependencies):

            # Apply membrane potentials
            self.applyMembranePot = tf.scatter_update(
                self.u, np.arange(nInput, nFull),
                tf.slice(self.u, [nInput], [-1]) +
                self.timeStep * tf.slice(self.uDiff, [nInput], [-1]))

            # Apply eligibility trace
            dependencies.append(
                self.eligibility.assign(self.eligibility +
                                        self.timeStep * self.eligibilityDiff))
            dependencies.append(
                self.regEligibility.assign(self.regEligibility +
                                           self.timeStep *
                                           self.regEligibilityDiff))

        with tf.control_dependencies(dependencies):
            # Apply decay to the elifibility trace
            self.applyEligibility = self.eligibility.assign(
                self.eligibility *
                tf.exp(-1. * one * self.timeStep / self.tauEligibility))
            self.applyRegEligibility = self.regEligibility.assign(
                self.regEligibility *
                tf.exp(-1. * one * self.timeStep / self.tauEligibility))