Model and simulation files for the paper "Enhancing reinforcement learning models by including direct and indirect pathways
improves performance on striatal dependent tasks", by Kim Blackwell and Kenji Doya, PLoS Comp Biol in press (2023)

To run the simulations, run one of the following programs in python3:
BanditTask.py
SequenceTask.py
Discrim2stateProbT_twoQtwoS.py
Note that due to stochastic nature of task, output will not be identical to that in published manuscript.

To change the output or simulation length/time, adjust the following values:
	* runs: number of agents, or how many times the task will be simulated. Equivalent to biological replicates
	* trials: number of actions the agent will take for a single run.  For a 3 step task (BanditTask and Discrim), the number of trials is approximately trials/3
	* save_data: set to False if you do not want to save the output to disk
	* plot_Qhx: higher number gives more graphs

See the README.md for more information about other parameter values.

To generate figures in the manuscript, use the simulation output files provided with the package, update a few parameters as follows and run the following programs in python3:

Figure 2A-B and 3A-B: set task='discrim' in TD2Q_manuscript_graphs.py, and then run TD2Q_manuscript_graphs.py 
Figure 2C-E: set task='discrim' in TD2Q_Qhx_graphs.py, and then run TD2Q_Qhx_graphs.py 
Figure 3C-G: set task='discrim' in TD2Q_Qhx_graphs.py, comment out this line (695): states=['Pport,6kHz 1'] and then run TD2Q_Qhx_graphs.py 
Figure 4A-C: set task='AIP' in TD2Q_manuscript_graphs.py, and test='AIP' in discrimFiles.py, then run TD2Q_manuscript_graphs.py 
Figure 4D-F: set task='AIP' in TD2Q_Qhx_graphs.py and then run TD2Q_Qhx_graphs.py 
Figure 6: set task='Bandit' in TD2Q_Qhx_graphs.py, and then run TD2Q_Qhx_graphs.py 
Figure 7B: set task='Bandit' in TD2Q_manuscript_graphs.py, and then run TD2Q_manuscript_graphs.py 
Figure 9A-D: set task='sequence' in TD2Q_manuscript_graphs.py, and then run TD2Q_manuscript_graphs.py 
Figure 9E: set task='sequence' and add_barplot=1 in TD2Q_manuscript_graphs.py, set test='inact' in sequenceFiles.py, then run TD2Q_manuscript_graphs.py 
Figure 10: set task='sequence' in TD2Q_Qhx_graphs.py, and then run TD2Q_Qhx_graphs.py