Sparsifying Parametric Models with L0 Regularization
Nicolò Botteghi Urban Fasel

Quantum features of the transport through ion channels in the soft knock-on model
    @
    September 9, 2024
=================================================================================













This document contains an educational introduction to the problem of sparsifying parametric models with L_0 regularization <cit.>. We utilize this approach together with dictionary learning to learn sparse polynomial policies for deep reinforcement learning to control parametric partial differential equations <cit.>. 

The document is organized as follows: in Section <ref> we introduce the L_0 regularization that we use in our method introduced in <cit.>. In Section <ref>, we introduce the general problem setting for sparsifying parametric models. In Section <ref>, we discuss in more detail all the critical steps to derive the L_0 regularization, and in Section <ref>, we show different ways to use the L_0 regularization in deep reinforcement learning. The code and a tutorial are provided here: <https://github.com/nicob15/Sparsifying-Parametric-Models-with-L0>.






§ SPARSIFYING NEURAL NETWORK LAYERS WITH L_0 REGULARIZATION


To sparsify the weight/coefficient matrix Ξ, the differentiable L_0 regularization method introduced in <cit.> can be used. The method relaxes the discrete nature of L_0 to allow efficient and continuous optimization. 

Let d be a continuous random variable distributed according to a distribution p(d| ψ), where ψ indicates the parameters of p(d| ψ). Given a sample from d∼ p(d|ψ), we can define the hard-sigmoid rectification:

    z = min(1, max(0, d)).

Equation (<ref>) allows z, i.e., the learnable binary gate, to be exactly zero.
Additionally, we can still compute the probability of the gate being active, i.e., non-zero, by utilizing the cumulative distribution function P:

    p(z≠ 0| ψ) = 1 - P(d≤ 0|ψ).

We choose as candidate distribution a binary concrete <cit.>. Thus, the random variable d is distributed in (0,1) with probability density p(d|ψ), cumulative density P(d|ψ), and learnable parameters ψ=[logα, β], with logα the location and β the temperature parameter. The distribution can be stretched to the interval (γ, ζ), where γ < 0 and ζ > 1. Then, the hard-sigmoid on the samples analogously to Equation (<ref>) can be applied:

    u   ∼𝒰(0, 1)  , 
    d   = σ((logu - log (1-u) + logα)/ β)   , 
    d̅   = d(ζ - γ) + γ  , 
    z   = min(1, max(0, d̅))  ,

where σ corresponds to the sigmoid activation function.
We can now optimize the parameters ψ of the distribution by minimizing the probability of the gate being active (see Equation (<ref>)). This optimization problem can be seen as stretching the distribution in (0,1). Using Equation (<ref>) and the binary concrete distribution in Equation (<ref>), we can conveniently introduce the L_0 regularization loss as:

    L_0(ψ) = ∑_j=1^|ξ| (1-P_d̅_j(0|ψ)) = ∑_j=1^|ξ|σ(logα_j - βlogγ/ζ),

where ξ are the parameters of the model we want to sparsify. At test time, we can estimate the sparse parameters ξ^0 by:

    z   = min(1, max(0, σ(logα)(ζ - γ) + γ))  , 
    ξ^0    = ξ⊙z.


In the following two sections, we introduce the above idea and derive the concepts in more detail.  



§ PROBLEM SETTINGS

Suppose we have a dataset 𝒟 of input and output pairs {(x_1, y_1), ⋯, (x_N, y_N)}. We consider a standard supervised learning setting, e.g., regression or classification, with an L_0 regularization to promote sparsity of the parameters ξ of a generic parametric model h:𝒳→𝒴; x ↦ h(x;ξ), e.g., a neural network. We can write the loss function for such a problem as:

    ℒ(ξ)    = 1/N∑_i=1^N F(h(x_i;ξ), y_i) + λ ||ξ||_0,  ||ξ||_0 = ∑_j=1^|ξ|𝕀[ξ_j ≠ 0]  , 
       = ℒ_E + λℒ_C   ,

where F(·) corresponds to a generic loss function, e.g., mean-squared error or cross-entropy, |ξ| is the dimensionality of the parameter vector ξ, λ is a weighting/scaling factor, and 𝕀 is the indicator function. The first term of the loss function ℒ_E corresponds to the error loss that measures how well the model fits the training data, while ℒ_C corresponds to the complexity loss that measures the complexity (or sparsity) of the model. The optimal parameters ξ^* can be found as:

    ξ^*=_ξℒ(ξ).

The L_0 norm penalizes nonzero entries of the parameter vector and encourages sparsity in ξ^*.

Unfortunately, the optimization constitutes an intractable brute force 2^|ξ| combinatorial search due to the nondifferentiability of the L_0 complexity loss function. Usually, L_1 or Lasso, or L_2 norms are used as proxy of the L_0 norm as they are differentiable and can be used with gradient-based optimization techniques. Examples of these norms are visualized in Figure <ref>.

However,  L_1 and L_2 induce an undesirable shrinkage of the parameter values that is not introduced when using L_0.






§ A GENERAL FRAMEWORK FOR L_0 REGULARIZATION

Consider the L_0 norm under reparametrization of ξ:

    ξ^0_j = ξ_j z_j,  z_j ∈{0, 1}, ξ_j ≠ 0,  ||ξ||_0 = ∑_j=1^|ξ|z_j  ,

where z_j corresponds to the binary gate that denotes whether a parameter is present or not. Under this reparametrization, the L_0 norm corresponds to the number of gates being active.
Using the Bernoulli distribution, we can reformulate the loss function in Equation (<ref>) as:

    ℒ(ξ, π) = 𝔼_q(z|π)[1/N∑_i=1^N F(h(x_i;ξ⊙z), y_i)] + λ ||ξ||_0,  ||ξ||_0 = ∑_j=1^|ξ|π_j  ,

where ⊙ indicated the element-wise product.
Analogously to Equation (<ref>), we can find the optimal parameters ξ^* and π^* by solving:

    ξ^*,π^* =_ξ, πℒ(ξ, π).

We could think of choosing a Bernoulli distribution over each gate z_j:

    q(z_j|π_j) = Bern(π_j)  .





[enhanced jigsaw,breakable,pad at break*=1mm,
  colback=gray!5!white,colframe=gray!75!black,title=Bernoulli Distribution]
A Bernoulli distribution Bern(π) describes the distribution of a random variable X taking value 1 with probability π or value 0 with probability 1-π, with π being the (learnable) parameter of the distribution:

    p(X=1) = 1-p(X=0) = π.







The optimization problem in Equation (<ref>) and (<ref>) is a special case of the variational lower bound over the parameters of the neural network involving spike and slab prior <cit.>. 



 §.§ Spike and Slab Distribution and Relation to Variational Inference


The spike and slab distribution, shown in Figure <ref>, is considered the gold standard in sparsity-promoting Bayesian inference/linear regression. 

It is defined as a mixture of a delta spike at 0 and a continuous distribution over the real line, e.g., a standard Gaussian:

    p(z) = Bern(π),  p(ξ|z=0)=δ(ξ),  p(ξ|z=1)=𝒩(ξ|0,1).





The true posterior distribution under the spike and slab prior is intractable. However, we can rely on variational inference <cit.>.

Let q(ξ, z) be a spike and slab approximate posterior over the parameters ξ and the gate variables z. We can write the variational-free energy under the spike and slab prior and approximated posterior over a parameter vector ξ as:

    ℱ   = -𝔼_q(z)q(ξ|z)[log p(𝒟|ξ)] + ∑_j=1^|ξ|KL(q(z_j)||p(z_j)) 
       + ∑_j=1^|ξ|(q(z_j=1)KL(q(ξ_j | z_j=1)||p(ξ_j|z_j=1))) 
       + ∑_j=1^|ξ|(q(z_j=0)KL(q(ξ_j | z_j=0)||p(ξ_j|z_j=0)))  ,

where KL(q(ξ_j | z_j=0)||p(ξ_j|z_j=0))=0 since the KL divergence of two Gaussian distributions with same mean and variance is equal to zero. Therefore, Equation (<ref>) can be rewritten as:

    ℱ   = -𝔼_q(z)q(ξ|z)[log p(𝒟|ξ)] + ∑_j=1^|ξ|KL(q(z_j)||p(z_j)) 
       + ∑_j=1^|ξ|(q(z_j=1)KL(q(ξ_j | z_j=1)||p(ξ_j|z_j=1)))  ,

where the term KL(q(z_j)||p(z_j)) corresponds to the KL from the Bernoulli prior p(z_j) and the Bernoulli approximate posterior q(z_j), and KL(q(ξ_j | z_j=1)||p(ξ_j|z_j=1)) can be interpreted as the amount of information the parameter ξ_j contains about the data 𝒟, measured by the KL divergence from the prior p(ξ_j|z_j=1). 
We can further simplify Equation (<ref>) by assuming, from an empirical Bayesian procedure, the existence of a hypothetical prior p(ξ_j|z_j=1) for each parameter ξ_j that adapts to q(ξ_j|z_j=1) in a way that we need approximately λ NATs (natural units of information) to transform p(ξ_j|z_j=1) to that particular q(ξ_j|z_j=1). Those λ NATs are thus the amount of information that q(ξ_j|z_j=1) can encode about the data if we had used p(ξ_j|z_j=1) as the prior. This assumption translates into KL(q(ξ_j | z_j=1)||p(ξ_j|z_j=1))=λ. The coefficient λ can be viewed as the amount of flexibility of that hypothetical prior. 
Eventually, if we consider optimizing over ξ, instead of integrating, we can write the variational-free energy as:

    ℱ = -𝔼_q(z)[log p(𝒟|ξ⊙z)] + ∑_j=1^|ξ|KL(q(z_j)||p(z_j)) + λ∑_j=1^|ξ|q(z_j=1)  ,

where ξ corresponds to the optimized ξ. Additionally, by using the positivity of the KL divergence, we can obtain the variational lower bound as:

    ℱ≥ -𝔼_q(z)[log p(𝒟|ξ⊙z)] + λ∑_j=1^|ξ|π_j  ,

which is equivalent to Equation (<ref>) if we take the negative log-probability of the data to be equal to the loss ℒ(·). This shows that the minimization of the L_0 norm is very close to the variational lower bound involving a spike and slab distribution over the parameters and a fixed cost/penalty for the parameters when the gates are active. Note that, if we are interested in quantifying uncertainties over the gate variables z, we should optimize Equation (<ref>) (rather than (<ref>)) as this will properly penalize the entropy of q(z). Equation (<ref>) also allows to incorporate prior information about the behaviour of the gates (e.g., being active 10% of the time on average).



 §.§ Efficient Gradient-based Optimization of the L_0 Norm

Minimizing Equation (<ref>) is still not straightforward. While the second term of the loss is easy to minimize, the first term is still challenging due to the discrete nature of the gates z, which does not allow for efficient gradient-based optimization. However, we can replace the Bernoulli distribution to smooth the objective function in Equation (<ref>) and allow for efficient gradient-based optimization of the expected L_0 norm along with zeros in the parameters ξ. 

[enhanced jigsaw,breakable,pad at break*=1mm,
  colback=gray!5!white,colframe=gray!75!black,title=Gradient Estimators]
In principle, it is possible to estimate the gradient using REINFORCE <cit.>. However, REINFORCE suffers from high variance of the estimates, requires auxiliary models, and multiple evaluations <cit.>. An alternative is to use the straight-through estimator <cit.> as done in <cit.> or the concrete distribution <cit.>. Unfortunately, the first method provides biased gradients (due to ignoring the Heaviside function in the likelihood during the gradient evaluation), while the second one does not allow for the gates to be exactly zero during the optimization (thus precluding the benefits of conditional computation <cit.>).





Let d be a continuous random variable with a distribution q(d) of parameters ψ. We can now define the gates z as hard-sigmoid rectifications of d such that:

    d   ∼ q(d|ψ)  , 
    z   = min(1, max(0, d))) = g(d).

In this way, the gate is allowed to be exactly zero. Due to the underlying continuous random variable d, we can still compute the probability of the gate being nonzero, i.e., active, from the cumulative distribution function (CDF) Q(d|ψ):

    q(z≠ 0| ψ) = 1 - Q(d≤ 0|ψ)  ,

where 1 - Q(d≤ 0|ψ) corresponds to the probability of d being positive. This allows to replace the binary Bernoulli gates in Equation (<ref>) with the CDF in Equation (<ref>):

    ℒ(ξ,ψ) = 𝔼_q(d|ψ)[1/N∑_i=1^N F(h(x_i; ξ⊙ g(d)), y_i)] + λ∑_j=1^|ξ|(1-Q(d_j≤0|ψ_j)).

Analogously to Equation (<ref>), we can find the optimal parameters as:

    ξ^*, ψ^* = _ξ, ψℒ(ξ,ψ)


For a continuous distribution q(d|ψ) that allows for the reparametrization trick <cit.>, we can express Equation (<ref>) as the expectation over a parameter-free noise distribution p(ϵ) and a deterministic and differentiable transformation f(·) of the parameters ψ and ϵ:

    ℒ(ξ,ψ) = 𝔼_p(ϵ)[1/N∑_i=1^N F(h(x_i; ξ⊙ g(f(ψ, ϵ))), y_i)] + λ∑_j=1^|ξ|(1-Q(d_j≤0|ψ_j))  .

This allows for a Monte Carlo approximation to the generally intractable expectation over the noise distribution p(ϵ):

    ℒ(ξ,ψ)    = 1/L∑_l=1^L (1/N∑_i=1^N F(h(x_i; ξ⊙z^(l)), y_i)) + λ∑_j=1^|ξ|(1-Q(d_j≤0|ψ_j))  ,  
    z^(l)   = g(f(ψ, ϵ^(l)))  , ϵ^(l)∼ p(ϵ) .

Equation (<ref>) is differentiable with respect to ψ and can be used with (stochastic) gradient-based optimization, while still allowing the parameters to be exactly zero. Additionally, we can choose an appropriate smoothing distribution q(d|ψ) for our problem. 

A choice that works well in practice is the binary concrete <cit.>, with CONCRETE derived from CONtinuous relaxation of disCRETE distribution <cit.>. 




  §.§.§ The Hard-Concrete Distribution

Assume to have  a binary concrete random variable d distributed in (0, 1) with a probability density function (PDF) q(d|ψ):

    q(d|ψ) = βα d^-β-1(1-d)^-β-1/(α d^-β+(1-d)^-β)^2  ,

 and a cumulative distribution function Q(d|ψ):
 
    Q(d|ψ) = σ(log d - log(1-d)+logα/β)  ,

where the distribution has parameters (logα, β), with logα the location and β the temperature, and σ indicates the sigmoid activation function.

We can stretch the distribution in the (γ, ξ) interval, with γ < 0 and ξ > 1 to obtain:

    d  = d(ξ - γ) + γ.

The stretching of d induces the following probability density function and cumulative distribution function q(d|ψ) and Q(d|ψ):

    q(d̅|ψ)    = 1/|ξ - γ|q(d-γ/ξ - γ|ψ)  ,
    
            Q(d̅|ψ)    = Q(d-γ/ξ - γ|ψ).


We can further rectify d with the hard-sigmoid  such that:

    z = min(1, max(0, d)).

We obtain the following distribution over z, which is referred to as the hard-concrete distribution:

    q(z|ψ) = Q_d(0|ψ) δ(z) + (1- Q_d(1|ψ))δ(z-1) + (Q_d(1|ψ)) - Q_d(0|ψ)))q_d(z|d∈ (0,1), ψ)  .

This distribution is composed of a delta peak at zero with probability Q_d(0|ψ)), a delta peak at one with probability 1-Q_d(1|ψ)), and a truncated version of q_d(d|ψ) in the range (0, 1).


In practice, we can sample a gate z from the hard-concrete distribution by first applying the reparametrization trick, then sampling from a uniform distribution 𝒰(0,1), feeding the sample to the binary concrete CDF, stretching it, and finally passing it to the hard-sigmoid rectification:

    u    ∼𝒰(0, 1) 
     
            d    = σ(log u - log(1-u)+logα/β) 
    d   = d(ξ - γ) + γ
    
            z    = min(1, max(0, d)).

An example of a resulting gate z is shown in Figure <ref>.




We can derive the CDF in Equation (<ref>) as:

    Q(d|ψ)    = ∫_-∞^s q(x|ψ) dx   , 
       = ∫_-∞^s βα x^-β-1(1-x)^-β-1/(α x^-β+(1-x)^-β)^2 dx.

Solving the integral is very challenging due to the form of the PDF. However, we can rely on the Gumbel-Max trick <cit.>. We start from the observation that Bernoulli random variables are a special case of discrete distributions with states in {0, 1}. Consider a discrete distribution D ∼Discrete(α), with α∈ (0, ∞) the parameter of the discrete distribution, and a two-state discrete random variable on {0, 1}^2 such that D_1 + D_2 = 1:

    Q(D_1=1) = α_1/α_1 + α_2.





[enhanced jigsaw,breakable,pad at break*=1mm,
  colback=gray!5!white,colframe=gray!75!black,title=The Gumbel-Max Trick]
The Gumbel-Max trick proceeds as follows:

    
  * sample U_k  ∼𝒰(0, 1) with k=1, 2,
    
  * find the index k maximizing log(log U_k) + logα_k, and
    
  * set D_k=1 and the remaining D_i≠ k=0.






Therefore, if we apply the Gumbel-Max trick, the case of D_1=1 corresponds to the event:

    log(log U_1) + logα_1    > log(log U_2) + logα_2   , 
    
        G_1+ logα_1    > log(log U_2) + logα_2 .

The difference between log(log U_1)=G_1 and log(log U_2)=G_2 is a logistic distribution L with CDF the logistic function σ(x) = 1/(1+exp^-x):

    G_1 - G_2 ∼ L(U) = log U - log(1-U)   ,

where U ∼𝒰(0,1).
Therefore, we can rewrite Equation (<ref>) as:

    Q(D_1=1)    = Q(log(log U_1) + logα_1 > log(log U_2) + logα_2) 
       = Q(G_1 - G_2 + logα_1 - logα_2) 
       = Q( log U - log(1-U) + logα > 0)  ,

where logα = logα_1 - logα_2.

Eventually, we can define the binary concrete random variable as:

    Z = σ(log U - log(1-U) + logα/β)

and its CDF as:

    Q(z)    = Q(Z ≤ z) = Q(σ(log U - log(1-U) + logα/β)≤ z) 
       = Q(log U - log(1-U) + logα/β≤σ(z)^-1) 
       = Q(log U - log(1-U) + logα/β≤log(z/1-z)) 
       = Q(log U - log(1-U) + logα≤βlog(z/1-z))

where σ(z)^-1=log(z/1-z) is the inverse of σ(z). Equation (<ref>) is equivalent to (<ref>).





§ DEEP REINFORCEMENT LEARNING WITH L_0 REGULARIZATION

To conclude the tutorial, we show how to use the differentiable L_0 regularization in the context of deep reinforcement learning to learn transition models, reward models, and control policies. We use as test case the pendulum environment from OpenAI Gym <cit.> <https://www.gymlibrary.dev/environments/classic_control/pendulum/>.



 §.§ Learning Transition and Reward Models




  §.§.§ Dataset Collection


We start by creating a training set of 1000 episodes and a test set of 100 episodes. These dataset are obtained by applying a random policy to the environment for 200 steps per episodes. The code is shown and explained in Code <ref> and <ref>.





[language=Python, caption=Creation of training set., label=lst:train_set]
import gymnasium as gym
from utils.replay_buffer import ReplayBuffer

# initialize gym environment
env = gym.make('Pendulum-v1', g=9.81)

# set maximum number of episodes and steps
max_episodes = 1000
max_steps = 200

# get dimension of the observations, actions, and memory buffer
obs_dim = env.observation_space.shape[0]
act_dim = env.action_space.shape[0]
buf_dim = int(max_episodes*max_steps)

# initialize training buffer (we will store the training data here)
training_buffer = ReplayBuffer(obs_dim=obs_dim, act_dim=act_dim, size=buf_dim)

# create training set
for episode in range(max_episodes):

    # reset the environment at the beginning of each episode
    observation, _ = env.reset() 
    
    for steps in range(max_steps+1):

        # select action according to a random policy
        action = env.action_space.sample() 

        # apply the action to the environment
        next_observation, reward, terminated, truncated, _ = env.step(action) 

        done = terminated or truncated

        # store data tuple in the training buffer
        training_buffer.store(observation, action, reward, next_observation, done)

        # set next observation as the current observation
        observation = next_observation

        if done:
            break



[language=Python, caption=Creation of test set., label=lst:test_set]
# set maximum number of episodes 
max_episodes_test = 100

# set dimension of memory buffer
buf_dim = int(max_episodes*max_steps)

# initialize testing buffer (we will store the test data here)
testing_buffer = ReplayBuffer(obs_dim=obs_dim, act_dim=act_dim, size=buf_dim)

# create test set
for episode in range(max_episodes_test):

    # reset the environment at the beginning of each episode
    observation, _ = env.reset()
    
    for steps in range(max_steps + 1):

        # select action according to a random policy
        action = env.action_space.sample()

        # apply the action to the environment
        next_observation, reward, terminated, truncated, _ = env.step(action)

        done = terminated or truncated

        # store data tuple in the testing buffer
        testing_buffer.store(observation, action, reward, next_observation, done)

        # set next observation as the current observation
        observation = next_observation

        if done:
            break





  §.§.§ Transition and reward models

We first show how to learn:

    
  * a transition model T: 𝒮×𝒜→𝒮, mapping a state-action pair (s_t, a_t) to the next state s_t+1, and
    
  * a reward model R: 𝒮×𝒜→ℝ, mapping a state-action pair (s_t, a_t) to the reward r_t.
 
We approximate T and R in three different ways:

    
  * using a fully-connected neural network with parameters ξ_T and ξ_R, respectively,
    
  * using a fully-connected neural network trained with L_0 regularization with parameters ξ_T and ψ_T and ξ_R and ψ_R, and
    
  * using a SINDy-like model <cit.> trained with L_0 regularization with parameters ξ_T and ψ_T and ξ_R and ψ_R (this is analogous to the sparse policy introduced in <cit.>). We utilize a polynomial-feature library, a Fourier library, and a generalized library composed of polynomial and Fourier features.


The fully-connected neural network model is shown in Code <ref>. The models is composed of three layers with ELU activation for the first two. For learning transition and reward model, we utilize the same architecture.
[language=Python, caption=Fully-connected neural network model., label=lst:FCNN_model]
class FCNN(nn.Module):
    def __init__(self, input_dim=3, output_dim=1, h_dim=256, use_bias=True):
        super(FCNN, self).__init__()

        self.use_bias = use_bias
        
        self.fc = nn.Linear(input_dim, h_dim, bias=use_bias)
        
        self.fc1 = nn.Linear(h_dim, h_dim, bias=use_bias)
        
        self.fc2 = nn.Linear(h_dim, output_dim, bias=use_bias)

    def forward(self, obs, act):
        # concatenate observation and action before feeding them to the input layer
        x = torch.cat([obs, act], dim=1)
        x = F.elu(self.fc(x))
        x = F.elu(self.fc1(x))
        # the output is either the next observation or the reward
        out = self.fc2(x)
        return out

Similarly to the fully-connected neural network, the sparse fully-connected neural network is composed of three layers with ELU activations. However, the layers include the L_0 mask introduced in <cit.>.
[language=Python, caption=Sparse fully-connected neural network model., label=lst:sparseFCNN_model]
class SparseFCNN(nn.Module):
    def __init__(self, input_dim=3, output_dim=1, h_dim=256, weight_decay=0.,                           droprate_init=0.5, temperature=2./3., lambda_coeff=1.):
        
        super(SparseFCNN, self).__init__()

        self.fc = L0Dense(in_features=input_dim, out_features=h_dim, bias=True,                             weight_decay=weight_decay, droprate_init=droprate_init,                            temperature=temperature, lamba=lambda_coeff, local_rep=False)
        
        self.fc1 = L0Dense(in_features=h_dim, out_features=h_dim, bias=True,                                weight_decay=weight_decay, droprate_init=droprate_init,                            temperature=temperature, lamba=lambda_coeff, local_rep=False)
        
        self.fc2 = L0Dense(in_features=h_dim, out_features=output_dim, bias=True,                           weight_decay=weight_decay, droprate_init=droprate_init,                            temperature=temperature, lamba=lambda_coeff, local_rep=False)

    def forward(self, obs, act):
        # concatenate observation and action before feeding them to the input layer
        x = torch.cat([obs, act], dim=1)
        x = F.elu(self.fc(x))
        x = F.elu(self.fc1(x))
        # the output is either the next observation or the reward
        out = self.fc2(x)
        return out

Eventually, in Code <ref>, we show the structure of the L_0 SINDy like model with three different feature libraries. The model makes use of the library features of the PySINDy library <https://github.com/dynamicslab/pysindy> and of a single sparse layer to learn the coefficients of each feature. The sparse layer allows for utilizing the L_0 regularization introduced in <cit.> to learn a sparse linear combination of the nonlinear library features.
[language=Python, caption=L_0 SINDy-like model., label=lst:L0SINDy_model]
    class L0SINDy_model(nn.Module):
    def __init__(self, input_dim=3, output_dim=1, weight_decay=0., droprate_init=0.5,                   temperature=2. / 3., lambda_coeff=1., degree=3, frequency=1,                           lib_type='polynomial'):
        
        super(L0SINDy_reward, self).__init__()

        if lib_type == 'polynomial':
            self.lib = PolynomialLibrary(degree=degree, include_bias=True,                                                 include_interaction=True)
            x = np.ones((1, input_dim))
            self.lib.fit(x)
            xf = self.lib.transform(x)
            coef_dim = xf.shape[1]

        if lib_type == 'fourier':
            self.lib = FourierLibrary(n_frequencies=frequency, include_sin=True,                                       include_cos=True, interaction_terms=True)
            x = np.ones((1, input_dim))
            self.lib.fit(x)
            xf = self.lib.transform(x)
            coef_dim = xf.shape[1]

        if lib_type == "polyfourier":
            poly_lib = PolynomialLibrary(degree=degree, include_bias=True,                                                 include_interaction=True)
            fourier_lib = FourierLibrary(n_frequencies=frequency, include_sin=True,                                        include_cos=True, interaction_terms=True)
            self.lib = GeneralizedLibrary([poly_lib, fourier_lib])
            x = np.ones((1, input_dim))
            self.lib.fit(x)
            xf = self.lib.transform(x)
            coef_dim = xf.shape[1]

        # single sparse layer learning the coefficents of each dictionary feature
        self.fc = L0Dense(in_features=coef_dim, out_features=output_dim, bias=False,                         weight_decay=weight_decay, droprate_init=droprate_init,                            temperature=temperature, lamba=lambda_coeff, local_rep=False)

    def forward(self, obs, act):
        # concatenate observation and action before feeding them to the input layer
        x = torch.cat([obs, act], dim=1)
        # compute the features of x using the chosen library functions
        xf = torch.from_numpy(self.lib.transform((x).cpu().numpy())).to(self.device)
        # feed the features to a single layer of sparse fully-connected neural network
        out = self.fc(xf)
        return out




  §.§.§ Training the Models

We indicate with ξ the parameters of the different models and with ψ the L_0-mask parameters.
We train the models to minimize the mean-squared error between the predictions and the ground truth. In particular, we optimize the fully-connected neural network parameters with:

    ℒ(ξ_T) = 𝔼_s_t, a_t, s_t+1[||s_t+1 - ŝ_t+1||^2_2]  ,

and

    ℒ(ξ_R) = 𝔼_s_t, a_t, r_t[||r_t - r̂_t||^2]  ,

where ŝ_t+1=T(s_t, a_t, ξ) is the prediction of the transition model T and  r̂_t=R(s_t, a_t, ξ) is the prediction of the reward model R, respectively.

The models using the L_0 regularization are instead optimized according to Equation (<ref>):

    ℒ(ξ_T, ψ_T) = 𝔼_s_t, a_t, s_t+1[||s_t+1 - ŝ_t+1||^2_2] + λL_0(ψ_T)  ,

and

    ℒ(ξ_R, ψ_R) = 𝔼_s_t, a_t, r_t[||r_t - r̂_t||^2] + λL_0(ψ_R).


Code <ref> and <ref>, we show the training loop for the transition and reward models.
[language=Python, caption=Transition model training loop., label=lst:dyn_model_training]
def train_dynamics_model(model, optimizer, train_loader, batch_size,                                                num_training_iterations, l0=False):

    for _ in range(num_training_iterations):

        # random sample a batch of data
        data = train_loader.sample_batch(batch_size)

        obs = torch.from_numpy(data['obs']).cuda()
        next_obs = torch.from_numpy(data['next_obs']).cuda()
        act = torch.from_numpy(data['act']).cuda()

        optimizer.zero_grad()

        # predict the next observation using the parametric model
        pred_next_obs = model(obs, act)

        # If we do not rely on the L0 regularization, we simply utilize the MSE loss
        if l0 == False:
            loss = torch.nn.functional.mse_loss(next_obs, pred_next_obs)
            total_loss = loss
        # If we on the L0 regularization, we add the L0 penalty to the MSE loss
        else:
            loss = torch.nn.functional.mse_loss(next_obs, pred_next_obs)
            reg = -(model.fc.regularization())
            total_loss = loss + reg

        # compute the loss gradients wrt the model parameters
        total_loss.backward()

        # update the parameters
        optimizer.step()


[language=Python, caption=Reward model training loop., label=lst:rew_model_training]
def train_reward_model(model, optimizer, train_loader, batch_size,                                                num_training_iterations, l0=False):

    for _ in range(num_training_iterations):

        # random sample a batch of data
        data = train_loader.sample_batch(batch_size)

        obs = torch.from_numpy(data['obs']).cuda()
        rew = torch.from_numpy(data['rew']).cuda().reshape(-1, 1)
        act = torch.from_numpy(data['act']).cuda()

        optimizer.zero_grad()

        # predict the reward using the parametric model
        pred_rew = model(obs, act)

        # If we do not rely on the L0 regularization, we simply utilize the MSE loss
        if l0 == False:
            loss = torch.nn.functional.mse_loss(rew, pred_rew)
            total_loss = loss
        # If we on the L0 regularization, we add the L0 penalty to the MSE loss
        else:
            loss = torch.nn.functional.mse_loss(rew, pred_rew)
            reg = -(model.fc.regularization())
            total_loss = loss + reg

        # compute the loss gradients wrt the model parameters
        total_loss.backward()

        # update the parameters
        optimizer.step()




  §.§.§ Results

In Figure <ref>, we show the accuracy of the different transition and reward models on the training and test sets after 500 training epochs. While the fully-connected neural network seems to be training faster than the sparse models, we need to keep in mind that the L_0 regularization makes the minimization of the loss slightly more challenging due to the need for minimizing not only the MSE but also the number of parameters. Therefore, especially for the sparse models, it is worth mentioning that the results can be further improved with longer training time and hyperparameter optimization techniques. 




 §.§ Learning Control Policies

We utilize the method proposed in <cit.> to learn sparse, and interpretable policies for the pendulum example. We replace the neural network-based policy of the twin-delayed deep deterministic policy gradient algorithm <cit.> with a polynomial, Fourier, and a polynomial and Fourier policies that are sparsified over training using the L_0 regularization (all the details can be found in our paper  <cit.>). It is worth mentioning that our method is independent of the deep reinforcement learning algorithm chosen and other algorithms can be used.



  §.§.§ Results


In Figure <ref>, we show the training and evaluation rewards collected by the different agents in the simple task of stabilizing the inverted pendulum in its unstable equilibrium. Due to the simplicity of the task, we do not see substantial differences among the agents with the exception of the Fourier agent. This is due to the fact that a policy only composed of sines and cosines is not suitable for a stabilization task in a single point. However, such a policy may be useful for periodic tasks. Due to the limited number of learnable parameters, the agent with polynomial feature library learns to solve the task slightly faster than the fully-connected neural network agent. In addition, the agents relying on the sparse dictionaries allow for deriving a closed-form equation of the learned policies, opening the door to a-posteriori stability and robustness analysis. 


unsrt
 .9
 
 
 
 
 






