Nengo PES learning rule and error signal dimensionality

[kshivvy]

I am utilizing a Nengo network to generate an updated activation array to control a robotic arm.

Given the current robot arm position as a (n,d) array, I want to learn the (10, ) activation array to reach some certain target position.

This is a general idea of the connections I have (without the PES learning):

input_node -> input_layer-> output_layer -> output_node

input_node: Represents the current arm position. The shape = (n, d) numpy array representing the current arm position. d = dimension (x, y, z) and n = positions along a dimension

input_layer: The spiking neurons to represent the input

output_layer: The spiking neurons to represent the output. This is of shape (10, ), representing the activation required. Each element has to be between [-1.0, and 1.0].

output_node: The decoded output from the spiking output_layer

The trouble comes when I connect the input_layer to the output_layer with the PES learning rule. My error signal is simply the distance of the end of the robot arm to the target position, but this does not match the dimensionality of the ensembles in the connection.

Below is the code:

global most_recent_activation
global learned_weights

most_recent_position = np.array(shearable_rod.position_collection)

n = most_recent_position.shape[0]
d = most_recent_position.shape[1]

# Define the nengo network
with nengo.Network() as net:

     # Returns the current arm position
     def get_rod_position(t):
         current_position = np.array(shearable_rod.position_collection)
         return current_position.ravel()

     # Calculates the distance between the end of the arm and the target
     def reward(t):
        end_position = most_recent_position[:,-1]
        distance = np.linalg.norm(end_position - target_position)
        #print("distance: ", distance)
        return -distance

     # Spiking neurons representing input rod position: get_rod_position() -> rod_posistion_node -> rod_position_ensemble
     input_node = nengo.Node(get_rod_position)
     input_layer = nengo.Ensemble(n*d, dimensions = n*d)

     # Spiking Neurons representing hidden layer
     #hidden_layer = nengo.Ensemble(n_neurons = int((2/3)*(n*d) + 13*2), dimensions = n*d)

     # Spiking neurons representing output layer
     output_layer = nengo.Ensemble(n_neurons = 128, dimensions = 10,
                                                     intercepts = Uniform(low=-1.0, high=1.0))
     output_node = nengo.Node(size_in = n*d, size_out = 10) # decodes the output of the spiking neurons

     # Weights to transform from input to output_layer
     weights = np.random.normal(size=(10, n*d))

     if (learned_weights is not None):
         weights = learned_weights

     # Conncetions: input_node -> input_layer -> hidden_layer -> output_layer -> output_node
     nengo.Connection(input_node, input_layer)
     learning_connection = nengo.Connection(input_layer, output_layer,
                                        transform = weights,
                                        learning_rule_type = nengo.PES(learning_rate=3e-4))

     #WeightSaver(learning_connection, "learning_connection_weights")
     # W2 = np.random.normal(size=(13*2, n*d))
     error_node = nengo.Node(output = reward, size_out = 1)
     nengo.Connection(error_node, learning_connection.learning_rule)

     nengo.Connection(output_layer, output_node)

     net.probe_output = nengo.Probe(output_node)

The error I get is:

My questions are as follows:

1. How can I connect the input_layer to the output_layer and learn with the scalar error signal? This is the only notion of error that I have, I can not use a pre - post error signal since I don’t know the true (10, ) activation array.

2. Is input_node -> input_layer-> output_layer -> output_node connection approach correct? The only reason I convert input_node to input_layer is because I am unable to connect the node to the ensemble with the PES learning rule. It seems to require two ensembles.

3. Do I need to have output_layer -> output_node to decode the output of the spiking neural network?

4. Later I will need to save the weights learned as I will be running this network in an online fashion. How can I do this? I could not find any good example or tutorials online.

Thank you!

[travis.dewolf]

Hello!

Sounds like an interesting project. You may find these blog posts I wrote a while back on setting up an arm simulation interface and building a adaptive spiking neural controller to be relevant.

Question 1
The error that your seeing is indeed caused because the error signal that your using on your learning connection is 1D and the network is expecting a 10D error signal. You can overcome this by setting a transform on this connection to broadcast it to a 10D vector. E.g. nengo.Connection(error_node, learning_connection.learning_rule, transform=np.ones((10,1))), which would apply the same error signal to all 10 dimensions. Setting this transform would be the same as changing your reward function to return -1 * np.ones(10) * distance.

But underlying this is a different issue in what exactly the expected behaviour of this learning will be. This depends on how exactly the learned 10D signal is used by the robot. For example, if the robot is 20 units away from the robot, then the learning signal for all 10 dimensions will be 20. This will cause the output for each dimension to start decreasing (because Nengo uses -1 * training signal). If this causes the robot to move further away, then the distance from the target increases and the system quickly becomes unstable.

The PES learning rule is a supervised learning rule, so you want to know what the target output of the network is or have worked through some Lyapunov stability analysis with the specific learning rule to prove stability.

Questions 2/3
For this kind of network you can have your set up with just one node. Inside this node you can both take in the learned signal from the neural network, and output the current state of the robot arm. So I would set it up as the following: interface_node -> ensemble -> interface_node

interface_node: Takes as input the 10D output from ensemble (which it can clip to [-1, 1], and gives as output the n*d feedback signal from the robotic arm

ensemble: Takes as input the n*d feedback signal from the robotic arm, and has a learning connection back to the interface node. On this connection it learns a 10D output signal that drives the arm.

The learning rules in Nengo require that the pre-synaptic object is a set of neurons, so that’s the reason that you’re seeing that error when trying to apply the PES rule to the connection from node -> ensemble. You can, however, apply the PES rule to the connection from ensemble -> node, because the post-synaptic object does not need to be neurons.

Question 4
Here’s a working example of how I tend to do this:

import nengo
import numpy as np
import matplotlib.pyplot as plt

def run_net(weights):
    with nengo.Network(seed=0) as n:
        ens = nengo.Ensemble(100, 1)
        out = nengo.Node(size_in=1)
        error = nengo.Node([-.1])
        # create the learning connection
        conn = nengo.Connection(ens.neurons, out, learning_rule_type=nengo.PES(1e-2), transform=weights)
        nengo.Connection(error, conn.learning_rule)
        probe = nengo.Probe(out)
        
    with nengo.Simulator(n) as sim:
        # run the sim and plot the output
        sim.run(1)
        plt.figure()
        plt.plot(sim.trange(), sim.data[probe])
        # return connection weights from the last time step
        weights = sim.signals[sim.model.sig[conn]["weights"]]
    return weights

weights = np.zeros((1, 100))
for ii in range(2):
    weights = run_net(weights)
    print(weights)
plt.show()

The important points:

a) Set the seed on the network so that all your parameters (encoders, gains, biases, etc) are consistent between runs. If you train your weights given some specific set of neurons, then try to use them again starting with a completely different set of neurons there will be issues.

b) Change the pre-synaptic object to ensemble.neurons, and specify an initial set of connection weights using the transform parameter (the example starts off with all zeros).

c) Build and run your simulation, then access your trained weights inside the context:

with nengo.Simulator(net) as sim:
    sim.run(runtime)
    ...
    weights = sim.signals[sim.model.sig[conn]["weights"]]

if you try to access it outside the simulator context

with nengo.Simulator(net) as sim:
    sim.run(runtime)
    ...
weights = sim.signals[sim.model.sig[conn]["weights"]]

you will get an error because when you exit the context your simulator closes and data from the runs is no longer accessible

d) Use these weights to initialize instead of the weights = np.zeros when running the network again. I usually save them to file with np.savez_compressed and then look for that file when building my network (and load in with np.load), and if there is no file then I use np.zeros instead.

Hopefully that answers most of your questions, if not please let me know! And what the learning rule should be will depend on what the 10D input to the robot controls, exactly (I’m happy to talk more about this).

Cheers,

[kshivvy]

Hi Travis,

Thank you so much for your extremely detailed answers! Your explanation about error dimensions and weight saving really helped me move forward in my project. I did have some follow up questions:

Question 1

Looking at the PES weight update equation, it makes sense that the error vector has to be 10D. I’ve played around with different reward functions, but the system becomes unstable as you have noted.

I think the next step would be to generate a bunch of training data and use nengo_dl for a more traditional neural network architecture.

1. Is the biological plausibility reduced if I simply use a deep neural network with spiking neurons?

2. Is there an alternative learning rule in Nengo that allows me to maximize the reward function (in a biologically plausible way), and transforms the 1D error signal into the 10D activation somehow?

Question 2
I tried implementing you suggestion: interface_node -> ensemble -> interface_node but I’ve ran into some issues. Below is my code and associated error message.

global most_recent_activation
global learned_weights
global elastica_iteration

most_recent_position = np.array(shearable_rod.position_collection)

n = most_recent_position.shape[0]
d = most_recent_position.shape[1]

nengo_learning_rate = 3e-2

#previous_weights = np.copy(learned_weights)

# Defines the nengo network
with nengo.Network(seed = 0) as net:

     # Returns the current arm position
     def get_rod_position(t):
         current_position = np.array(shearable_rod.position_collection)
         return current_position.ravel()

     # Calculates the distance between the end of the arm and the target
     def reward(t):
        end_position = most_recent_position[:,-1]
        distance = np.linalg.norm(end_position - target_position)
        #print("distance: ", distance)
        reward = np.ones(10) * distance
        return reward

     interface_node  = nengo.Node(size_out = n*d, output = get_rod_position)

     ensemble  = nengo.Ensemble(n_neurons = 128, dimensions = 10, intercepts = Uniform(low=-1.0, high=1.0))

     # Weights to transform from input_layer to output_layer
     if (learned_weights is None):
         learned_weights = np.zeros((10, 10))

     learning_connection = nengo.Connection(ensemble, interface_node,
                                        transform = learned_weights,
                                        learning_rule_type = nengo.PES(learning_rate=nengo_learning_rate))

     nengo.Connection(interface_node, ensemble)

     # Error node takes reward and connects it to the PES learning rule
     error_node = nengo.Node(output = reward, size_out = 10)
     nengo.Connection(error_node, learning_connection.learning_rule)

     net.probe_output = nengo.Probe(ensemble)

# Run the nengo simulation
nengo_sim_dt = 0.001 # normally 0.001
nengo_iterations = 10
nengo_runtime = nengo_iterations * nengo_sim_dt
activation = None

# Only run Nengo network every 1000 elastica iterations
with nengo.Simulator(net, progress_bar=False, dt=nengo_sim_dt) as sim:
    sim.run(nengo_runtime)
    learned_weights = sim.signals[sim.model.sig[learning_connection]["weights"]]

1. I’m a bit confused about the circular connections between interface_node and ensemble. Why does this circular method work?

2. Does my original input_node -> input_layer-> output_layer technique pose issues since I’m converting between the input node and ensemble? I was able to eliminate the output_ensemble.

Thank you!
Keshav

[travis.dewolf]

Hi Keshav,

Glad you found it helpful!

Question 1
I need to know what the 10D vector input to the robot represents before I can comment on this approach working or not.

1 - The method of training is less biologically plausible, but not necessarily the end result. There are a bunch of researchers working on this kind of approach to modeling currently, for example: https://twitter.com/JonAMichaels/status/1165464126803644416
2 - For sure, one example would be this work from our lab: http://compneuro.uwaterloo.ca/publications/Rasmussen2017.html

Question 2
I reworked the code for you (won’t be able to do this most times but had time today), please find it attached. I haven’t run it so there are possibly a couple of bugs, but it should address the problem that you were facing (and a couple others that were going to come up). The error you’re seeing is because you’re trying to connect into a Node and by default they don’t accept any input. To set up a Node to accept input you have to specify size_in in its constructor. Also, the node function (in this case your get_rod_position function) has to be set up to accept both time t and an input signal x. I’ve done that and there’s a comment for where you should send that signal out to the robot.

This might answer your sub-questions here. The architecture that you were using with input_node -> input_layer -> output_layer would also work, but you don’t strictly need the output_layer. In the interest of keeping the model as simple as possible I removed it. Do you mean that you were able to eliminate output_node?

scrap2.py (2.2 KB)