Help with Basal Ganglia and Thalamus Model

I’m fairly new to using Nengo, and I want to create a simple action selector model that I can train to control the movement of an agent in a simulated environment. I seem to be having trouble with the dimensions of my input ensemble not matching the dimension of my Basal Ganglia dimension (I have set the dimension for the basal ganglia ensemble and thalamus ensemble to be the number of actions possible ie: move straight, move left, etc.). This is the code I’ve made so far and I can’t get this to run:

#Global Variables
NUM_INPUTS = 6
NUM_ACTIONS = 4
DEFAULT_ENSEMBLE_NEURONS = 100
BASALGANGLIA_NEURONS = 100 #default = 100
THALAMUS_NEURONS = 50 #default = 50



model = nengo.Network()
with model:
    
    #Define model ensembles
    input_ensemble = nengo.Ensemble(n_neurons=DEFAULT_ENSEMBLE_NEURONS * NUM_INPUTS, dimensions=NUM_INPUTS)
    basal_ganglia = nengo.networks.BasalGanglia(n_neurons_per_ensemble=BASALGANGLIA_NEURONS, dimensions=NUM_ACTIONS)
    thalamus = nengo.networks.Thalamus(n_neurons_per_ensemble=THALAMUS_NEURONS, dimensions=NUM_ACTIONS)
    action_ensemble = nengo.Ensemble(n_neurons=DEFAULT_ENSEMBLE_NEURONS * NUM_ACTIONS, dimensions=NUM_ACTIONS)
    
    #Define connections between ensembles
    nengo.Connection(input_ensemble, basal_ganglia.input, synapse=None)
    nengo.Connection(basal_ganglia.output, thalamus.input, synapse=0.01)
    nengo.Connection(thalamus.output, action_ensemble, synapse=0.01)

    #Define model probes
    output_probe = nengo.Probe(thalamus.output, synapse=0.01)
    bas_probe = nengo.Probe(basal_ganglia.output, synapse=0.01)
    


with nengo.Simulator(model) as sim:
    sim.run(1)

Can someone please help me understand how I should design my network so that I can feed n inputs into it and have it select and action of of m possible actions assuming n does not equal m?

Hello,

To my knowledge, the BG-TH networks operate as a argmax (Winner-Take-All) function on a given dimensionality vector. As a result, to select one of the resulting dimensions in the BG/TH actions - ie) the one with the highest value, a transformation first needs to be applied to convert the inputs into actions in order for it to perform the WTA. In some literature, this intermediate transformation result is called the utility value and the action associated with the highest utility value is selected as the optimal decision. This would also convert the N dimensional input into the M possible actions.

So injecting a random utility transform matrix could be something such as below. The “argmax” is then performed over the resulting utility ensemble output.

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

#Global Variables
NUM_INPUTS = 6
NUM_ACTIONS = 4
DEFAULT_ENSEMBLE_NEURONS = 100
BASALGANGLIA_NEURONS = 100 #default = 100
THALAMUS_NEURONS = 50 #default = 50

model = nengo.Network()
with model:
    
    #Define model ensembles
    input_ensemble = nengo.Ensemble(n_neurons=DEFAULT_ENSEMBLE_NEURONS * NUM_INPUTS, dimensions=NUM_INPUTS)
    basal_ganglia = nengo.networks.BasalGanglia(n_neurons_per_ensemble=BASALGANGLIA_NEURONS, dimensions=NUM_ACTIONS)
    thalamus = nengo.networks.Thalamus(n_neurons_per_ensemble=THALAMUS_NEURONS, dimensions=NUM_ACTIONS)
    action_ensemble = nengo.Ensemble(n_neurons=DEFAULT_ENSEMBLE_NEURONS * NUM_ACTIONS, dimensions=NUM_ACTIONS)
    utility_ensemble = nengo.Ensemble(n_neurons=DEFAULT_ENSEMBLE_NEURONS * NUM_ACTIONS, dimensions=NUM_ACTIONS)
    
    #Define connections between ensembles
    nengo.Connection(input_ensemble, utility_ensemble, function=lambda x: np.dot(x, np.random.rand(NUM_INPUTS, NUM_ACTIONS)))
    nengo.Connection(utility_ensemble, basal_ganglia.input, synapse=None)
    nengo.Connection(basal_ganglia.output, thalamus.input, synapse=0.01)
    nengo.Connection(thalamus.output, action_ensemble, synapse=0.01)

    #Define model probes
    utility_probe = nengo.Probe(utility_ensemble, synapse=0.2)
    output_probe = nengo.Probe(thalamus.output, synapse=0.01)
    bas_probe = nengo.Probe(basal_ganglia.output, synapse=0.01)
    
with nengo.Simulator(model) as sim:
    sim.run(1)

#Plot utility value
plt.figure()
plt.plot(sim.trange(), sim.data[utility_probe])
plt.legend([f"Action {action_id} Utility" for action_id in range(NUM_ACTIONS)])
plt.title("Action Utility")

#Plot BG-TH output value
plt.figure()
plt.plot(sim.trange(), sim.data[output_probe])
plt.legend([f"Action {action_id}" for action_id in range(NUM_ACTIONS)])
plt.title("Thalamus Output")

Thanks for the reply, this is really helpful and has opened up a whole new avenue of research for me! Is it standard practice with SNNs to have a random utility function, or is there a way to guide the utility function with certain factors you want the model to incorporate? For example, if I want my model to control an agent to eat food and avoid predators, could I incorporate some kind of quantitative values for various factors such as energy cost, time, difficulty, etc. into my utility function to influence the model’s perceived utility of each action?

I’m glad you found that to be helpful. Certainly taking a random action is not an optimal solution but was moreso for the sake of building and running a sample model all the way through. I wouldn’t particularly attribute this to being a standard practice with SNNs but rather taking some concepts from other fields of research and utilizing the NEF to realize them in an SNN.

I think the main thing here is to note the functionality of the BG-TH network as a WTA/argmax component when incorporating it into an overall network. The argmax function is pretty common in decision making systems.

To comment on transformations that can be applied, this is a large field by itself, some classical approaches could include a creating decision-logic/decision tree, crafting a Q-table, state-transition tables, etc. More complex approaches add additional components such as learning and internal state which requires transformation adjustments (using weights typically for SNNs) for learning and/or memory to represent internal state. Reinforcement learning usually applies here. In many solutions, the temporal aspect is, for the most part, disregarded which is a key feature of SNNs - so more research needs to be done here in my opinion.