# PyTorch Implementation of Differentiable ODE Solvers

This library provides ordinary differential equation (ODE) solvers implemented in PyTorch. Backpropagation through all solvers is supported using the adjoint method. For usage of ODE solvers in deep learning applications, see [1].

As the solvers are implemented in PyTorch, algorithms in this repository are fully supported to run on the GPU.

---

## Installation
```
git clone https://github.com/rtqichen/torchdiffeq.git
cd torchdiffeq
pip install -e .
```

## Examples
Examples are placed in the [`examples`](./examples) directory.

We encourage those who are interested in using this library to take a look at [`examples/ode_demo.py`](./examples/ode_demo.py) for understanding how to use `torchdiffeq` to fit a simple spiral ODE.

## Basic usage
This library provides one main interface `odeint` which contains general-purpose algorithms for solving initial value problems (IVP), with gradients implemented for all main arguments. An initial value problem consists of an ODE and an initial value,
```
dy/dt = f(t, y)    y(t_0) = y_0.
```
The goal of an ODE solver is to find a continuous trajectory satisfying the ODE that passes through the initial condition.

To solve an IVP using the default solver:
```
from torchdiffeq import odeint

odeint(func, y0, t)
```
where `func` is any callable implementing the ordinary differential equation `f(t, x)`, `y0` is an _any_-D Tensor or a tuple of _any_-D Tensors representing the initial values, and `t` is a 1-D Tensor containing the evaluation points. The initial time is taken to be `t[0]`.

Backpropagation through `odeint` goes through the internals of the solver, but this is not supported for all solvers. Instead, we encourage the use of the adjoint method explained in [1], which will allow solving with as many steps as necessary due to O(1) memory usage.

To use the adjoint method:
```
from torchdiffeq import odeint_adjoint as odeint

odeint(func, y0, t)
```
`odeint_adjoint` simply wraps around `odeint`, but will use only O(1) memory in exchange for solving an adjoint ODE in the backward call.

The biggest **gotcha** is that `func` must be a `nn.Module` when using the adjoint method. This is used to collect parameters of the differential equation.

### Keyword Arguments
 - `rtol` Relative tolerance.
 - `atol` Absolute tolerance.
 - `method` One of the solvers listed below.

#### List of ODE Solvers:

Adaptive-step:
 - `dopri5` Runge-Kutta 4(5) [default].
 - `adams` Adaptive-order implicit Adams.

Fixed-step:
 - `euler` Euler method.
 - `midpoint` Midpoint method.
 - `rk4` Fourth-order Runge-Kutta with 3/8 rule.
 - `explicit_adams` Explicit Adams.
 - `fixed_adams` Implicit Adams.

### References
[1] Ricky T. Q. Chen, Yulia Rubanova, Jesse Bettencourt, David Duvenaud. "Neural Ordinary Differential Equations." *Advances in Neural Processing Information Systems.* 2018. [[arxiv]](https://arxiv.org/abs/1806.07366)

---

If you found this library useful in your research, please consider citing
```
@article{chen2018neural,
  title={Neural Ordinary Differential Equations},
  author={Chen, Ricky T. Q. and Rubanova, Yulia and Bettencourt, Jesse and Duvenaud, David},
  journal={Advances in Neural Information Processing Systems},
  year={2018}
}
```

multivariate time series forecasting
point processes
multivariate time series imputation
time series


This repository provides a set of tools and libraries for solving ordinary differential equations (ODEs) using numerical methods. The `Dopri5` and `VariableCoefficientAdamsBashforth` classes in `torchdiffeq/_impl/dopri5.py` and `torchdiffeq/_impl/adams.py`, respectively, implement adaptive stepsize integration methods for solving ODEs. The `OdeSolver` class in `torchdiffeq/_impl/ode_solver.py` defines the interface for all ODE solvers in TorchDif.

The data used by this repository is likely to be related to scientific or engineering applications that involve solving ODEs, such as modeling physical systems, predicting population dynamics, or optimizing control systems. The repository tackles the problem of solving ordinary differential equations (ODEs) using numerical methods and provides a set of tools and libraries for doing so.


# Parameters from Tsitouras (2011).
...
class Tsit5Solver(AdaptiveStepsizeODESolver):
...

    def before_integrate(self, t):
...

    def advance(self, next_t):
...

    def _adaptive_tsit5_step(self, rk_state):
...
########################################################
#                      Assertions                      #
########################################################
...
########################################################
#                     Error Ratio                      #
########################################################

# approximately equal to 6pi
...
# add 1 all timestamps to avoid division by 0
...
# generate clock-wise and counter clock-wise spirals in observation space
# with two sets of time-invariant latent dynamics
...
# sample starting timestamps
...
# don't sample t0 very near the start or the end
...
# uniformly select rotation
...
# batching for sample trajectories is good for RNN; batching for original
# trajectories only for ease of indexing
...
class LatentODEfunc(nn.Module):
...

    def __init__(self, latent_dim=4, nhidden=20):
...

    def forward(self, t, x):
...
class RecognitionRNN(nn.Module):
...

    def __init__(self, latent_dim=4, obs_dim=2, nhidden=25, nbatch=1):
...

    def forward(self, x, h):
...

    def initHidden(self):
...
class Decoder(nn.Module):
...

    def __init__(self, latent_dim=4, obs_dim=2, nhidden=20):
...

    def forward(self, z):
...
class RunningAverageMeter(object):
...

    def __init__(self, momentum=0.99):

    def reset(self):

    def update(self, val):
...
# generate toy spiral data
...
# model
...
# backward in time to infer q(z_0)
...
# forward in time and solve ode for reconstructions
...
# compute loss
...
# sample from trajectorys' approx. posterior

class TestSolverError(unittest.TestCase):
...

    def test_euler(self):
...

    def test_midpoint(self):
...

    def test_rk4(self):
...

    def test_explicit_adams(self):
...

    def test_adams(self):
...

    def test_dopri5(self):
...

    def test_adjoint(self):
...
class TestSolverBackwardsInTimeError(unittest.TestCase):
...

    def test_euler(self):
...

    def test_midpoint(self):
...

    def test_rk4(self):
...

    def test_explicit_adams(self):
...

    def test_adams(self):
...

    def test_dopri5(self):
...

    def test_adjoint(self):
...
class TestNoIntegration(unittest.TestCase):
...

    def test_midpoint(self):
...

    def test_rk4(self):
...

    def test_explicit_adams(self):
...

    def test_adams(self):

class ResBlock(nn.Module):
...

    def __init__(self, inplanes, planes, stride=1, downsample=None):
...

    def forward(self, x):
...
class ConcatConv2d(nn.Module):
...

    def __init__(self, dim_in, dim_out, ksize=3, stride=1, padding=0, dilation=1, groups=1, bias=True, transpose=False):
...

    def forward(self, t, x):
...
class ODEfunc(nn.Module):
...

    def __init__(self, dim):
...

    def forward(self, t, x):
...
class ODEBlock(nn.Module):
...

    def __init__(self, odefunc):
...

    def forward(self, x):
...

    def nfe(self):
...

    def nfe(self, value):
...
class Flatten(nn.Module):
...

    def __init__(self):
...

    def forward(self, x):
...
class RunningAverageMeter(object):
...

    def __init__(self, momentum=0.99):

    def reset(self):

    def update(self, val):

class TestGradient(unittest.TestCase):
...

    def test_midpoint(self):
...

    def test_rk4(self):
...

    def test_dopri5(self):
...

    def test_adams(self):
...

    def test_adjoint(self):
...
# reg_y0_grad = y0.grad
...
# adj_y0_grad = y0.grad
...
# self.assertLess(max_abs(reg_y0_grad - adj_y0_grad), eps)
...
class TestCompareAdjointGradient(unittest.TestCase):
...

    def problem(self):
class Odefunc(torch.nn.Module):
...

            def __init__(self):
...

            def forward(self, t, y):
...

    def test_dopri5_adjoint_against_dopri5(self):

class _VCABMState(collections.namedtuple('_VCABMState', 'y_n, prev_f, prev_t, next_t, phi, order')):
...
class VariableCoefficientAdamsBashforth(AdaptiveStepsizeODESolver):
...

    def before_integrate(self, t):
...

    def advance(self, final_t):
...

    def _adaptive_adams_step(self, vcabm_state, final_t):
...
# Explicit predictor step.
...
# Update phi to implicit.
...
# Implicit corrector step.
...
# Error estimation.
...
# Retry with adjusted step size if step is rejected.
...
# We accept the step. Evaluate f and update phi.

class TestCollectionState(unittest.TestCase):
...

    def test_dopri5(self):
...

    def test_dopri5_gradient(self):
...

    def test_adams(self):

class Lambda(nn.Module):
...

    def forward(self, t, y):
...
# (M, D)
# (T)
# (T, M, D)
...
class ODEFunc(nn.Module):
...

    def __init__(self):
...

    def forward(self, t, y):
...
class RunningAverageMeter(object):

class ConstantODE(torch.nn.Module):
...

    def __init__(self, device):
...

    def forward(self, t, y):
...

    def y_exact(self, t):
...
class SineODE(torch.nn.Module):
...

    def __init__(self, device):
...

    def forward(self, t, y):
...

    def y_exact(self, t):
...
class LinearODE(torch.nn.Module):
...

    def __init__(self, device, dim=10):
...

    def forward(self, t, y):
...

    def y_exact(self, t):


Using summaries of '000Justin000/torchdiffeq' files from Context, we can infer that this repository tackles the problem of solving ordinary differential equations (ODEs) using numerical methods. The `Dopri5` and `VariableCoefficientAdamsBashforth` classes in `torchdiffeq/_impl/dopri5.py` and `torchdiffeq/_impl/adams.py`, respectively, implement adaptive stepsize integration methods for solving ODEs. The `OdeSolver` class in `torchdiffeq/_impl/ode_solver.py` defines the interface for all ODE solvers in TorchDif.

Therefore, this repository provides a set of tools and libraries for solving ordinary differential equations using numerical methods, such as the Dormand-Prince (Dopri5) and Adams-Bashforth (ABM) methods. The data used by this repository is likely to be related to scientific or engineering applications that involve solving ODEs, such as modeling physical systems, predicting population dynamics, or optimizing control systems.


1. `torchdiffeq/_impl/dopri5.py`: This file contains the implementation of the Dormand-Prince (Dopri5) ODE solver, which is a widely used numerical method for solving ordinary differential equations (ODEs). The code defines the `Dopri5` class that inherits from the `AdaptiveStepsizeODESolver` class and implements the necessary methods to perform adaptive stepsize integration.
2. `torchdiffeq/_impl/adams.py`: This file contains the implementation of the Adams-Bashforth (ABM) ODE solver, which is another popular numerical method for solving ODEs. The code defines the `VariableCoefficientAdamsBashforth` class that inherits from the `AdaptiveStepsizeODESolver` class and implements the necessary methods to perform adaptive stepsize integration with variable coefficients.
3. `torchdiffeq/_impl/ode_solver.py`: This file contains the implementation of the base `OdeSolver` class, which defines the interface for all ODE solvers in TorchDif


examples/latent_ode.py, examples/odenet_mnist.py, examples/ode_demo.py, torchdiffeq/_impl/adams.py, torchdiffeq/_impl/tsit5.py, generate_spiral2d, TestSolverBackwardsInTimeError, construct_problem, ConstantODE, SineODE, batch_time, y_exact, TupleFunc, _decreasing, _interp_coeff_tsit5


examples/latent_ode.py, examples/odenet_mnist.py, examples/ode_demo.py, torchdiffeq/_impl/adams.py, torchdiffeq/_impl/tsit5.py, generate_spiral2d, TestSolverBackwardsInTimeError, construct_problem, ConstantODE, SineODE, batch_time, y_exact, TupleFunc, _decreasing, _interp_coeff_tsit5## tasks:
#Lipschitz control, Backwards differentiation, Poisson equation(dense


Lipschitz control, Backwards differentiation, Poisson equation(dense


examples/latent_ode.py:
⋮...
│class RunningAverageMeter(object):
│    """Computes and stores the average and current value"""
│
⋮...
│    def update(self, val):
⋮...

examples/ode_demo.py:
⋮...
│def makedirs(dirname):
⋮...
│class RunningAverageMeter(object):
│    """Computes and stores the average and current value"""
│
⋮...
│    def update(self, val):
⋮...

examples/odenet_mnist.py:
⋮...
│def norm(dim):
⋮...
│def makedirs(dirname):
⋮...

tests/problems.py:
⋮...
│class ConstantODE(torch.nn.Module):
│
│    def __init__(self, device):
│        super(ConstantODE, self).__init__()
│        self.a = torch.nn.Parameter(torch.tensor(0.2).to(device))
⋮...
│    def y_exact(self, t):
⋮...
│class SineODE(torch.nn.Module):
│
│    def __init__(self, device):
⋮...
│    def y_exact(self, t):
⋮...
│class LinearODE(torch.nn.Module):
│
│    def __init__(self, device, dim=10):
│        super(LinearODE, self).__init__()
│        self.dim = dim
│        U = torch.randn(dim, dim).to(device) * 0.1
│        A = 2 * U - (U + U.transpose(0, 1))
│        self.A = torch.nn.Parameter(A)
⋮...
│    def y_exact(self, t):
⋮...

torchdiffeq/_impl/adams.py:
⋮...
│def compute_implicit_phi(explicit_phi, f_n, k):
⋮...

torchdiffeq/_impl/adjoint.py:
⋮...
│class OdeintAdjointMethod(torch.autograd.Function):
│
│    @staticmethod
│    def forward(ctx, *args):
│        assert len(args) >= 8, 'Internal error: all arguments required.'
│        y0, func, t, flat_params, rtol, atol, method, options = \
│            args[:-7], args[-7], args[-6], args[-5], args[-4], args[-3], args[-2], args[-1]
│
│        ctx.func, ctx.rtol, ctx.atol, ctx.method, ctx.options = func, rtol, atol, method, options
│
│        with torch.no_grad():
│            ans = odeint(func, y0, t, rtol=rtol, atol=atol, method=method, options=options)
⋮...
│    @staticmethod
│    def backward(ctx, *grad_output):
│
⋮...

torchdiffeq/_impl/misc.py:
⋮...
│def _scaled_dot_product(scale, xs, ys):
⋮...
│def _convert_to_tensor(a, dtype=None, device=None):
⋮...
│def _is_iterable(inputs):
⋮...

torchdiffeq/_impl/odeint.py:
⋮...
│def odeint(func, y0, t, rtol=1e-7, atol=1e-9, method=None, options=None):
⋮...

torchdiffeq/_impl/rk_common.py:
⋮...
│def rk4_alt_step_func(func, t, dt, y, k1=None):
⋮...

torchdiffeq/_impl/solvers.py:
⋮...
│class AdaptiveStepsizeODESolver(object, metaclass=abc.ABCMeta):
│    def __init__(self, func, y0, atol, rtol, **unused_kwargs):
│        _handle_unused_kwargs(self, unused_kwargs)
│        del unused_kwargs
│
│        self.func = func
│        self.y0 = y0
│        self.atol = atol
⋮...
│    def integrate(self, t):
⋮...
│class FixedGridODESolver(object, metaclass=abc.ABCMeta):
│    def __init__(self, func, y0, step_size=None, grid_constructor=None, **unused_kwargs):
│        unused_kwargs.pop('rtol', None)
│        unused_kwargs.pop('atol', None)
│        _handle_unused_kwargs(self, unused_kwargs)
│        del unused_kwargs
│
│        self.func = func
│        self.y0 = y0
│
│        if step_size is not None and grid_constructor is None:
⋮...
│    def integrate(self, t):
⋮...


TIMIT_preparation.py:
⋮...
│def ReadList(list_file):
⋮...
│def copy_folder(in_folder,out_folder):
⋮...

compute_d_vector.py

data_io.py:
⋮...
│def ReadList(list_file):
⋮...
│def read_conf():
│
⋮...
│def str_to_bool(s):
⋮...
│def create_batches_rnd(batch_size,data_folder,wav_lst,N_snt,wlen,lab_dict,fact_amp):
│
⋮...
│def read_conf_inp(cfg_file):
│
⋮...

dnn_models.py:
⋮...
│def flip(x, dim):
⋮...
│def sinc(band,t_right):
⋮...
│class SincConv_fast(nn.Module):
│    """Sinc-based convolution
│    Parameters
│    ----------
│    in_channels : `int`
│        Number of input channels. Must be 1.
│    out_channels : `int`
│        Number of filters.
│    kernel_size : `int`
│        Filter length.
│    sample_rate : `int`, optional
⋮...
│    @staticmethod
│    def to_mel(hz):
⋮...
│    @staticmethod
│    def to_hz(mel):
⋮...
│    def __init__(self, out_channels, kernel_size, sample_rate=16000, in_channels=1,
⋮...
│class sinc_conv(nn.Module):
│
│    def __init__(self, N_filt,Filt_dim,fs):
│        super(sinc_conv,self).__init__()
│
│        # Mel Initialization of the filterbanks
│        low_freq_mel = 80
│        high_freq_mel = (2595 * np.log10(1 + (fs / 2) / 700))  # Convert Hz to Mel
│        mel_points = np.linspace(low_freq_mel, high_freq_mel, N_filt)  # Equally spaced in Mel scal
│        f_cos = (700 * (10**(mel_points / 2595) - 1)) # Convert Mel to Hz
│        b1=np.roll(f_cos,1)
│        b2=np.roll(f_cos,-1)
⋮...
│def act_fun(act_type):
│
⋮...
│class LayerNorm(nn.Module):
│
│    def __init__(self, features, eps=1e-6):
│        super(LayerNorm,self).__init__()
│        self.gamma = nn.Parameter(torch.ones(features))
│        self.beta = nn.Parameter(torch.zeros(features))
⋮...
│class MLP(nn.Module):
│    def __init__(self, options):
│        super(MLP, self).__init__()
│
│        self.input_dim=int(options['input_dim'])
│        self.fc_lay=options['fc_lay']
│        self.fc_drop=options['fc_drop']
│        self.fc_use_batchnorm=options['fc_use_batchnorm']
│        self.fc_use_laynorm=options['fc_use_laynorm']
│        self.fc_use_laynorm_inp=options['fc_use_laynorm_inp']
│        self.fc_use_batchnorm_inp=options['fc_use_batchnorm_inp']
⋮...
│class SincNet(nn.Module):
│
│    def __init__(self,options):
│       super(SincNet,self).__init__()
│
│       self.cnn_N_filt=options['cnn_N_filt']
│       self.cnn_len_filt=options['cnn_len_filt']
│       self.cnn_max_pool_len=options['cnn_max_pool_len']
│
│
│       self.cnn_act=options['cnn_act']
│       self.cnn_drop=options['cnn_drop']
│
⋮...

inference.py

similarity.py

speaker_id.py:
⋮...
│def create_batches_rnd(batch_size,data_folder,wav_lst,N_snt,wlen,lab_dict,fact_amp):
│
⋮...


speech emotion recognition, agents.py This is the main source code

speaker recognition
speaker identification
speaker verification


compute_d_vector.py, speaker_id.py, dnn_models.py, similarity.py, TIMIT_preparation.py, read_conf_inp, sinc_conv, SincConv_fast, SincNet, ig_f, OptionParser, OptionParser, cfg, ConfigParser, ConfigParser## tasks:
#speech emotion recognition, agents.py This is the main source code


compute_d_vector.py, speaker_id.py, dnn_models.py, similarity.py, TIMIT_preparation.py, read_conf_inp, sinc_conv, SincConv_fast, SincNet, ig_f, OptionParser, OptionParser, cfg, ConfigParser, ConfigParser


# compute_d_vector.py
# Mirco Ravanelli
# Mila - University of Montreal
...
# Feb 2019
...
# Description:
# This code computes d-vectors using a pre-trained model
...
# Model to use for computing the d-vectors
# This is the model to use for computing the d-vectors (it should be pre-trained using the speaker-id DNN)
# Config file of the speaker-id experiment used to generate the model
#te_lst='data_lists/TIMIT_test.scp' # List of the wav files to process
...
# output dictionary containing the a sentence id as key as the d-vector as value
...
# Avoid frames with an energy that is 1/10 over the average energy
...
#device = None
...
# Reading cfg file
...
#[data]
...
#[windowing]
...
#[cnn]
...
#[dnn]
...
#[class]
...
# Folder creation
...
# loss function
...
# Converting context and shift in samples
...
# Batch_dev
...
# Feature extractor CNN
...
# Amplitude normalization
...
# computing energy on each frame:
...
# split signals into chunks
...
# averaging and normalizing all the d-vectors
...
# checks for nan
...
# saving the d-vector in a numpy dictionary

#import scipy.io.wavfile
...
# Mandatory
...
#[data]
...
#[windowing]
...
#[cnn]
...
#[dnn]
...
#[class]
...
#[optimization]
...
# Initialization of the minibatch (batch_size,[0=>x_t,1=>x_t+N,1=>random_samp])
...
# select a random sentence from the list  (joint distribution)
...
# accesing to a random chunk
...
#randint(0, snt_len-2*wlen-1)
...
# Current Frame
...
#[data]
...
#[windowing]
...
#[cnn]
...
#[dnn]
...
#[class]

# speaker_id.py
# Mirco Ravanelli
# Mila - University of Montreal
...
# July 2018
...
# Description:
# This code performs a speaker_id experiments with SincNet.
...
# How to run it:
# python speaker_id.py --cfg=cfg/SincNet_TIMIT.cfg
...
#import scipy.io.wavfile
...
# Initialization of the minibatch (batch_size,[0=>x_t,1=>x_t+N,1=>random_samp])
...
# select a random sentence from the list
#[fs,signal]=scipy.io.wavfile.read(data_folder+wav_lst[snt_id_arr[i]])
#signal=signal.astype(float)/32768
...
# accesing to a random chunk
...
#randint(0, snt_len-2*wlen-1)
...
# Reading cfg file
...
#[data]
...
#[windowing]
...
#[cnn]
...
#[dnn]
...
#[class]
...
#[optimization]
...
# training list
...
# test list
...
# Folder creation
...
# setting seed
...
# loss function
...
# Converting context and shift in samples
...
# Batch_dev
...
# Feature extractor CNN
...
# Loading label dictionary
...
# Full Validation  new
...
#[fs,signal]=scipy.io.wavfile.read(data_folder+wav_lst_te[i])
#signal=signal.astype(float)/32768

#di1=np.load("d_vect_speaker_f.npy",allow_pickle=True).item()
...
#a=[di1["3"]]
...
#di2=np.load("d_vect_speaker_m.npy",allow_pickle=True).item()

# test
...
#user_label={0:"adarsh",1:"anuj",2:"piyush",3:"peeyush",4:"sameer",5:"rajat",6:"rachna",7:"sangram",8:"shashikant",9:"karan",10:"eram",11:"anjani",12:"akash"}
#[windowing]
...
#[cnn]
...
#[dnn]
...
#[class]
...
#[optimization]
...
# setting seed
...
# loss function
...
# Converting context and shift in samples
...
# Batch_dev

class SincConv_fast(nn.Module):
...

    def to_mel(hz):
...

    def to_hz(mel):
...
#msg = (f'SincConv only support one input channel '
#       f'(here, in_channels = {in_channels:d}).')
...
# Forcing the filters to be odd (i.e, perfectly symmetrics)
...
# initialize filterbanks such that they are equally spaced in Mel scale
...
# filter lower frequency (out_channels, 1)
...
# filter frequency band (out_channels, 1)
...
# Hamming window
#self.window_ = torch.hamming_window(self.kernel_size)
# computing only half of the window
...
# (1, kernel_size/2)
...

    def forward(self, waveforms):
...
# Due to symmetry, I only need half of the time axes
...
# Equivalent of Eq.4 of the reference paper (SPEAKER RECOGNITION FROM RAW WAVEFORM WITH SINCNET). I just have expanded the sinc and simplified the terms. This way I avoid several useless computations.
...
class sinc_conv(nn.Module):
...

    def __init__(self, N_filt,Filt_dim,fs):
...
# Mel Initialization of the filterbanks
...
# Convert Hz to Mel
# Equally spaced in Mel scale
# Convert Mel to Hz
...

    def forward(self, x):
...
# Filter window (hamming)
...
# initializzed like this, but not used in forward!
...
class LayerNorm(nn.Module):
...

    def __init__(self, features, eps=1e-6):
...

    def forward(self, x):
...
class MLP(nn.Module):
...

    def __init__(self, options):
...
# input layer normalization
...
# input batch normalization
...
# Initialization of hidden layers
...
# dropout
...
# activation
...
# layer norm initialization
...
# Linear operations
...
# weight initialization
...

    def forward(self, x):
...
# Applying Layer/Batch Norm
...
class SincNet(nn.Module):
...

    def __init__(self,options):
...
# dropout
...
# activation
...
# layer norm initialization

#!/usr/bin/env python3
...
# TIMIT_preparation
# Mirco Ravanelli
# Mila - University of Montreal
...
# July 2018
...
# Description:
# This code prepares TIMIT for the following speaker identification experiments.
# It removes start and end silences according to the information reported in the *.wrd files and normalizes the amplitude of each sentence.
...
# How to run it:
# python TIMIT_preparation.py $TIMIT_FOLDER $OUTPUT_FOLDER data_lists/TIMIT_all.scp
...
# NOTE: This script expects filenames in lowercase (e.g, train/dr1/fcjf0/si1027.wav" rather than "TRAIN/DR1/FCJF0/SI1027.WAV)
...
# Read List file
...
# Replicate input folder structure to output folder
...
# Speech Data Reverberation Loop
...
# Open the wav file
...
# Signal normalization
...
# Read wrd file
...
# Remove silences


This repository tackles the problem of speaker recognition from raw waveform using a PyTorch implementation of the SincNet model. The data used is the TIMIT dataset, which contains speech recordings with different speakers and sentences. The repository provides a custom PyTorch module for the convolutional layer in SincNet, which uses the sinc activation function to compute the output of each filter. Additionally, it includes a script for preparing the TIMIT dataset for speaker identification experiments by removing start and end silences according to the information in the *.wrd files and normalizing the amplitude of each sentence.


This repository tackles the problem of speaker recognition from raw waveform using a PyTorch implementation of the SincNet model. The data used is the TIMIT dataset, which contains speech recordings with different speakers and sentences.


# SincNet
SincNet is a neural architecture for processing **raw audio samples**. It is a novel Convolutional Neural Network (CNN) that encourages the first convolutional layer to discover more **meaningful filters**. SincNet is based on parametrized sinc functions, which implement band-pass filters.

In contrast to standard CNNs, that learn all elements of each filter, only low and high cutoff frequencies are directly learned from data with the proposed method. This offers a very compact and efficient way to derive a **customized filter bank** specifically tuned for the desired application.

This project releases a collection of codes and utilities to perform speaker identification with SincNet.
An example of speaker identification with the TIMIT database is provided. If you are interested in **SincNet applied to speech recognition you can take a look into the PyTorch-Kaldi github repository (https://github.com/mravanelli/pytorch-kaldi).**

<img src="https://github.com/mravanelli/SincNet/blob/master/SincNet.png" width="400" img align="right">

[Take a look into our video introduction to SincNet](https://www.youtube.com/watch?v=mXQBObRGUgk&feature=youtu.be)

## Cite us
If you use this code or part of it, please cite us!

*Mirco Ravanelli, Yoshua Bengio, âSpeaker Recognition from raw waveform with SincNetâ* [Arxiv](http://arxiv.org/abs/1808.00158)

## Prerequisites
- Linux
- Python 3.6/2.7
- pytorch 1.0
- pysoundfile (``` conda install -c conda-forge pysoundfile```)
- We also suggest using the anaconda environment.

## Updates
Feb, 16 2019:
- We replaced the old "sinc_conv"  with "SincConv_fast". The latter is 50% faster.
- In the near future, we plan to support SincNet based speaker-id within the [PyTorch-Kaldi project](https://github.com/mravanelli/pytorch-kaldi) (the current version of the project only supports SincNEt for speech recognition experiments). This will allow users to perform speaker recognition experiments in a faster and much more flexible environment. The current repository will anyway remain as a showcase.

## How to run a TIMIT experiment
Even though the code can be easily adapted to any speech dataset, in the following part of the documentation we provide an example based on the popular TIMIT dataset.

**1. Run TIMIT data preparation.**

This step is necessary to store a version of TIMIT in which start and end silences are removed and the amplitude of each speech utterance is normalized. To do it, run the following code:

``
python TIMIT_preparation.py $TIMIT_FOLDER $OUTPUT_FOLDER data_lists/TIMIT_all.scp
``

where:
- *$TIMIT_FOLDER* is the folder of the original TIMIT corpus
- *$OUTPUT_FOLDER* is the folder in which the normalized TIMIT will be stored
- *data_lists/TIMIT_all.scp* is the list of the TIMIT files used for training/test the speaker id system.

**2. Run the speaker id experiment.**

- Modify the *[data]* section of *cfg/SincNet_TIMIT.cfg* file according to your paths. In particular, modify the *data_folder* with the *$OUTPUT_FOLDER* specified during the TIMIT preparation. The other parameters of the config file belong to the following sections:
 1. *[windowing]*, that defines how each sentence is split into smaller chunks.
 2. *[cnn]*,  that specifies the characteristics of the CNN architecture.
 3. *[dnn]*,  that specifies the characteristics of the fully-connected DNN architecture following the CNN layers.
 4. *[class]*, that specify the softmax classification part.
 5. *[optimization]*, that reports the main hyperparameters used to train the architecture.

- Once setup the cfg file, you can run the speaker id experiments using the following command:

``
python speaker_id.py --cfg=cfg/SincNet_TIMIT.cfg
``

The network might take several hours to converge (depending on the speed of your GPU card). In our case, using an *nvidia TITAN X*, the full training took about 24 hours. If you use the code within a cluster is crucial to copy the normalized dataset into the local node, since the current version of the code requires frequent accesses to the stored wav files. Note that several possible optimizations to improve the code speed are not implemented in this version since are out of the scope of this work.

**3. Results.**

The results are saved into the *output_folder* specified in the cfg file. In this folder, you can find a file (*res.res*) summarizing training and test error rates. The model *model_raw.pkl* is the SincNet model saved after the last iteration.
Using the cfg file specified above, we obtain the following results:
```
epoch 0, loss_tr=5.542032 err_tr=0.984189 loss_te=4.996982 err_te=0.969038 err_te_snt=0.919913
epoch 8, loss_tr=1.693487 err_tr=0.434424 loss_te=2.735717 err_te=0.612260 err_te_snt=0.069264
epoch 16, loss_tr=0.861834 err_tr=0.229424 loss_te=2.465258 err_te=0.520276 err_te_snt=0.038240
epoch 24, loss_tr=0.528619 err_tr=0.144375 loss_te=2.948707 err_te=0.534053 err_te_snt=0.062049
epoch 32, loss_tr=0.362914 err_tr=0.100518 loss_te=2.530276 err_te=0.469060 err_te_snt=0.015152
epoch 40, loss_tr=0.267921 err_tr=0.076445 loss_te=2.761606 err_te=0.464799 err_te_snt=0.023088
epoch 48, loss_tr=0.215479 err_tr=0.061406 loss_te=2.737486 err_te=0.453493 err_te_snt=0.010823
epoch 56, loss_tr=0.173690 err_tr=0.050732 loss_te=2.812427 err_te=0.443322 err_te_snt=0.011544
epoch 64, loss_tr=0.145256 err_tr=0.043594 loss_te=2.917569 err_te=0.438507 err_te_snt=0.009380
epoch 72, loss_tr=0.128894 err_tr=0.038486 loss_te=3.009008 err_te=0.438005 err_te_snt=0.019481
....
epoch 320, loss_tr=0.033052 err_tr=0.009639 loss_te=4.076542 err_te=0.416710 err_te_snt=0.006494
epoch 328, loss_tr=0.033344 err_tr=0.010117 loss_te=3.928874 err_te=0.415024 err_te_snt=0.007215
epoch 336, loss_tr=0.033228 err_tr=0.010166 loss_te=4.030224 err_te=0.410034 err_te_snt=0.005051
epoch 344, loss_tr=0.033313 err_tr=0.010166 loss_te=4.402949 err_te=0.428691 err_te_snt=0.009380
epoch 352, loss_tr=0.031828 err_tr=0.009238 loss_te=4.080747 err_te=0.414066 err_te_snt=0.006494
epoch 360, loss_tr=0.033095 err_tr=0.009600 loss_te=4.254683 err_te=0.419954 err_te_snt=0.005772
```
The converge is initially very fast (see the first 30 epochs). After that the performance improvement decreases and oscillations into the sentence error rate performance appear. Despite these oscillations an average improvement trend can be observed for the subsequent epochs. In this experiment, we stopped our training  at epoch 360.
The fields of the res.res file have the following meaning:
- loss_tr: is the average training loss (i.e., cross-entropy function) computed at every frame.
- err_tr: is the classification error (measured at frame level) of the training data. Note that we split the speech signals into chunks of 200ms with 10ms overlap. The error is averaged for all the chunks of the training dataset.
- loss_te is the average test loss (i.e., cross-entropy function) computed at every frame.
- err_te: is the classification error (measured at frame level) of the test data.
- err_te_snt: is the classification error (measured at sentence level) of the test data. Note that we split the speech signal into chunks of 200ms with 10ms overlap. For each chunk, our SincNet performs a prediction over the set of speakers. To compute this classification error rate we averaged the predictions and, for each sentence, we voted for the speaker with the highest average probability.

[You can find our trained model for TIMIT here.](https://bitbucket.org/mravanelli/sincnet_models/)

## Where SincNet is implemented?
To take a look into the SincNet implementation you should open the file *dnn_models.py* and read the classes *SincNet*, *sinc_conv* and the function *sinc*.

## How to use SincNet with a different dataset?
In this repository, we used the TIMIT dataset as a tutorial to show how SincNet works.
With the current version of the code, you can easily use a different corpus. To do it you should provide in input the corpora-specific input files (in wav format) and your own labels. You should thus modify the paths into the *.scp files you find in the data_lists folder.

To assign to each sentence the right label, you also have to modify the dictionary "*TIMIT_labels.npy*".
The labels are specified within a python dictionary that contains sentence ids as keys (e.g., "*si1027*") and speaker_ids as values. Each speaker_id is an integer, ranging from 0 to N_spks-1. In the TIMIT dataset, you can easily retrieve the speaker id from the path (e.g., *train/dr1/fcjf0/si1027.wav* is the sentence_id "*si1027*" uttered by the speaker "*fcjf0*"). For other datasets, you should be able to retrieve in such a way this dictionary containing pairs of speakers and sentence ids.

You should then modify the config file (*cfg/SincNet_TIMIT.cfg*) according to your new paths. Remember also to change the field "*class_lay=462*" according to the number of speakers N_spks you have in your dataset.

**The version of the Librispeech dataset used in the paper is available upon request**. In our work, we have used only 12-15 seconds of training material for each speaker and we processed the original librispeech sentences in order to perform amplitude normalization. Moreover, we used a simple energy-based VAD to avoid silences at the beginning and end of each sentence as well as to split in multiple chunks the sentences that contain longer silence

## References

[1] Mirco Ravanelli, Yoshua Bengio, âSpeaker Recognition from raw waveform with SincNetâ [Arxiv](http://arxiv.org/abs/1808.00158)


1. `sincnet.py`: This file contains a PyTorch implementation of the SincNet model for speaker recognition from raw waveform. It defines the architecture of the model, including the number of filters, filter dimensions, and the sinc activation function used in the convolutional layers. The forward pass computes the output of the network by applying the sinc activation function to the input waveform and then passing it through a series of linear layers with batch normalization and dropout.
2. `sinc_conv.py`: This file defines a custom PyTorch module for the convolutional layer in SincNet, which uses the sinc activation function to compute the output of each filter. The module takes the input waveform as an argument and applies the sinc activation function to it, followed by a series of linear layers with batch normalization and dropout.
3. `TIMIT_preparation.py`: This file prepares the TIMIT dataset for speaker identification experiments by removing start and end silences according to the information in the *.wrd files and normalizing the amplitude of each sentence. It reads a list file containing all the filenames in the dataset, replicates the input folder structure


alg/alg_baseline.py, alg/alg_baseline_checkers.py, alg/alg_credit.py, alg/alg_credit_checkers.py, alg/alg_qmix_checkers.py, Q_global, Q_global_1output, Q_global_sumo, Q_global_checkers, Q_credit, list_initialize_credit_ops, check_actions, check_actions, list_initialize_target_ops, test_checkers## tasks:
#reinforcement learning, convolution network, avatai,

efficient exploration
autonomous vehicles
multi agent reinforcement learning
reinforcement learning


alg/networks.py:
⋮...
│def fc3(t_input, n_hidden1=64, n_hidden2=64, n_outputs=9,
⋮...
│def convnet_1(t_input, f1=4, k1=[5,3], s1=[1,1], scope='convnet_1'):
⋮...
│def get_variable(name, shape):
│
⋮...

alg/replay_buffer.py:
⋮...
│class Replay_Buffer():
│
│    def __init__(self, size=1e6):
│        self.memory = []
│        self.maxsize = int(size)
⋮...
│    def sample(self, length):
⋮...

env/multiagent-particle-envs/multiagent/core.py:
⋮...
│class EntityState(object):
⋮...
│class AgentState(EntityState):
⋮...
│class Action(object):
⋮...
│class Landmark(Entity):
⋮...
│class World(object):
│    def __init__(self):
│        # list of agents and entities (can change at execution-time!)
│        self.agents = []
│        self.landmarks = []
│        # communication channel dimensionality
│        self.dim_c = 0
│        # position dimensionality
│        self.dim_p = 2
│        # color dimensionality
│        self.dim_color = 3
⋮...
│    def apply_action_force(self, p_force):
⋮...
│    def apply_environment_force(self, p_force):
⋮...
│    def integrate_state(self, p_force):
⋮...
│    def update_agent_state(self, agent):
⋮...
│    def get_collision_force(self, entity_a, entity_b):
⋮...

env/multiagent-particle-envs/multiagent/multi_discrete.py:
⋮...
│class MultiDiscrete(gym.Space):
│    """
│    - The multi-discrete action space consists of a series of discrete action spaces with different
│    - It can be adapted to both a Discrete action space or a continuous (Box) action space
│    - It is useful to represent game controllers or keyboards where each key can be represented as
│    - It is parametrized by passing an array of arrays containing [min, max] for each discrete acti
│       where the discrete action space can take any integers from `min` to `max` (both inclusive)
│    Note: A value of 0 always need to represent the NOOP action.
│    e.g. Nintendo Game Controller
│    - Can be conceptualized as 3 discrete action spaces:
│        1) Arrow Keys: Discrete 5  - NOOP[0], UP[1], RIGHT[2], DOWN[3], LEFT[4]  - params: min: 0,
⋮...
│    def sample(self):
⋮...

env/multiagent-particle-envs/multiagent/scenarios/simple_spread.py:
⋮...
│class Scenario(BaseScenario):
│    def make_world(self):
│        world = World()
│        # set any world properties first
│        world.dim_c = 2
│        num_agents = 3
│        num_landmarks = 3
│        world.collaborative = True
│        # add agents
│        world.agents = [Agent() for i in range(num_agents)]
│        for i, agent in enumerate(world.agents):
⋮...
│    def reset_world(self, world):
⋮...

env/observation.py:
⋮...
│class Observation(object):
│
│    def __init__(self, sim, eid, left, right, front, back, res_forward, res_side=0.8, occlusion=Fal
│        """__init__
│
│        Initialize observation channels
│        sim - shared simulation object
│        eid - ID of ego-car that receives this observation
│        left - number of visible sublanes on the left
│        right - number of visible sublanes on the right
│        front - distance (m) visible in front of vehicle
│        back - distance (m) visible behind vehicle
⋮...
│    def occlude_cell(self, r, c):
⋮...


alg/alg_baseline.py, alg/alg_baseline_checkers.py, alg/alg_credit.py, alg/alg_credit_checkers.py, alg/alg_qmix_checkers.py, Q_global, Q_global_1output, Q_global_sumo, Q_global_checkers, Q_credit, list_initialize_credit_ops, check_actions, check_actions, list_initialize_target_ops, test_checkers


reinforcement learning, convolution network, avatai,


Using summaries of '011235813/cm3' files from Context, we can infer that this repository tackles a machine learning problem related to value function estimation for a simulation environment. The repository contains classes for implementing global and local critic models, which are used to evaluate the value of states in the environment.

The data used by this repository is likely to be derived from the simulation environment itself, as the classes for implementing the models require access to state-action pairs and their corresponding rewards or returns. The data may include information about the state of the environment at different points in time, as well as the actions taken by agents interacting with it.

Overall, this repository appears to be focused on developing and testing machine learning models for value function estimation in a simulation environment, with the goal of improving the performance of these models over time through training on large datasets.


class Alg(object):
...
# Global state
...
# Agent observations
...
# Dimensions for image input
...
# Initialize computational graph
...

    def create_networks(self, stage):
...
# Placeholders
...
# Individual agent networks
# output dimension is [time * n_agents, q-values]
...
# To extract Q-value from agent_qs and agent_qs_target; [batch*n_agents, l_action]
...
# [batch, n_agents]
...
# Mixing network
...

    def get_assign_target_ops(self, list_vars):
...
# ops for equating main and target
...
# ops for slow update of target toward main
...
# create op that assigns value of main variable to
# target variable of the same name
...
# incremental update of target towards main
...
# ops for equating main and target
...
# create op that assigns value of main variable to
# target variable of the same name
...
# ops for slow update of target toward main
...
# incremental update of target towards main
...

    def run_actor(self, actions_prev, obs_others, obs_self_t, obs_self_v, goals, epsilon, sess):
...
# convert to batch
...

    def create_train_op(self):
...
# TD target calculated in train_step() using Mixer_target
...

    def create_summary(self):
...

    def process_actions(self, n_steps, actions):
...
# Each row of actions is one time step,
# row contains action indices for all agents
# Convert to [time, agents, l_action]
# so each agent gets its own 1-hot row vector
...
# Convert to format [time*agents, agents-1, l_action]
# so that the set of <n_agent> actions at each time step
# is duplicated <n_agent> times, and each duplicate
# now contains all <n_agent>-1 actions representing
# the OTHER agents actions
...
# extract all actions except agent n's action
...
# interleave
...
# In-place reshape of actions to [time*n_agents, l_action]
...

    def process_batch(self, batch):
...
# shapes are [time, ...original dims...]
# [time, grid]
# [time, agents, l_state_one_agent]
# note that *_local objects have shape
# [time, agents, ...original dim...]
# [time,agents,h,w,c] or [time, agents, obs_others]
# [time,agents,row,column,channel]
# [time,agents,l_obs_self]
...
# [time,agents]
# [time]
# [time,agents]
# [time, grid]
# [time, agents, l_state_one_agent]
# [time,agents,h,w,c]
# [time,agents,row,column,channel]
# [time,agents,l_obs_self]
# [time]
# [time, agents, l_goal]
...
# For all global quantities, for each time step,
# duplicate values <n_agents> times for
# batch processing of all agents
...
# In-place reshape for *_local quantities,
# so that one time step for one agent is considered
# one batch entry
...

    def process_goals(self, goals, n_steps):
...
# Reshape so that one time step for one agent is one batch entry
...
# Reshape to be [n_agents * n_steps, (n_agents-1)*l_goal]
...

    def process_global_state(self, v_global, n_steps):
...
# Reshape into 2D, each block of <n_agents> rows correspond to one time step
...
# Reshape into 2D, each row is state of all other agents, each block of
# <n_agents> rows correspond to one time step
...
# Each agent for each time step is now a batch entry
...
# Get argmax actions from target networks
...
# [batch*n_agents]
# Convert to 1-hot

# Particle
...
# SUMO
...
# Checkers
...
# ----------- Alg parameters ----------------- #
...
# Curriculum stage
...
# If 1, then uses Q-net and global reward
...
# If 1, then restores variables from same stage
...
# If 1, then does not restore variables, even if stage > 1
...
# Name of model to restore
...
# Total number of training episodes
...
# Number of evaluation episodes to run every <period>
...
# Probability of using random configuration
...
# ----------- SUMO parameters ---------------- #
...
# Create entire computational graph
# Creation of new trainable variables for new curriculum
# stage is handled by networks.py, given the stage number
...
# restore only those variables that were not
# just created at this curriculum stage
...
# Copy weights of Q_global to Q_credit at the start of Stage 2
...
# initialize target networks to equal main networks
...
# save everything without exclusion
...
# For computing average over 100 episodes
...
# Write log headers
...
# Each iteration is a training episode
...
# SUMO time functions return negative values afer 24 days (in millisecond) of simulation time
# Hence use 0 for departure time, essentially triggering an immediate departure
...
# Goals for input to policy and value function
...
# Random settings for route, lane and goal
...
# Use predetermined values for route, lane, goal
...
# Random actions when filling replay buffer
...
# Run actor network for all agents as batch
...
# check feasible actions
...
# step environment
...
# store transition into memory
...
# Sample batch of transitions from replay buffer
...
# Write TF summary every <period> episodes, for the first minibatch
...
# Decrease exploration only after training
...
# Clear buffer

class Alg(object):
...
# Global state
...
# Agent observations
...
# Dimensions for image input
...
# Initialize computational graph
...
# TF summaries
...

    def create_networks(self, stage):
...
# Placeholders
...
# Actor network
...
# probs is normalized
...
# V(s,g^n)
...
# Q(s, a^{-n}, g^n, g^{-n}, n, o^n)
...

    def get_assign_target_ops(self, list_vars):
...
# ops for equating main and target
...
# ops for slow update of target toward main
...
# create op that assigns value of main variable to
# target variable of the same name
...
# incremental update of target towards main
...
# For policy
...
# ops for equating main and target
...
# ops for slow update of target toward main
...
# incremental update of target towards main
...
# Repeat for Q if needed
...
# ops for equating main and target
...
# create op that assigns value of main variable to
# target variable of the same name
...
# ops for slow update of target toward main
...
# incremental update of target towards main
...

    def run_actor(self, actions_prev, obs_others, obs_self_t, obs_self_v, goals, epsilon, sess):
...
# convert to batch
...

    def run_actor_target(self, actions_prev, obs_others, obs_self_t, obs_self_v, goals, epsilon, sess):
...

    def create_local_critic_train_op(self):
...
# TD target calculated in train_step() using V_target
...

    def create_global_critic_train_op(self):
...
# TD target calculated in train_step() using Q_target
...
# Get Q-value of action actually taken by point-wise mult
...

    def create_policy_gradient_op(self):
...
# batch of 1-hot action vectors
...
# self.probs has shape [batch, l_action]
...
# --------------- COMA -------------------- #
...
# Q-values of the action actually taken [batch_size]
...
# Use all Q-values at output layer [batch_size]
...
# -------------- End COMA ----------------- #
...

    def create_summary(self):
...

    def process_actions(self, n_steps, actions):
...
# Each row of actions is one time step,
# row contains action indices for all agents
# Convert to [time, agents, l_action]
# so each agent gets its own 1-hot row vector
...
# Convert to format [time*agents, agents-1, l_action]
# so that the set of <n_agent> actions at each time step
# is duplicated <n_agent> times, and each duplicate
# now contains all <n_agent>-1 actions representing
# the OTHER agents actions
...
# extract all actions except agent n's action
...
# interleave
...
# In-place reshape of actions to [time*n_agents, l_action]
...

    def process_batch(self, batch):
...
# shapes are [time, ...original dims...]
# [time, grid]
# [time, agents, l_state_one_agent]
# note that *_local objects have shape
# [time, agents, ...original dim...]
# [time,agents,h,w,c] or [time, agents, obs_others]
# [time,agents,row,column,channel]
# [time,agents,l_obs_self]
...
# [time,agents]
# [time]
# [time,agents]
# [time, grid]
# [time, agents, l_state_one_agent]
# [time,agents,h,w,c]
# [time,agents,row,column,channel]
# [time,agents,l_obs_self]
# [time]
# [time, agents, l_goal]
...
# For all global quantities, for each time step,
# duplicate values <n_agents> times for
# batch processing of all agents
...
# In-place reshape for *_local quantities,
# so that one time step for one agent is considered
# one batch entry
...

    def process_goals(self, goals, n_steps):
...
# Reshape so that one time step for one agent is one batch entry
...
# Reshape to be [n_agents * n_steps, (n_agents-1)*l_goal]
...

    def process_global_state(self, v_global, n_steps):
...
# Reshape into 2D, each block of <n_agents> rows correspond to one time step
...
# Reshape into 2D, each row is state of all other agents, each block of
# <n_agents> rows correspond to one time step
...
# Each agent for each time step is now a batch entry
...
# Create 1-hot agent labels [n_steps*n_agents, n_agents]
...
# ------------ Train local critic ----------------#
...
# V_target(o^n_{t+1}, g^n). V_next_res = V(o^n_{t+1},g^n) used in policy gradient
...
# Get V_target(s',g). V_next_res = V(s',g) is used later in policy gradient
...
# if true, then 0, else 1
...
# Run optimizer for local critic
...
# Get V_res = V(s,g) for use later in policy gradient
...
# ----------- Train global critic -------------- #
...
# Need a_{t+1} to evaluate TD target for Q
# actions is single dimension [n_steps*n_agents]
...
# Now each row is one time step, containing action
# indices for all agents
...
# if true, then 0, else 1
...
# Run optimizer for global critic
...
# --------------- Train policy ------------- #
...
# Already computed V_res when running V_op above

class Alg(object):
...
# global state is all agents' normalized (x,y,speed)
...
# Dimensions for image input
...
# Dimension of agent's observation of itself
...
# global state is all agents' velocity and position (landmarks are fed via goals)
...
# position and velocity
...
# other agents' position and velocity
...
# relative position and velocity
...
# agent's own velocity and position
...
# Initialize computational graph
...
# Use Q_global when n_agents == 1 (the choice is arbitrary,
# since both networks share the same Stage 1 architecture)
...

    def create_networks(self, stage):
...
# Placeholders
...
# Actor network
...
# probs is normalized
...
# Q_n(s,\abf)
...
# Q_n(s,a^m)
...
# V(s,g^n), used as ablation at stage 2
...

    def get_assign_target_ops(self, list_vars):
...
# ops for equating main and target
...
# ops for slow update of target toward main
...
# create op that assigns value of main variable to target variable of the same name
...
# incremental update of target towards main
...
# For policy
...
# op that assigns value of main variable to target variable
...
# incremental update of target towards main
...
# Repeat for Q_credit
...
# create op that assigns value of main variable to target variable of the same name
...
# incremental update of target towards main
...
# create op that assigns value of main variable to target variable of the same name
...
# incremental update of target towards main
...

    def get_assign_global_to_credit_ops(self):
...
# op that assigns value of Q_global variable to Q_credit variable of the same name
...

    def run_actor(self, local_others, local_v, goals, epsilon, sess):
...
# convert to batch
...

    def run_actor_target(self, local_others, local_v, goals, epsilon, sess):
...

    def create_Q_credit_train_op(self):
...
# TD target calculated in train_step() using V_target
...
# Q_credit network has only one output
...

    def create_Q_global_train_op(self):
...
# TD target calculated in train_step() using Q_target
...
# Q_global network has only one output
...

    def create_V_train_op(self):
...

    def create_policy_gradient_op(self):
...
# batch of 1-hot action vectors
...
# self.probs has shape [batch, l_action]
...
# If  stage==2, must be [batch*n_agents, n_agents], consecutive <n_agents> rows are same
...
# First dim is n_agents^2 * batch;
# If n_agents=1, first dim is batch; second dim is l_action
# For <n_agents> > 1, the rows are
# Q_1(s,a^1),...,Q_N(s,a^1),Q_1(s,a^2),...,Q_N(s,a^2), ... , Q_1(s,a^N),...,Q_N(s,a^N)
# where each row contains <l_action> Q-values, one for each possible action
# Note that all Q networks have only one output, and the <l_action> dimension is due to evaluating all possible actions before feeding in feed_dict
...
# First dim is n_agents^2 * batch;
# If n_agents=1, first dim is batch; second dim is l_action
...
# For the general case of any number of agents (covers n_agents==1)
...
# [batch*n_agents, n_agents]
...
# [batch*n_agents, n_agents], each block of nxn is matrix A_{mn} at one time step
...
# [batch, n_agents]
...
# [batch, n_agents]
# [batch, n_agents]
...

    def create_summary(self):
...

    def process_actions(self, n_steps, actions):
...
# Each row of actions is one time step,
# row contains action indices for all agents
# Convert to [time, agents, l_action]
# so each agent gets its own 1-hot row vector
...
# Convert to format [time*agents, agents-1, l_action]
# so that the set of <n_agent> actions at each time step
# is duplicated <n_agent> times, and each duplicate
# now contains all <n_agent>-1 actions representing
# the OTHER agents actions
...
# extract all actions except agent n's action
...
# interleave
...
# In-place reshape of actions to [time*n_agents, l_action]
...

    def process_batch(self, batch):
...
# shapes are [time, ...original dims...]
# [time, agents, l_state_one_agent]
# note that *_local objects have shape
# [time, agents, ...original dim...]
# [time,agents,h,w,c] or [time, agents, obs_others]
# [time,agents,l]
# [time,agents]
# [time]
# [time,agents]
# [time, agents, l_state_one_agent]
# [time,agents,h,w,c]
# [time,agents,l]
# [time]
# [time, agents, l_goal]
...
# For all global quantities, for each time step,
# duplicate values <n_agents> times for
# batch processing of all agents
...
# In-place reshape for *_local quantities,
# so that one time step for one agent is considered
# one batch entry
...

    def process_goals(self, goals, n_steps):
...
# Reshape so that one time step for one agent is one batch entry
...
# Reshape to be [n_agents * n_steps, (n_agents-1)*l_goal]
...

    def process_global_state(self, v_global, n_steps):
...
# Reshape into 2D, each block of <n_agents> rows correspond to one time step
...
# Reshape into 2D, each row is state of all other agents, each block of
# <n_agents> rows correspond to one time step
...
# Each agent for each time step is now a batch entry
...
# Create 1-hot agent labels [n_steps*n_agents*n_agents, n_agents]
# because Q_n(s,a^m) requires both n and m = 1,..., n_agents
# agent_labels = np.tile(self.agent_labels, (n_steps*n_agents,1))
...
# ------------ Train Q^{\pibf}_n(s,\abf) --------------- #
# ------------------------------------------------------ #
...
# Need a_{t+1} for all agents to evaluate TD target for Q
# actions is single dimension [n_steps*n_agents]
...
# Now each row is one time step, containing action
# indices for all agents
...
# Get target Q_n(s',\abf')
...
# if true, then 0, else 1
...
# Run optimizer
...
# Get Q_global_res = Q_n(s,\abf) for use later in policy gradient
...
# For feeding into Q_actual
...
# ----------- Train Q^{\pibf}_n(s,a^m) -------------- #
# --------------------------------------------------- #
...
# When going down rows, n is the inner index and m is the outer index
...
# Repeat the things indexed by n, so that the group of n_agents rows
# at each time step is repeated n_agent times
...
# Duplicate the things indexed by m so that consecutive n_agents rows are the same (so that when summing over n, they are the same)
...
# Get target Q_n(s',a'^m)
...
# reward_local is n_steps*n_agents
# Duplicate into n_steps*n_agents*n_agents so that each time step is
# [r^1,r^2,r^3,r^4,...,r^1,r^2,r^3,r^4] (assume N=4)
# Now n_steps*n_agents*n_agents
...
# if true, then 0, else 1
...
# Duplicate the things indexed by n
...
# Duplicate the things indexed by m
...
# Run optimizer for global critic
...
# Get target V_n(s')
...
# if true, then 0, else 1
...
# Run optimizer
...
# Get V_res = V_n(s) for use later in policy gradient
...
# For feeding into advantage function
...
# --------------- Train policy ------------- #
# ------------------------------------------ #
...
# Stage 1
...
# compute Q(s,a=all possible,g)
...
# not used
...
# not used
...
# Compute values for probs_evaluated
...
# Duplicate everything by number of possible actions, to prepare for batch
# computation of counterfactuals
...
# Counterfactual actions
...
# Compute Q_n(s,a^m) for all n and m and all actions a^m=i, for all steps
...
# n_steps * n_agents^2 * l_action

class Checkers(object):
...
# Only n_rows and n_columns have green and orange squares
...
# Total grid size is larger so that agents' observations are valid
# when they are located on the boundary
...
# Used to determine episode termination
...
# Initial agent locations, situated in expanded grid
...

    def populate_world(self):
...
# Fill in invalid cells
...
# Agent locations are also invalid, i.e. cannot move there
...
# Fill in rewards
...

    def get_valid_grid(self):
...
# the +1 is to include agents' starting column
...
# No need to include the invalid channel, which is only used
# for agents' observations
...

    def get_global_state(self):
...
# each agent's coordinates and number of collected orange and green
...

    def get_obs(self, agent_location):
...
# edit the "invalid" channel so that agent's current location is valid
...

    def normalize(self, location):
...

    def get_local_observation(self):
...
# [n_obs+1, n_obs+1] grid centered on agent location
# observation grid
# Compute normalized agent coordinates
...
# Concatenate with normalized number of collected things
...
# list_obs_self.append(obs_self)
...
# coordinates of all other agents
# If only one agent, use its own location (needed for placeholder
# to work, but it is ignored)
...
# Stopped here
...

    def agent_act(self, idx, action):
...
# Current location
...
# up
...
# down
...
# left
...
# right
...
# Penalty for trying to move to invalid location
...

    def get_reward(self, idx, goal_idx):
...
# agent_act() was called, so this is the agent's new location
...
# green
...
# orange
...
# green
...
# orange
...

    def step(self, actions):
...
# local_rewards.append(penalty+self.get_reward(idx))
...

    def reset(self, goals):
...
# If single agent, initial row is random 0,1,2
...
# each row is one agent's (r,c) coordinates

class Alg(object):
...
# global state is all agents' normalized (x,y,speed)
...
# Dimensions for image input
...
# Dimension of agent's observation of itself
...
# global state is all agents' velocity and position (landmarks are fed via goals)
...
# position and velocity
...
# other agents' position and velocity
...
# relative position and velocity
...
# agent's own velocity and position
...
# Initialize computational graph
...
# TF summaries
...

    def create_networks(self, stage):
...
# Placeholders
...
# Actor network
...
# probs is normalized
...
# V(s,g^n)
...
# Q(s, a^{-n}, g^n, g^{-n}, n, o^n)
...

    def get_assign_target_ops(self, list_vars):
...
# ops for equating main and target
...
# ops for slow update of target toward main
...
# create op that assigns value of main variable to
# target variable of the same name
...
# incremental update of target towards main
...
# For policy
...
# ops for equating main and target
...
# ops for slow update of target toward main
...
# incremental update of target towards main
...
# Repeat for Q if needed
...
# ops for equating main and target
...
# create op that assigns value of main variable to
# target variable of the same name
...
# ops for slow update of target toward main
...
# incremental update of target towards main
...

    def run_actor(self, local_others, local_v, goals, epsilon, sess):
...
# convert to batch
...

    def run_actor_target(self, local_others, local_v, goals, epsilon, sess):
...

    def create_local_critic_train_op(self):
...
# TD target calculated in train_step() using V_target
...

    def create_global_critic_train_op(self):
...
# TD target calculated in train_step() using Q_target
...
# Get Q-value of action actually taken by point-wise mult
...

    def create_policy_gradient_op(self):
...
# batch of 1-hot action vectors
...
# self.probs has shape [batch, l_action]
...
# --------------- COMA -------------------- #
...
# Q-values of the action actually taken [batch_size]
...
# Use all Q-values at output layer [batch_size]
...
# -------------- End COMA ----------------- #
...

    def create_summary(self):
...

    def process_actions(self, n_steps, actions):
...
# Each row of actions is one time step,
# row contains action indices for all agents
# Convert to [time, agents, l_action]
# so each agent gets its own 1-hot row vector
...
# Convert to format [time*agents, agents-1, l_action]
# so that the set of <n_agent> actions at each time step
# is duplicated <n_agent> times, and each duplicate
# now contains all <n_agent>-1 actions representing
# the OTHER agents actions
...
# extract all actions except agent n's action
...
# interleave
...
# In-place reshape of actions to [time*n_agents, l_action]
...

    def process_batch(self, batch):
...
# shapes are [time, ...original dims...]
# [time, agents, l_state_one_agent]
# note that *_local objects have shape
# [time, agents, ...original dim...]
# [time,agents,h,w,c] or [time, agents, obs_others]
# [time,agents,l]
# [time,agents]
# [time]
# [time,agents]
# [time, agents, l_state_one_agent]
# [time,agents,h,w,c]
# [time,agents,l]
# [time]
# [time, agents, l_goal]
...
# For all global quantities, for each time step,
# duplicate values <n_agents> times for
# batch processing of all agents
...
# In-place reshape for *_local quantities,
# so that one time step for one agent is considered
# one batch entry
...

    def process_goals(self, goals, n_steps):
...
# Reshape so that one time step for one agent is one batch entry
...
# Reshape to be [n_agents * n_steps, (n_agents-1)*l_goal]
...

    def process_global_state(self, v_global, n_steps):
...
# Reshape into 2D, each block of <n_agents> rows correspond to one time step
...
# Reshape into 2D, each row is state of all other agents, each block of
# <n_agents> rows correspond to one time step
...
# Each agent for each time step is now a batch entry
...
# Create 1-hot agent labels [n_steps*n_agents, n_agents]
...
# ------------ Train local critic ----------------#
...
# V_target(o^n_{t+1}, g^n). V_next_res = V(o^n_{t+1},g^n) used in policy gradient
...
# Get V_target(s',g). V_next_res = V(s',g) is used later in policy gradient
...
# if true, then 0, else 1
...
# Run optimizer for local critic
...
# Get V_res = V(s,g) for use later in policy gradient
...
# ----------- Train global critic -------------- #
...
# Need a_{t+1} to evaluate TD target for Q
# actions is single dimension [n_steps*n_agents]
...
# Now each row is one time step, containing action
# indices for all agents
...
# if true, then 0, else 1
...
# Run optimizer for global critic
...
# --------------- Train policy ------------- #
...
# Already computed V_res when running V_op above

class EgoCar(car.Car):
...
# absolute sublane number of center of goal lane
...
# Car encounters merge situation if it started from
# lane 0 of straight road and goal lane is 0 on edge 3,
# or if it started from entrance road
...
# All speed checks off
...
# No checks at all
...
# List of IDs of edges along vehicle's route
...
# One-hot vector, see update_current_edge_type()
...
# Vehicle speed limits
...
# 9
# Actions indices. LEFT and RIGHT cause sublane change,
# not a compete lane change
...
# if vtype == 'ego_car_type':
...
# acceleration
# deceleration
...
# Local observation object
...
# Maximum number of steps allowed before termination
# Equivalent to requiring minimum average speed of 20m/s
# over entire route
...
# minimum seconds between lane change operations
# round(1/self.dt)
...
# False
...
# time-to-collision tolerance threshold
# original 10.0
...
# True
...
# Distance from next road junctions
# Will be part of the observation vector
...
# Distance to junctions along route
# Consider goal position of the final edge as a junction
...
# Distance to goal_pos and goal_lane on final edge
...

  def __str__(self):

  def update_current_edge_type(self):
...

  def update_state(self, map_car_values, action):
...
# This means the car has reached end of
# final edge and was removed by simulator.
...
# -10
...
# update basic state
...
# Index of edge along vehicle's route
...
# Number of lanes of current edge
...
# absolute distance from junctions
...
# last element is distance to goal position
...
# normalized distance to next junction
...
# normalized distance to goal
...
# difference from goal lane
...
# allow lane change after time max_lc_step has passed
...
# local observation update
...
# self.update_current_edge_type()
...
# Assign instantaneous reward
...
# -10.0
...
# -10.0
...

  def execute_action(self, action):
...
# parse action
...
# self.NOOP
...
# finite difference to change velocity
...
# bound velocity
...
# apply velocity
...

  def get_feasible_actions(self, map_car_values):
...
# min, max velocity
...
# lane changing
...
# f_action[self.LEFT_L_ON] = np.nan
...
# f_action[self.RIGHT_R_ON] = np.nan
...

  def getBaselineAction(self):
...

  def getRandomAction(self):
...

  def wait_until_ready(self):

class Replay_Buffer():
...

    def __init__(self, size=5e4):
...

    def add(self, episode, is_bad=False):
...

    def add_1(self, transition):
...

    def add_2(self, transition):
...

    def sample_batch(self, size):
...
# Enough bad transitions but not enough good ones
...
# Not enough bad transitions but enough good ones

class Alg(object):
...
# Global state
...
# Agent observations
...
# Dimensions for image input
...
# Dimension of agent's observation of itself
...
# Initialize computational graph
...
# Use Q_global when n_agents == 1 (the choice is arbitrary,
# since both networks share the same Stage 1 architecture)
...
# TF summaries
...

    def create_networks(self, stage):
...
# Placeholders
...
# Actor network
...
# probs is normalized
...
# Q_n(s,\abf)
...
# Q_n(s,a^m)
...
# V(s,g^n), used as ablation at stage 2
...

    def get_assign_target_ops(self, list_vars):
...
# ops for equating main and target
...
# ops for slow update of target toward main
...
# create op that assigns value of main variable to target variable of the same name
...
# incremental update of target towards main
...
# For policy
...
# op that assigns value of main variable to target variable
...
# incremental update of target towards main
...
# Repeat for Q_credit
...
# create op that assigns value of main variable to target variable of the same name
...
# incremental update of target towards main
...
# create op that assigns value of main variable to target variable of the same name
...
# incremental update of target towards main
...

    def get_assign_global_to_credit_ops(self):
...
# op that assigns value of Q_global variable to Q_credit variable of the same name
...

    def run_actor(self, actions_prev, obs_others, obs_self_t, obs_self_v, goals, epsilon, sess):
...
# convert to batch
...

    def run_actor_target(self, actions_prev, obs_others, obs_self_t, obs_self_v, goals, epsilon, sess):
...

    def create_Q_credit_train_op(self):
...
# TD target calculated in train_step() using V_target
...
# Q_credit network has only one output
...

    def create_Q_global_train_op(self):
...
# TD target calculated in train_step() using Q_target
...
# Q_global network has only one output
...

    def create_V_train_op(self):
...

    def create_policy_gradient_op(self):
...
# batch of 1-hot action vectors
...
# self.probs has shape [batch, l_action]
...
# if stage==2, must be [batch*n_agents, n_agents], consecutive <n_agents> rows are same
...
# First dim is n_agents^2 * batch;
# If n_agents=1, first dim is batch; second dim is l_action
# For <n_agents> > 1, the rows are Q_1(s,a^1),...,Q_N(s,a^1),Q_1(s,a^2),...,Q_N(s,a^2), ... , Q_1(s,a^N),...,Q_N(s,a^N)
# where each row contains <l_action> Q-values, one for each possible action
# Note that all Q networks have only one output, and the <l_action> dimension is due to evaluating all possible actions before feeding in feed_dict
...
# First dim is n_agents^2 * batch;
# If n_agents=1, first dim is batch; second dim is l_action
...
# For the general case of any number of agents (covers n_agents==1)
...
# [batch*n_agents, n_agents]
...
# [batch*n_agents, n_agents], each block of nxn is matrix A_{mn} at one time step
...
# [batch, n_agents]
...
# [batch, n_agents]
# [batch, n_agents]
...

    def create_summary(self):
...

    def process_actions(self, n_steps, actions):
...
# Each row of actions is one time step,
# row contains action indices for all agents
# Convert to [time, agents, l_action]
# so each agent gets its own 1-hot row vector
...
# Convert to format [time*agents, agents-1, l_action]
# so that the set of <n_agent> actions at each time step
# is duplicated <n_agent> times, and each duplicate
# now contains all <n_agent>-1 actions representing
# the OTHER agents actions
...
# extract all actions except agent n's action
...
# interleave
...
# In-place reshape of actions to [time*n_agents, l_action]
...

    def process_batch(self, batch):
...
# shapes are [time, ...original dims...]
# [time, grid]
# [time, agents, l_state_one_agent]
# note that *_local objects have shape
# [time, agents, ...original dim...]
# [time,agents,h,w,c] or [time, agents, obs_others]
# [time,agents,row,column,channel]
# [time,agents,l_obs_self]
...
# [time,agents]
# [time]
# [time,agents]
# [time, grid]
# [time, agents, l_state_one_agent]
# [time,agents,h,w,c]
# [time,agents,row,column,channel]
# [time,agents,l_obs_self]
# [time]
# [time, agents, l_goal]
...
# Try to free memory
...
# For all global quantities, for each time step,
# duplicate values <n_agents> times for
# batch processing of all agents
...
# In-place reshape for *_local quantities,
# so that one time step for one agent is considered
# one batch entry
...

    def process_goals(self, goals, n_steps):
...
# Reshape so that one time step for one agent is one batch entry
...
# Reshape to be [n_agents * n_steps, (n_agents-1)*l_goal]
...

    def process_global_state(self, v_global, n_steps):
...
# Reshape into 2D, each block of <n_agents> rows correspond to one time step
...
# Reshape into 2D, each row is state of all other agents, each block of
# <n_agents> rows correspond to one time step
...
# Each agent for each time step is now a batch entry
...
# ------------ Train Q^{\pibf}_n(s,\abf) --------------- #
# ------------------------------------------------------ #
...
# Need a_{t+1} for all agents to evaluate TD target for Q
# actions is single dimension [n_steps*n_agents]
...
# Now each row is one time step, containing action
# indices for all agents
...
# Get target Q_n(s',\abf')
...
# if true, then 0, else 1
...
# Run optimizer
...
# Get Q_global_res = Q_n(s,\abf) for use later in policy gradient
...
# For feeding into Q_actual
...
# ----------- Train Q^{\pibf}_n(s,a^m) -------------- #
# --------------------------------------------------- #
...
# When going down rows, n is the inner index and m is the outer index
...
# Repeat the things indexed by n, so that the group of n_agents rows
# at each time step is repeated n_agent times
...
# Duplicate the things indexed by m so that consecutive n_agents rows are the same (so that when summing over n, they are the same)
...
# Duplicate shared state
...
# Get target Q_n(s',a'^m)
...
# reward_local is n_steps*n_agents
# Duplicate into n_steps*n_agents*n_agents so that each time step is
# [r^1,r^2,r^3,r^4,...,r^1,r^2,r^3,r^4] (assume N=4)
# Now n_steps*n_agents*n_agents
...
# if true, then 0, else 1
...
# Duplicate the things indexed by n
...
# Duplicate the things indexed by m
...
# Duplicate shared state
...
# Run optimizer for global critic
...
# Get target V_n(s')
...
# if true, then 0, else 1
...
# Run optimizer
...
# Get V_res = V_n(s) for use later in policy gradient
...
# For feeding into advantage function
...
# --------------- Train policy ------------- #
# ------------------------------------------ #
...
# Stage 1
...
# compute Q(s,a=all possible,g)
...
# not used
...
# Compute values for probs_evaluated
...
# Duplicate everything by number of possible actions, to prepare for batch
# computation of counterfactuals
...
# Counterfactual actions
...
# Compute Q_n(s,a^m) for all n and m and all actions a^m=i, for all steps
...
# n_steps * n_agents^2 * l_action

class Car(object):
...

  def __init__(self, sim, vel, lane, pos=0, cid='car'):
...

  def __str__(self):

  def updateState(self):


# Cooperative Multi-Stage Multi-Goal Multi-Agent Reinforcement Learning (CM3)

This repository provides code for experiments in the paper [CM3](https://arxiv.org/abs/1809.05188)[1], published in ICLR 2020.
It contains the main algorithm and baselines, and the three simulation Markov games on which algorithms were evaluated.

## Dependencies

- All experiments were run on Ubuntu 16.04
- Python 3.6
- TensorFlow 1.10
- [SUMO](https://github.com/eclipse/sumo)
- pygame: `sudo apt-get install python-pygame`
- OpenAI Gym 0.12.1

## Project structure

- `alg`: Implementation of algorithms and config files. `config.json` is the main config file. `config_particle_*.json` specifies various instances of the cooperative navigation task. `config_sumo_stage{1,2}.json` specifies agent initial/goal lane configurations for SUMO. `config_checkers_stage{1,2}.json` specifies parameters of the Checkers game.
- `env`: Python wrappers/definitions of the simulation environments.
- `env_sumo`: XML files that define the road and traffic for the underlying SUMO simulator.
- `log`: Each experiment run will create a subfolder that contains the reward values logged during the training or test run.
- `saved`: Each experiment run will create a subfolder contains trained TensorFlow models.

## Environments

There are three simulations, selected by the `experiment` field in `alg/config.json`.

1. Cooperative navigation: particles must move to individual target locations while avoiding collisions.
   - Environment code located in `env/multiagent-particle-envs/`
2. SUMO
   - Stage 1: single agent on empty road. Corresponds to setting `"stage" : 1`
   - Stage 2: two agents on empty road. Corresponds to setting `"stage" : 2`
   - Python wrappers located in `env/`. Entry point is `env/multicar_simple.py`
   - SUMO topology and traffic defined in `env_sumo/simple/`
3. Checkers: two agents cooperate to collect rewards while avoiding penalties in a checkered map.
   - Implemented in `env/checkers.py`

## Environment setup

1. Cooperative navigation: run `pip install -e .` inside `env/multiagent-particle-envs/`
2. SUMO: Install [SUMO](https://github.com/eclipse/sumo) and add the following to your `.bashrc`
 - `export PYTHONPATH=$PYTHONPATH:path/to/sumo`
 - `export PYTHONPATH=$PYTHONPATH:path/to/sumo/tools`
 - `export SUMO_HOME="path/to/sumo"`
3. Checkers: None required

## Training

### Environment-specific examples

**Cooperative navigation**

- In `config.json`, set
  - `experiment`: "particle"
  - `particle_config` should be one of `config_particle_stage1.json`, `config_particle_stage2_antipodal.json`, `config_particle_stage2_cross.json`, `config_particle_stage2_merge.json`
- Inside `alg/`, execute
  `python train_onpolicy.py`

**SUMO**

- In `config.json`, set
  - `experiment`: "sumo"
  - `port`: if multiple SUMO experiments are run in parallel, each experiment must have its unique number
- Inside `alg/`, execute
  `python train_offpolicy.py --env ../env_sumo/simple/merge.sumocfg`
- Include the option `--gui` to show SUMO GUI while training (at the cost of increased runtime)

**Checkers**

- In `config.json`, set
  - `experiment` : "checkers"
- Inside `alg/`, execute
  `python train_offpolicy.py`

### General notes for running Stage 1 and 2 of CM3
- `stage`: either 1 or 2
- `dir_restore`: for Stage 2 of CM3, this must be equal to the string for `dir_name` when Stage 1 was run.
- `use_alg_credit`: 1 for CM3
- `use_Q_credit`: 1 for CM3. 0 for ablation that uses value function baseline.
- `train_from_nothing`: 1 for Stage 1 of CM3, or the ablation that omits the curriculum. 0 to allow restoring a trained Stage 1 model.
- `model_name`: when training Stage 2 and restoring a Stage 1 model, this must be the name of the model in Stage 1.
- `prob_random`: 1.0 for Stage 1, 0.2 for Stage 2. Not applicable for Checkers.

## Citation

<pre>
@inproceedings{yang2019cm3,
  title={CM3: Cooperative Multi-goal Multi-stage Multi-agent Reinforcement Learning},
  author={Yang, Jiachen and Nakhaei, Alireza and Isele, David and Fujimura, Kikuo and Zha, Hongyuan},
  booktitle={International Conference on Learning Representations},
  year={2019}
}
</pre>

## License

See [LICENSE](LICENSE).

SPDX-License-Identifier: MIT


This repository tackles a machine learning problem related to value function estimation for a simulation environment. The repository contains classes for implementing global and local critic models, which are used to evaluate the value of states in the environment. The data used by this repository is likely to be derived from the simulation environment itself, as the classes for implementing the models require access to state-action pairs and their corresponding rewards or returns.


* `env/car.py`: This file contains a class called `Car` that represents a car object in the simulation environment. It has methods for updating its state, such as its velocity and position, and for printing information about it.
* `model/q_global.py`: This file contains a class called `QGlobal` that implements a global critic model used to evaluate the value of states in the simulation environment. It has methods for training the model on data from the simulation and for using it to compute the expected return for a given state-action pair.
* `model/q_local.py`: This file contains a class called `QLocal` that implements a local critic model used to evaluate the value of states in the simulation environment. It has methods for training the model on data from the simulation and for using it to compute the expected return for a given state-action pair.
* `model/v_global.py`: This file contains a class called `VGlobal` that implements a global value function model used to evaluate the value of states in the simulation environment. It has methods for training the model on data from the simulation and for using it to compute the expected return for a given state-action pair.


# Hierarchial Cooperative Multi-Agent Reinforcement Learning with Skill Discovery (HSD)

This is the code for experiments in the paper [Hierarchial Cooperative Multi-Agent Reinforcement Learning with Skill Discovery](https://arxiv.org/abs/1912.03558), published in AAMAS 2020. Ablations and baselines are included.

## Prerequisites

- Python version >= 3.5.2
- TensorFlow 1.13.1
- PyGame 1.9.4
- [STS2 1.0.0](https://github.com/electronicarts/SimpleTeamSportsSimulator). In case of future API changes, our algorithm is compatible with at least this [submit](https://github.com/electronicarts/SimpleTeamSportsSimulator/tree/80b4c688a4b8fb3113f466fae5de060e29c79fbe).

## Project structure

- `alg` : implementation of algorithm, neural networks, `config.json` containing all hyperparameters.
- `env` : implementation of multi-agent wrapper around STS2 simulator.
- `results` : each experiment will create a subfolder that contains log files recorded during training and eval.
- `test` : test scripts

## Training

Each algorithm named `alg_*.py` is run through a script with name `train_*.py`.
The pairings are as follows:
- `train_hsd.py` runs `alg_hsd.py` (HSD)
- `train_offpolicy.py` runs `alg_qmix.py` (QMIX) and `alg_iql.py` (IQL)
- `train_hsd_scripted.py` runs `alg_hsd_scripted.py`

To do multi-seed runs that sweep over the initial random seed, set appropriate choices in config.json and use `train_multiprocess.py`. See example below.

For all algorithms,
- Activate your TensorFlow (if using `virtualenv`) and allocate GPU using `export CUDA_VISIBLE_DEVICES=<n>` where `n` is some GPU number.
- `cd` into the `alg` folder
- Execute training script, e.g. `python train_hsd.py`
- Periodic training progress is logged in `log.csv`, along with saved models, under `results/<dir_name>`.

### Example 1: training HSD

- Select correct settings in `alg/config.json`. Refer to `config_hsd.json` for an example. The key parameters to set are
  - `"alg_name" : "hsd"`
  - everything under `"h_params"`
  - neural network parameters under `"nn_hsd"`

### Example 2: training QMIX

- Select correct settings in `alg/config.json`. Refer to `config_qmix.json` for an example. The key parameters to set are
  - `"alg_name" : "qmix"`
  - neural network parameters under `"nn_qmix"`

### Example 3 for multi-seed runs

For example, to conduct 5 parallel runs with seeds 12341,12342,...,12345 and save into directory names hsd_1, hsd_2,...,hsd_3 (all under `results/`), set the following parameters in config.json:
- `"N_seeds" : 5`
- `"seed" : 12341`
- `"dir_name" : "hsd"`
- `"dir_idx_start" : 1`

## Testing

### Example 1 for testing HSD

- Choose appropriate settings in `alg/config.json`.

  - `"dir_name" : "hsd_1"`
  - `"model_name" : "model_good.ckpt-<some number>"`
  - `"render" : true` (to see PyGame)
  - `"N_test" : 100` (for 100 test episodes)
  - `"measure" : true` (to enable generation of additional .csv files for analysis of behavior)

- `cd` into the `alg` folder. Execute test script `python test.py`

- Results will be stored in `test.csv` under `results/<dir_name>/`. If `"measure" : true`, then files `matrix_role_counts.pkl`, `count_skills.pkl` and `count_low_actions.pkl` will also be generated.

## Citation

<pre>
@inproceedings{yang2020hierarchical,
  title={Hierarchical Cooperative Multi-Agent Reinforcement Learning with Skill Discovery},
  author={Yang, Jiachen and Borovikov, Igor and Zha, Hongyuan},
  booktitle={Proceedings of the 19th International Conference on Autonomous Agents and MultiAgent Systems},
  pages={1566--1574},
  year={2020}
}
</pre>

## License

HSD is distributed under the terms of the BSD-3 license. All new contributions must be made under this license.

See [LICENSE](LICENSE) for details.

SPDX-License-Identifier: BSD-3

q learning
multi agent reinforcement learning
reinforcement learning


This repository tackles the problem of training a hierarchical multi-agent reinforcement learning (HSD) algorithm, which involves training both a high-level policy that selects subgoals and a low-level policy that takes actions based on those subgoals. The file also includes code for computing intrinsic rewards and updating the high-level state.

The repository uses a dataset of observations collected during training episodes to train both the high-level policy and low-level policies using stochastic gradient descent (SGD).


# print(s)
...
# number of times each of 8 event occurs under each skill
# 8 hardcoded events
# number of times each skill was selected, recorded separately for cases with and without possession
...
# number of primitive actions under each skill
...
# one sequence of skill integer for each agent
...
# count of skill usage over (x,y) position, one 2D matrix for each skill
...
# if not None, then it's a length-2 list of [agent idx, skill idx]

# Each <steps_per_assign> is one "step" for the high-level policy
# This means we train the high-level policy once for every
# <steps_per_train> high-level steps
...
# Buffer for high level role assignment policy
...
# Buffer for low level agent policy
...
# Logging
...
# Variables with suffix _h are high-level quantities for training the role assignment policy
# These are the high-level equivalent of the s_t in a usual transition tuple (s_t, a_t, s_{t+1})
...
# Cumulative discounted reward for high-level policy
...
# Action taken by high-level role assignment policy
...
# steps within an episode
...
# The environment state at this point, e.g. <state_home>,
# acts like the "next state" for the high-level policy
# All of the intervening environment steps act as a single step for the high-level policy
...
# Get new role assignment, i.e. take high-level action
...
# Conduct training of high-level policy
...
# Update high-level state
...
# Take low-level actions, conditioned on roles
...
# Train low-level policies
...
# Since the episode is done, we also terminate the current role assignment period,
# even if not all <steps_per_assign> have been completed

class Alg(object):
...

    def __init__(self, config_alg, n_agents, l_state, l_obs, l_action, nn):
...
# Initialize computational graph
...
# TF summaries
...

    def create_networks(self):
...
# Placeholders
...
# Individual agent networks
# output dimension is [time * n_agents, q-values]
...
# To extract Q-value from agent_qs and agent_qs_target
# [batch*n_agents, l_action]
...
# [batch*n_agents, 1]
...

    def get_assign_target_ops(self, list_vars):
...
# ops for equating main and target
...
# ops for slow update of target toward main
...
# create op that assigns value of main variable to
# target variable of the same name
...
# incremental update of target towards main
...

    def run_actor(self, list_obs, epsilon, sess):
...
# convert to batch
...

    def create_train_op(self):
...
# TD target calculated in train_step() using Mixer_target
...

    def create_summary(self):
...

    def process_actions(self, n_steps, actions):
...
# Each row of actions is one time step,
# row contains action indices for all agents
# Convert to [time, agents, l_action]
# so each agent gets its own 1-hot row vector
...
# In-place reshape of actions to [time*n_agents, l_action]
...

    def process_batch(self, batch):
...
# shapes are [time, ...original dims...]
# [time, l_state]
# [time, agents, l_obs]
# [time, agents]
# [time]
# [time, l_state]
# [time, agents, l_obs]
# [time]
...
# Try to free memory
...
# Duplicate values <n_agents> times for batch processing of all agents
...
# In-place reshape for obs, so that one time step
# for one agent is considered one batch entry
...

    def train_step(self, sess, batch, step_train=0, summarize=False, writer=None):
...
# Each agent for each time step is now a batch entry

class Alg(object):
...

    def __init__(self, alg_name, config_alg, n_agents, l_state, l_obs, l_action, N_roles, nn):
...
# Combinatorial action space, but we don't allow duplicate roles
...
# Initialize computational graph
...
# TF summaries
...

    def populate_list_list_roles(self, agent_idx, list_indices):
...
# Skip over duplicate roles
...

    def create_networks(self):
...
# Placeholders
...
# Low-level Q-functions
...
# Low level action
...
# High-level Q-functions
...
# Individual agent networks
# output dimension is [time * n_agents, q-values]
...
# To extract Q-value from agent_qs and agent_qs_target
# [batch*n_agents, N_roles]
...
# [batch*n_agents, 1]
...
# [batch, n_agents]
...
# Mixing network
...
# Standard Q-learning for role assignment
...

    def get_assign_target_ops(self):
...
# ops for equating main and target
...
# ops for slow update of target toward main
...
# ops for slow update of low-level target toward low-level main
...
# create op that assigns value of main variable to
# target variable of the same name
...
# incremental update of target towards main
...
# ops for equating main and target
...
# create op that assigns value of main variable to
# target variable of the same name
...
# ops for slow update of target toward main
...
# incremental update of target towards main
...
# create op that assigns value of main variable to
# target variable of the same name
...
# incremental update of target towards main
...
# create op that assigns value of main variable to
# target variable of the same name
...
# incremental update of target towards main
...

    def run_actor(self, list_obs, roles, epsilon, sess):
...
# convert to batch
...

    def assign_roles(self, list_obs, epsilon, sess):
...

    def assign_roles_centralized(self, state, epsilon, sess):
...

    def create_train_op(self):
...
# TD target calculated in train_step() using Mixer_target
...
# Treat the role assignment as single-agent Q-learning
...

    def create_train_op_IQL(self):
...

    def create_summary(self):
...

    def process_actions(self, n_steps, actions, n_actions):
...
# Each row of actions is one time step,
# row contains action indices for all agents
# Convert to [time, agents, N_roles]
# so each agent gets its own 1-hot row vector
...
# In-place reshape of actions to [time*n_agents, N_roles]
...

    def process_batch(self, batch):
...
# shapes are [time, ...original dims...]
...
# [time, l_state]
# [time, agents, l_obs]
# [time, agents]
# [time]
# [time, l_state]
# [time, agents, l_obs]
# [time]
...
# [time, l_state]
# [time]
# [time]
# [time, l_state]
# [time]
...
# Try to free memory
...
# In-place reshape for obs, so that one time step
# for one agent is considered one batch entry
...

    def process_batch_low(self, batch):
...
# shapes are [time, ...original dims...]
# state = np.stack(batch[:,0]) # [time, l_state]
# [time, agents, l_obs]
# [time, agents]
# [time, agents]
# state_next = np.stack(batch[:,4]) # [time, l_state]
# [time, agents, l_obs]
# done = np.stack(batch[:,6]) # [time]
# [time, agents, N_roles]
...
# Try to free memory
...
# In-place reshape for obs, so that one time step
# for one agent is considered one batch entry
...

    def train_step(self, sess, batch, step_train=0, summarize=False, writer=None):
...
# Each agent for each time step is now a batch entry
...
# Get argmax actions from target networks
...
# [batch*n_agents]
# Convert to 1-hot
...
# Get Q_tot target value
...

    def train_step_low(self, sess, batch, step_train=0, summarize=False, writer=None):

class Replay_Buffer():
...

    def __init__(self, size=1e6):
...

    def add(self, transition):
...

    def sample(self, length):

class Alg(object):
...

    def __init__(self, config_alg, config_h, n_agents, l_state, l_obs, l_action, l_z, nn):
...
# Domain-specific removal of information from agent observation
# Either none (deactivate) or scalar index where obs should be truncated for use by decoder
...
# Initialize computational graph
...
# TF summaries
...

    def create_networks(self):
...
# Placeholders
...
# Decoder p(z|tau)
...
# Low-level policy
...
# Low-level Q-functions
...
# High-level QMIX
# Individual agent networks
# output dimension is [time * n_agents, q-values]
...
# To extract Q-value from agent_qs and agent_qs_target
# [batch*n_agents, N_roles]
...
# [batch*n_agents, 1]
...
# [batch, n_agents]
...
# Mixing network
...

    def get_assign_target_ops(self):
...
# ops for equating main and target
...
# ops for slow update of target toward main
...
# ops for slow update of low-level target toward low-level main
...
# create op that assigns value of main variable to
# target variable of the same name
...
# incremental update of target towards main
...
# ops for equating main and target
...
# create op that assigns value of main variable to
# target variable of the same name
...
# ops for slow update of target toward main
...
# incremental update of target towards main
...

    def run_actor(self, list_obs, roles, epsilon, sess):
...
# convert to batch
...

    def assign_roles(self, list_obs, epsilon, sess, N_roles_current):
...
# Limit the number of activated options based on curriculum
...

    def create_train_op_high(self):
...

    def create_train_op_low(self):
...
# self.probs shape is [batch size * traj length, l_action]
# now log_probs shape is [batch size * traj length]
...
# Rehape to [batch size, traj length]
...
# E [ \sum_t \log \pi(a_t|o_t,z) * R ]
...
# Critic train op
...
# Policy train op
...

    def create_train_op_decoder(self):
...

    def create_summary(self):
...

    def process_actions(self, n_steps, actions, n_actions):
...
# Each row of actions is one time step,
# row contains action indices for all agents
# Convert to [time, agents, N_roles]
# so each agent gets its own 1-hot row vector
...
# In-place reshape of actions to [time*n_agents, N_roles]
...

    def process_batch(self, batch):
...
# shapes are [time, ...original dims...]
# [time, l_state]
# [time, agents, l_obs]
# [time, agents]
# [time]
# [time, l_state]
# [time, agents, l_obs]
# [time]
...
# Try to free memory
...
# In-place reshape for obs, so that one time step
# for one agent is considered one batch entry
...

    def train_policy_high(self, sess, batch, step_train, summarize=False, writer=None):
...
# Each agent for each time step is now a batch entry
...
# Get argmax actions from target networks
...
# [batch*n_agents]
# Convert to 1-hot
...
# Get Q_tot target value
...

    def process_batch_low(self, batch):
...
# shapes are [time, ...original dims...]
# [time, agents, l_obs]
# [time, agents]
# [time, agents]
# [time, agents, l_obs]
# [time, agents, N_roles]
# [time]
...
# Try to free memory
...
# In-place reshape for obs, so that one time step
# for one agent is considered one batch entry
...

    def train_policy_low(self, sess, batch, step_train=0, summarize=False, writer=None):
...
# Get target values
...

    def train_decoder(self, sess, dataset, step_train, summarize=False, writer=None):
...
# Downsample obs along the traj dimension
# obs has shape [batch, traj, l_obs]
...
# use the difference between consecutive states in a trajectory, rather than the state
...
# Train decoder
...
# decoder_probs has shape [batch, N_roles]
...

    def process_dataset(self, dataset):
...
# shapes are [batch, ...original dims...]
# [batch, traj, l_obs]
# [batch, traj, l_action]
# [batch, traj]
# [batch, traj, l_obs]
# [batch, traj]
# [batch, N_roles]
...

    def compute_reward(self, sess, agents_traj_obs, z):
...
# Downsample along traj dimension
...
# use the difference between consecutive states in a trajectory, rather than the state
...

    def train_policy_and_decoder(self, sess, dataset, alpha, epsilon, step_train, summarize=False, writer=None):
...
# Downsample obs along the traj dimension
# obs has shape [batch, traj, l_obs]
...
# use the difference between consecutive states in a trajectory, rather than the state
...
# Train decoder
...
# Compute mean and std of batch of P(z|traj)
...
# decoder_probs has shape [batch, N_roles]
...
# Reshape quantities for low-level policy training
...
# duplicate s.t. each time step in trajectory gets the same z
...
# Compute intrinsic reward for each batch entry
...
# Environment reward, sum along trajectory
...
# Train low-level policy
...
# NOTE: intrinsic reward not implemented yet
# Train critic
...
# if true, then 0, else 1

# print(state_home)
...
# print("Team score", state_home[idx_start])
# idx_start += 1
...
# print("Opponent score", state_home[idx_start])
# idx_start += 1
...
# input("Enter: ")

class Alg(object):
...

    def __init__(self, config_alg, n_agents, l_state, l_obs, l_action, nn):
...
# Initialize computational graph
...
# TF summaries
...

    def create_networks(self):
...
# Placeholders
...
# Individual agent networks
# output dimension is [time * n_agents, q-values]
...
# To extract Q-value from agent_qs and agent_qs_target
# [batch*n_agents, l_action]
...
# [batch*n_agents, 1]
...
# [batch, n_agents]
...
# Mixing network
...

    def get_assign_target_ops(self, list_vars):
...
# ops for equating main and target
...
# ops for slow update of target toward main
...
# create op that assigns value of main variable to
# target variable of the same name
...
# incremental update of target towards main
...
# ops for equating main and target
...
# create op that assigns value of main variable to
# target variable of the same name
...
# ops for slow update of target toward main
...
# incremental update of target towards main
...

    def run_actor(self, list_obs, epsilon, sess):
...
# convert to batch
...

    def create_train_op(self):
...
# TD target calculated in train_step() using Mixer_target
...

    def create_summary(self):
...
# mixer_main_variables = [v for v in tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, 'Mixer_main')]
...

    def process_actions(self, n_steps, actions):
...
# Each row of actions is one time step,
# row contains action indices for all agents
# Convert to [time, agents, l_action]
# so each agent gets its own 1-hot row vector
...
# In-place reshape of actions to [time*n_agents, l_action]
...

    def process_batch(self, batch):
...
# shapes are [time, ...original dims...]
# [time, l_state]
# [time, agents, l_obs]
# [time, agents]
# [time]
# [time, l_state]
# [time, agents, l_obs]
# [time]
...
# Try to free memory
...
# In-place reshape for obs, so that one time step
# for one agent is considered one batch entry
...

    def train_step(self, sess, batch, step_train=0, summarize=False, writer=None):
...
# Each agent for each time step is now a batch entry
...
# Get argmax actions from target networks
...
# [batch*n_agents]
# Convert to 1-hot

# print(state_home)
...
# print("Team score", state_home[idx_start])
# idx_start += 1
...
# print("Opponent score", state_home[idx_start])
# idx_start += 1
...
# input("Enter: ")

# Final number of roles
...
# Number of roles increases according to a curriculum
...
# Reward coefficient
...
# Number of single-agent trajectory segments used for each decoder training step
...
# Replay buffer for high level policy
...
# Buffer for low level agent policy
...
# Dataset of [obs traj, z] for training decoder
...
# Variables with suffix '_h' are used for training the high level policy
...
# Cumulative discounted reward for high-level policy
...
# Action taken by high-level policy, limited to current number of roles in curriculum
...
# List of lists, where each sublist is a trajectory of observations by an agent
# Each trajectory contains a history of length at most <steps_per_assign> previous steps,
# Used for computing intrinsic reward
...
# steps within an episode
...
# Run high-level policy as usual to select the subgoal
...
# The environment state at this point, e.g. <state_home>,
# acts like the "next state" for the high-level policy
# All of the intervening environment steps act as a single step for the high-level policy
...
# Store all agents' observation trajectories into dataset for training decoder
...
# Compute intrinsic reward for all agents, returns an np.array of N_agents scalar values
...
# Take high-level action for all agents
...
# Conduct training of high-level policy
...
# Update high-level state
...
# Take low-level actions, conditioned on high level role assignment
...
# Ignore local_rewards from environment. Compute local rewards using intrinsic and global environment reward
...
# Collect low-level observation into trajectories
...
# Limit to be max length <steps_per_assign>
...
# Train low-level policies
...
# Since the episode is done, we also terminate the current role assignment period,
# even if not all <steps_per_assign> have been completed
...
# Append trajectories into dataset, so that decoder sees states that get termination reward
...
# If dataset size is large enough, then train decoder
...
# Decide whether to increase the number of subgoals
...
# Empty the dataset
...
# Evaluation episodes