Symbolic-Model-Based Reinforcement Learning
Abstract
We investigate using symbolic regression (SR) to model dynamics with mathematical expressions in model-based reinforcement learning (MBRL).
While the primary promise of MBRL is to enable sample-efficient learning, most popular MBRL algorithms rely, in order to learn their approximate world model, on black-box over-parametrized neural networks, which are known to be data-hungry and are prone to overfitting in low-data regime.
In this paper, we leverage the fact that a large collection of environments considered in RL is governed by physical laws that compose elementary operators e.g sin , √ , exp , ddt , and we propose to search a world model in the space of interpretable mathematical expressions with SR.
We show empirically on simple domains that MBRL can benefit from the extrapolation capabilities and sample efficiency of SR compared to neural models.
Introduction
Motivated by real-world applications, such as robotic controlling [1, 2], one of the main goals of the field of Reinforcement Learning (RL) is to learn to control systems with non-linear (potentially stochastic) dynamics.
Most control systems, irrespective of the task to solve (i e. rewards maximization) have in common the fact that their dynamics are governed by physical laws.
They are usually expressed with mathematical equations connecting the next state with the system’s past states and controller’s actions with operators such as time-derivatives, trigonometric operators, power functions.
Solving a task usually involves implicitly a good understanding of the dynamics, e.g. goal reaching.
For instance, in the classic CartPole environment [3, 4] the system’s state is described by ( x, ˙ x, θ, ˙ θ ) where x (resp. ˙
x ) denotes the position (resp. speed) of the cart along the x -axis and θ (resp. ˙
θ ) is the angle (resp.
rotation) of the pole w.r.t the cart.
The agent’s action a affects the system’s state according to the following equations [4, 5]:
[Equation]
where g, K mag , m p , m c , l are all constants.
Systems with impossible states, safety or physical constraints, e.g. the joint of a robotic arm cannot exceed a certain angle [6, 7], can be expressed via piece-wise expressions.
When the environment is stochastic, probabilistic distributions appear.
Note that the parametrization of the state and action spaces has an impact on the symbols that are necessary to describe the system dynamics; for instance, if ˙ x and ˙ θ were not in the CartPole agent’s observation, derivative operators would have to appear in the equation.
Generally, the larger the state-space parametrization, the "shorter" equation is.
Model-Based Reinforcement Learning (MBRL) is a class of RL algorithms that involves a two-step procedure repeated until task termination; a) learn from data a forward dynamics model (possibly stochastic) function f that maps current state s t and action a t to next state s t +1 and b) derive a policy from this model.
Though they have been shown to learn faster than model-free algorithms theoretically [8] and in certain applications [9, 10], MBRL algorithms are often hard to train in practice [11].
In this work, we challenge the go-to approach of using over-parametrized feed-forward neural networks to approximate f as they are prone to overfitting when collected data does not have enough coverage of state and action spaces.
We propose to leverage prior knowledge on operators that could appear in the environment dynamics equations, e.g. sin , √ , exp , ddt .
Manipulating expressions to fit data is exactly the objective of SR algorithms that select f within a large family of expressions through composition of operators, constants and variables, as opposed to gradient-descent of over-parametrized models, e.g. neural networks (NNs).
The latter have more degrees of freedom and are easier to optimize, but prone to overfitting in low-data regimes, whereas
SR has recently shown excellent extrapolation capabilities [12, 13, 14].
Furthermore, SR provides an interpretable and differential model, interesting properties for RL, which we develop in section 4.
Contributions.
We propose a novel approach to dynamics modelling in control problems, which we call Symbolic-Model-Based RL, that uses mathematical expressions to model dynamics.
To our knowledge, this is the first work that proposes to leverage SR to find an interpretable function f that best maps state-action pairs ( s t , a t ) to the next state s t +1 .
We provide empirical evidences in simple domains, where our method largely outperforms the over-parametrized approaches, that SR provides faster better dynamics models that generalize to unseen states-action pairs.
Related work
MBRL.
We study the RL problem where an agent interacts with an unknown environment, formalized by a Markov Decision Process ( S , A , p, r ) where S (resp.
A ) represents the state (resp. action) space, p the transition dynamics and r the reward functions.
MBRL algorithms are a class of RL algorithms that ground control policies on a model of the dynamics of environment.
Most of approaches alternate two steps repeatedly: i) collect data D with the current policy and learn an approximate model f of the environment’s dynamics, fitting D as in supervised learning (SL), i.e.:
[Equation]
where F is a family of functions, e.g. neural networks or Gaussian processes, and L a loss function that depends on the nature of f ; ii) then, simulate transitions with the model f and optimize the policy accordingly.
Note that in this work we consider reward and termination functions learned in the same way as dynamics, even though MBRL algorithms [15, 16] often consider them provided to the agent.
We refer the reader to an exhaustive MBRL review [17].
Step i) usually faces classic under/over-fitting issues of SL (w.r.t the state and action space), causing sub-optimal task performance.
The design of predictive model has proven to be very challenging; Gaussian Processes (GPs) can under-fit on complex dynamical systems [18], whereas over-parametrized functions, i.e. neural networks, can express complex (and high-dimensional) dynamics, but are prone to overfitting.
To avoid overfitting, one can a) acquire more data, but this often comes with exploration challenges, b) use regularization, i.e. making the model simpler, in the form of priors, e.g. GPs’ kernel function or Bayesian neural networks [19, 20], or c) use ensembles [21, 22, 23].
[15] uses probabilistic NNs (b) in combination with ensembles (c) to make MBRL agents uncertainty-aware for better planning.
Symbolic regression.
SR, the science of inferring mathematical laws from experimental data by searching over the space of interpretable mathematical expression by manipulating operators, constants and variables, has already proved useful in a variety of domain such as physics [24, 25, 26, 27].
The leading algorithms for SR, excellently reviewed in [12], mainly rely on genetic programming (GenP) [28, 29, 30] by iterating through steps of selection, mutation and crossover until a satisfactory accuracy level is achieved.
More recently, neural networks have been applied to SR to identify qualitative patterns to reduce the search space, either in combination with GenP algorithms [24, 31] or without [32, 33, 34, 35, 13].
Recently, SR has been applied to model-free RL to learn interpretable policies explicitly [36, 37, 38] or via the associated value function [39], but to our knowledge we are the first to consider SR for MBRL.
In addition to the immediate interpretability benefit, considering the dynamics model search space F in Eq.
(2) to be the family of short mathematical expressions that contain constants, variables and operators from a given dictionary has the advantages of injecting prior knowledge, smoothness properties, as well as to significantly reduce the size of the search space.
Experiments
Description of the algorithm.
We propose to use an expression optimized via SR, instead of the usual NN dynamics model, to fit data in a MBRL’s algorithm (step i)); though the symbolic approach is applicable to most MBRL algorithm, we use Probabilistic Ensembles with Trajectory
Sampling (PETS) [15] implemented with MBRL-lib [11] as our base MBRL algorithm and Operon [30], an efficient GenP algorithm, our base SR algorithm (details in Appendix A.1) Preliminary experiments on the considered environments, showed that Operon had better performance that numerous symbolic regressors, e.g. gplearn [40], PS-Tree [41] and AI-Feynman [24]; they either had inaccurate predictions, overly complex expressions or inference was too long.
As in PETS, we maintain an ensemble of 7 dynamics models and select the sequence of actions that maximize rewards using the cross-entropy method [42] on the dynamics simulated with trajectory sampling.
We call this model Symbolic-PETS and compare it to MLP-PETS , the original version of PETS.
Our experiments investigate the following questions: i) Can SR algorithms help learn predictive models with less samples, ii)
Do the obtained dynamics equations have interesting properties?
iii)
How is this manifesting in terms of performance (task solving)?
We consider deterministic environments, which already represent a substantial of environments in the RL literature [6, 7, 43].
An illustrating example.
For illustrative purposes, we consider a simple one-dimensional state
MDP where an agent moves on the horizontal axis x while observing its position, with episode length 10 and the following dynamics:
[Equation]
As illustrated on Fig. solving this MDP is challenging as it requires sufficient exploration to learn the reward function and avoid falling into local minimas i e. staying in relative integers x values without going to better ones.
As shown in Fig.
MLP-PETS’ dynamics models over-fits on the training distribution.
Even using an ensemble of 7 models, the uncertainty of the ensemble is not
[Figure]
n
[Figure]
t
[Figure]
p s
[Figure]
o d e l e rr o r
( r a n d o m p
o li c y )
[Figure] good enough to be leveraged for efficient exploration, leading to a sub-optimal policy. On the other hand, Symbolic-PETS’ dynamics models extrapolate really well on unseen states, thus achieving optimal behavior in just one update.
In Appendix B.1.2, we present the evolution of Symbolic-PETS every episodes after just 20 observed random transitions.
CartPole.
We consider the continuous CartPole [6], where the agent state is s t = ( x t , θ t , ˙ x t , ˙ θ t ) . As in [15, 11], the termination and reward function are made available to the agent, therefore control is restricted to be a problem of dynamics modelling.
We define the model error L in Eq.
(2) as the MSE averaged over output dimensions.
In a first experiment, we explore the capabilities of both symbolic and MLP regressors on the data generated by a random policy on CartPole.
We collect an evaluation dataset of 5 e 4 transitions, enough to have good state-action coverage.
We then train the two dynamics models on training datasets of growing sizes and plot the model error in Fig.
5 (App. B.2.2).
The symbolic model predicts the following equations with the accuracy reached by the MLP in two order of magnitude less interactions:
[Equation]
[Table]
Interestingly, predicted equations are a bit different than in Eq.
(1), though we can notice constants such as the time-discretization interval 0 .
What could be as missing terms can be explained by limited development as θ and x have small values because of CartPole’s constraints.
In Fig 3, we demonstrate that Symbolic-PETS is able to solve CartPole in just 20 interactions with the environment, a state-of-the-art performance to our knowledge.
Discussion
We demonstrated on simple environments that SR can learn better predictive models with many less samples than neural network and that they extrapolate well to unseen state-action pairs much better.
Though [12] showed promising results on real-world regression datasets with input dimensions up to 124 (with a single output dimension), SR still remains to be scaled to higher output dimensions, with challenges including parallelizing regressor training of each output dimension or pixel observations.
SR can have significant impact on multiple RL research topics, e.g partially-observable environments, or meta- and continual-RL thanks to sample efficiency of SR.
Having an interpretable (and differentiable) dynamics model can enable constrained reinforcement learning [44], environment design [45] and domain randomization [46] to enable "Sim2Real" transfer to new environments.
Yet, SR current tools are not as modulable as spatio-temporal NNs.
To be useful in the aforementioned RL domains, we aim at designing algorithms and their implementation to respect these non-exhaustive properties of neural networks; i) can express distributions to handle epistemic and aleatoric uncertainties, ii) learning can be initialized to some known expression, iii) allow multi-step predictions and learning and iv) preprocessing of inputs (raw pixel environments). A Method details A.1 Operon
We used Operon [30] with the following allowed operators add , sub , mul , div , sin , cos , pow for nodes.
Leaves are either variables, i.e. a state feature or a numerical constant.
We use the following hyper-parameters: 5 local iterations, a population of size 5000 , a total of 10000 generations andthreads.
B Experiment details
We perform our experiments with 1 GPU and 1 CPU.
B.1 Toy example
B.1.1 Model hyper-parameters.
MLP-specific hyper-parameters.
We consider a deterministic neural network with 4 hidden layers of size 200 with SiLU activation, trained with Adam optimizer during a maximum of 2000 epochs with batch size 256 and patience epochs 25 (meaning training stops when loss/ evaluation score does not progress more than 0 .
01 relatively), learning rate 7 .
5 e − 4 and weight decay 3 e − 5 .
Inputs are normalized.
Symbolic-specific hyper-parameters.
We use
Operon with local iterations, generations, threads, population size and allowed symbols are add,sub,mul,div,constant,variable,sin,exp,abs .
Action optimizer.
We use CEM with planning horizon 3 , 10 iterations, elite ratio 0 .
1 , population size 1000 , alpha 0 . 1 and clipped normal action distribution.
Out of the ensemble of 7 predictive models, only the 3 elite (w.r.t evaluation set) ones are used.
B.1.2 Results
We present the evolution of Symbolic-PETS every episodes after just 20 random transitions in Fig.
B.2 CartPole
B.2. 1 Model hyper-parameters
MLP-specific hyper-parameters.
Same as for the toy example.
Symbolic-specific hyper-parameters.
We use
Operon with local iterations, generations, threads, population size and allowed symbols are add,sub,mul,div,constant,variable,sin,cos,pow .
Action optimizer.
We use CEM with planning horizon 15 , 5 iterations, elite ratio 0 .
1 , population size 350 , alpha 0 .
Out of the ensemble of 7 predictive models, only the 3 elite (w.r.t evaluation set) ones are used.
B.2.2 Analysis of results
Interestingly, one can notice in Fig.
3 the performance (in terms of reward and model error) fluctuate a bit as the number of interactions grows (contrary to Fig. 5).
This can be explained by the fact that random samples are the most informative transitions in terms on environment understanding, whereas on-policy transitions (whose proportion grows during learning) are all located in a very narrow part of the state-space (pole standing still)
[Figure]
u ff
[Figure]
r
[Figure]
d
[Figure]
i o n
[Figure]
i
[Figure]
n T
i
[Figure]
R e p l a y b u ff e r s t a t e d
i s t r i b u t i o n
[Figure]
v
[Figure]
t
[Figure]
T
[Figure]
R e p l a y b u ff e r s t a t e d
i s t r i b u t i o n
[Figure] [Figure]