# Efficient Skill Acquisition for Complex Manipulation Tasks
in Obstructed Environments ††thanks: 1Applied AI Lab (A2I), Oxford Robotics
Institute, University of Oxford††thanks: Correspondence to:
<EMAIL_ADDRESS>
Jun Yamada1, Jack Collins1, Ingmar Posner1

###### Abstract

Data efficiency in robotic skill acquisition is crucial for operating robots
in varied small-batch assembly settings. To operate in such environments,
robots must have robust obstacle avoidance and versatile goal conditioning
acquired from only a few simple demonstrations. Existing approaches, however,
fall short of these requirements. Deep reinforcement learning (RL) enables a
robot to learn complex manipulation tasks but is often limited to small task
spaces in the real world due to sample inefficiency and safety concerns.
Motion planning (MP) can generate collision-free paths in obstructed
environments, but cannot solve complex manipulation tasks and requires goal
states often specified by a user or object-specific pose estimator. In this
work, we propose a system for efficient skill acquisition that leverages an
object-centric generative model (OCGM) for versatile goal identification to
specify a goal for MP combined with RL to solve complex manipulation tasks in
obstructed environments. Specifically, OCGM enables one-shot target object
identification and re-identification in new scenes, allowing MP to guide the
robot to the target object while avoiding obstacles. This is combined with a
skill transition network, which bridges the gap between terminal states of MP
and feasible start states of a sample-efficient RL policy. The experiments
demonstrate that our OCGM-based one-shot goal identification provides
competitive accuracy to other baseline approaches and that our modular
framework outperforms competitive baselines, including a state-of-the-art RL
algorithm, by a significant margin for complex manipulation tasks in
obstructed environments.

##  I Introduction

Teaching new skills to robots using limited supervision is essential for
maximising the up-time and productivity of robots, leading to faster return on
investment. Small-batch manufacturing, where there are a limited number of
parts to be produced, is an exemplary environment that would greatly benefit
from efficient skill acquisition. In a small-batch setting, a robot must learn
to manipulate new objects while maintaining data efficiency in potentially
arbitrarily obstructed environments. However, existing methods such as motion
planning and reinforcement learning struggle to satisfy such requirements.

Motion planning (MP) [1, 2, 3] generates collision-free paths capable of
guiding a robot safely in obstructed environments given an explicit state of
the environment and goal. However, MP is not designed to plan through complex
manipulation tasks requiring environmental interaction. Furthermore, MP
necessitates the specification of a goal state in the robot’s frame of
reference, which is typically accomplished through manual engineering [4],
template matching [5], or an object-specific pose estimator [6] trained on
manually labeled data.

Deep reinforcement learning (RL), on the other hand, has shown promising
outcomes in learning to control a robot for complex manipulation tasks such as
grasping [7, 8] and insertion [9]. However, prior works often limit operation
to simulated environments [10] or heavily restrict and regulate operating
spaces by executing with a short horizon without obstructions [9, 8] due to
the sample inefficiency and potential of executing unsafe policies.

Figure 1:  Task setup. We solve complex manipulation tasks within the entire
operational space of a robot by leveraging an OCGM for versatile and efficient
goal acquisition paired with MP and RL. Note that obstacles and a socket are
randomly placed on the table.

Combining MP and RL has been investigated by several prior works [11, 6] and
shows the potential of leveraging the strengths of both methods to solve
manipulation tasks in obstructed environments. Yet, goal specification for MP
in prior work has relied on either sample-inefficient interaction with the
environment or an object-specific pose estimator, which needs re-training for
each new target object.

In this work, we take advantage of both MP and RL to achieve efficient skill
acquisition for complex manipulation tasks in environments with arbitrary
obstacles. We propose a system that leverages an object-centric generative
model (OCGM) [12] for versatile, one-shot goal identification and re-
identification as a key component in solving complex manipulation tasks from a
few simple human demonstrations. Specifically, we identify a target object
from a single demonstration using an OCGM, pre-trained on diverse synthetic
scenes, leading to robust re-identification of that object in new scenes by
matching to its object-centric representation. Using the object’s position as
a goal, the motion planner generates a collision-free path to the target
object while avoiding obstacles before a learned RL policy is executed to
complete the complex manipulation tasks. We train an RL policy for each
manipulation skill from a sparse reward to preclude specialist knowledge for
reward engineering and guide exploration using a handful of easy-to-collect
demonstrations. To maximise performance, we also introduce a skill transition
network to reduce failures that occur when transitioning from MP to the
learned RL policy.

The contributions of our work are fourfold: (1) we propose a system for
efficient skill acquisition in obstructed environments that leverages an OCGM
for object-agnostic, one-shot goal specification, (2) we introduce a
transition network that smoothly interpolates between terminal states of
motion planning and feasible start states of a learned RL policy, (3) we show
that our OCGM-based one-shot goal specification method achieves comparable
accuracy against several goal identification baselines, and (4) we demonstrate
that our system performs significantly better in real-world environments
compared to baselines, including a state-of-the-art RL algorithm.

##  II Related Works

Recent success in deep RL [7, 10, 13, 14] enables a robot to learn complex
manipulation tasks such as grasping [7, 8] and insertion [9, 15, 16, 17, 18]
driven by a reward. To avoid the requirement of specialist knowledge for
reward engineering, several prior works have proposed sample-efficient RL
methods that can learn complex manipulation skills from a sparse reward by
leveraging a small number of demonstrations for guided exploration [8, 9, 19,
15]. However, due to the sample inefficiency of sparse rewards, studies have
been primarily conducted in simulated environments or within limited task
spaces in the real world. Learning from demonstration (LfD) [20, 21, 22] is an
alternative method for a robot to learn manipulation tasks by imitating
behaviour in expert demonstrations collected by a human operator, but it often
requires a large number of demonstrations to acquire manipulation skills.
While InsertionNet [23] enables a robot to solve industrial insertion tasks
within the entire operational space of a robot manipulator from a small number
of demonstrations, it is evaluated in a clean environment without obstruction.
Successful insertion is also made possible by a small initiation set for the
learnt skill. Adaptive LfD for insertion has also been proposed [24], allowing
a policy to quickly adapt to a new industrial object from the same category
seen in training using only a single demonstration and the object mesh.
However, such mesh information is not readily available, limiting real-world
applications. In our work, a manipulation skill is learnt using Framework for
Efficient Robot Learning (FERM)[8] but is essentially interchangeable with any
method capable of learning the needed skill efficiently.

Motion planning (MP) [1, 25, 2, 26, 3] can effectively generate a collision-
free path from a robot’s initial configuration to a goal pose using an
explicit model of the robot and environment. However, such a goal pose is
often specified by a user or object-specific pose estimator to generate a
path. Further, MP does not model the dynamics of the surrounding environment
or objects and therefore complex manipulation tasks are out of the scope for
MP.

Several previous works [11, 6, 27, 28, 29] combine MP and RL to leverage the
benefits of both methods to solve manipulation tasks in unstructured
environments. However, these preceding works limit their real-world
applicability by requiring a large number of samples to learn a goal estimator
[11] or by retraining an object-specific predictor for each new goal object
[6].

Our work leverages OCGMs [12, 30, 31] to find a target object for MP, negating
the need for object-specific goal estimators. OCGMs learn structured
representations of objects within complex scenes and provide object-specific
encodings useful for matching, however, most OCGMs have never been applied to
real world environments. In contrast to goal specification methods that
require human intervention or a large, object-specific datasets, i.e. template
matching or object classifier, OCGMs hold the promise of versatile target
object identification. This work leverages APEX [12] by training the model in
an unsupervised manner on a wide distribution of simulated data to assist with
generalisation to real-world environments.

##  III Methodology

Figure 2: Our framework architecture. (a) We leverage an OCGM to re-identify a
target object such that its object-centric representation matches one
extracted from a single demonstration. The goal state is specified in the
robot’s reference frame using an external RGB-D camera with calibrated
extrinsics. (b) Given the goal state acquired in (a), a motion planner
generates a collision-free path to the goal. (c) A skill transition network
guides the arm from the terminal state of the motion planning (MP) to the
initiation set of the RL policy. (d) Given a wrist camera image
$\boldsymbol{I}^{wrist}_{t}$ and robot’s internal state $\boldsymbol{x}_{t}$,
a learned RL policy executes the final interaction until task completion.

In this work, we present an efficient solution for solving complex
manipulation tasks in obstructed, real-world environments by leveraging an
OCGM. We demonstrate our method on several industrial insertion tasks as they
require learning complex insertion skills and also require the identification
of the target socket to complete the tasks (see Figure 1 for our task setup).
We break our method down into pre-training, in-situ skill training and
execution in the following subsections. The pre-training component of our
method is only completed once and can be reused for all future insertion
tasks. In-situ skill training is required for each new insertion task and
execution describes the process for autonomous task execution after training.

###  III-A Pre-training

Pre-training is required for APEX [12], our choice of unsupervised OCGM, to
achieve versatile one-shot target object acquisition, but it only needs to be
done once as it is trained on a diverse dataset to encourage generalise to a
variety of real-world objects. APEX is formulated as a set of VAEs, and takes
a video sequence $\boldsymbol{I}_{1:T}$ as input. Each frame is decomposed
into a set of latent representations for each discovered object $j$ consisting
of object location $\boldsymbol{z}^{where}_{t,j}$, appearance
$\boldsymbol{z}^{what}_{t,j}$, and presence $z^{pre}_{t,j}\in[0,1]$, where $T$
is the number of frames in a video sequence. We train APEX on a synthetic
dataset consisting of a set of trajectories in which a robot interacts with a
diverse set of primitive shapes of differing colour and size. Domain
randomisation [32] is a common technique for transferring models trained on
synthetic data to the real world, however, APEX is unable to handle drastic
background changes present in domain randomised datasets because a changing
background is hard to distinguish from a moving foreground object. Instead, we
add a small amount of noise to the camera pose for each trajectory, leading to
variations in the images. As a result, APEX is successfully applied to real-
world scenes with similar background textures. See Appendix A-A for further
information regarding APEX.

###  III-B In-Situ Skill Training

This section details the task-specific data and training required by our
method. The data must be collected for each new task that the robot is taught,
however, the supervised component only requires about 10 minutes to collect.
First, a single demonstration,
${\mathcal{D}^{goal}=\\{(\boldsymbol{I}^{ext}_{t},\boldsymbol{x}^{ee}_{t}),\dots\\}}$
consisting of a sequence of images $\boldsymbol{I}^{ext}$ from the third-
person camera and robot end-effector positions $\boldsymbol{x}^{ee}_{t}$, of a
successful task completion from anywhere within the robot’s operational space
is collected for goal specification using the pre-trained APEX. Additionally,
25 demonstrations,
$D^{RL}=\\{(\boldsymbol{I}^{wrist}_{t},\boldsymbol{x}_{t},\boldsymbol{a}_{t}),\dots\\}^{|25|}$,
of the complex interaction with objects for efficient RL training, are
collected from within a limited task space such that the connector is always
within sight of the wrist camera, where $\boldsymbol{I}^{wrist}_{t}$,
$\boldsymbol{x}_{t}$, $\boldsymbol{a}_{t}$ are a wrist camera image, robot
states, and action at time step $t$.

We train the RL policy with FERM [8] to complete the insertion task, which
takes between $60$ to $90$ minutes to train on a desktop computer with an i7
processor and a Nvidia Titan X GPU. FERM is composed of Soft Actor Critic
[10], contrastive learning [33, 34], and image augmentation [35]. In addition
to a gray-scale image from the wrist-mounted camera
($\boldsymbol{I}^{wrist[H\times W]}$ where $H$ and $W$ are 64 pixels), the
policy takes as input the end-effector Cartesian velocity and F/T sensor data
(see Figure 2 (d)) and outputs the desired 3-dimensional Cartesian end-
effector velocity for the robot. Because the policy takes as input local
information, it is able to generalise to any location in the robot’s
operational space. The RL policy is trained using a sparse reward
$r_{t}=\mathbb{1}[\boldsymbol{s}\in S_{g}]$ where $S_{g}$ is a set of goal
states defined as a goal pose of the the robot with $1cm$ tolerance determined
by simply teleoperating the robot to the goal pose before training the policy.
To accelerate training of the RL policy, we leverage the task demonstrations
$D^{RL}$ to initialise a replay buffer for guided exploration similar to FERM
and train the policy and critic asynchronously, as inspired by prior work [9].
We also limit the task space for this stage, such that the socket is always
within sight of the wrist camera for training, improving sample efficiency and
reducing the chance of unsafe interactions. For further details, see Appendix
B-A. While we adopt FERM to train a low-level manipulation policy, it is
interchangeable with any approach, including imitation learning [20, 21], that
can learn complex manipulation skills efficiently from a small number of
demonstrations. Note, however, that the learned policy often requires a large
coverage of states around a target object such as a socket for insertion tasks
to increase the success rate.

We also introduce a skill transition network, inspired by prior work [36], for
each insertion task to improve the task success rate. The terminal state of
the MP is not guaranteed to be within the feasible start states of the RL
policy, defined as _the initiation set of the skill_ (represented as a blue
box in Figure 2 (c)), due to the accumulation of errors in camera extrinsics
and estimation of the MP goal state 3D. To mitigate this issue, a simple
convolutional neural network (CNN) is trained on data collected in a self-
supervised manner to predict the Cartesian offset required to move the end-
effector from the terminal states of MP to the initiation set of the skill.
The dataset is collected in less than 30 minutes by sampling random Cartesian
poses around the target object and recording the sequences of wrist camera
images $\boldsymbol{I}^{wrist}_{t}$ and offsets between the current end-
effector pose and an initial pose used for RL training. The collected data
contains only local information conditioned on the wrist camera which allows
the transition network to generalise to unseen target positions. For further
details, see Appendix B-B.

###  III-C Execution

The execution of the task can be completed from _anywhere_ within the robot’s
operational space to any goal location. Execution follows four steps (see
Figure 2) that are completed autonomously: (i) goal identification via OCGM,
(ii) MP, (iii) skill transition network, and (iv) RL policy.

Figure 3: Target object identification and re-identification using an OCGM.
(a) We leverage a pre-trained OCGM to extract an object-centric representation
from a single task demonstration. The target object is identified in the
demonstration as the object mask closest to the robot end-effector position at
the end of the trajectory $\boldsymbol{x}^{ee}_{T}$ (see Eq 1). (b) Given a
new scene, the OCGM is used to acquire object-centric representations of all
objects and compare these with the already identified target object
representation to re-identify the target object.

#### Goal Identification via OCGM

As MP requires a goal pose to plan a collision-free path through the scene, we
leverage the pretrained OCGM to identify the target object from the single
demonstration $D^{goal}$ and re-identify it in the current scene to specify
the goal for MP. To identify the target object from $D^{goal}$, we first
acquire a set of object-centric representations by encoding the first external
camera image $\boldsymbol{I}^{ext}_{0}$ in the demonstration $D^{goal}$ (see
Figure 3 (a)). We determine the target object-centric representation
$\boldsymbol{z}^{what}_{target}$ such that the object is present at the
beginning of the trajectory, i.e. $p(z^{pre}_{0,j})\geq 0.5$ and such that it
is the closest to the robot end-effector position $\boldsymbol{x}^{ee}_{T}$ at
the end of the demonstration $D^{goal}$. To calculate the 3D poses of objects
in the robot’s reference frame, the centre of the object mask predicted by the
OCGM is converted to Cartesian coordinates using the RGB-D camera’s depth
plane and the known camera extrinsics. The closest object to the robot end-
effector position at the end of the demonstration $D^{goal}$ is calculated
using $L_{2}$ distance:

$target=\operatorname*{arg\,min}_{j,j=1..N}||\boldsymbol{x}^{ee}_{T}-\boldsymbol{o}_{j}||_{2}$
(1)

where $o_{j}$ is the 3D object positions in the robot reference frame and $N$
is the number of objects discovered by the OCGM in the scene
$\boldsymbol{I}^{ext}_{0}$. In order to re-identify the target object in the
current scene (see Figure 3 (b)), we compare the target object-centric
representation $\boldsymbol{z}^{what}_{target}$ with each object-centric
representation $\boldsymbol{\hat{z}}^{what}_{j}$ discovered in the new scene
(see Figure 3 (b)) using the $L_{2}$ distance and choose the object that has
the most similar representation:

$\hat{\boldsymbol{z}}^{what}_{target}=\operatorname*{arg\,min}_{j=1..N}||\boldsymbol{z}^{what}_{target}-\hat{\boldsymbol{z}}^{what}_{j}||_{2}$
(2)

#### MP + Transition Policy + RL Policy

Using the target object’s pose $\boldsymbol{o}_{target}$ in the robot’s
reference frame, we use an RRT-connect motion planner to guide the robot’s
end-effector to the location of the target object (see Figure 2 (b)). To avoid
collisions during the MP phase, an occupancy map, OctoMap [37], is created
using the point clouds captured by the calibrated external camera. After the
execution of MP, we leverage the trained skill transition network to guide the
arm into the initiation set of the skill (see Figure 2(c)) to maximise the
outcomes of the RL policy. Finally, the learned RL policy completes the
manipulation task.

##  IV Experiments

Our experiments are designed to answer the following guiding questions: (1)
does the use of an OCGM achieve versatile and efficient target object
identification for MP in the real world? (2) how well does our system perform
complex manipulation tasks, specifically industrial insertion tasks, in
obstructed environments? (3) does a skill transition network increase task
success rate?

Figure 4: Insertion tasks. We evaluate our framework on four assembly tasks
found in industry. Each socket is attached to a mount of varying size and
colour to demonstrate the versatility and efficiency of our one-shot goal
specification using an OCGM.

| VGA | RJ45 | E-model | USB-A | #Data | Intervention
---|---|---|---|---|---|---
Method | Accuracy | WSI | Accuracy | WSI | Accuracy | WSI | Accuracy | WSI |  |
Template matching | 70.0% | 54.7/81.9% | 35.0% | 22.1/50.5% | 87.5% | 73.9/94.5% | 55.0% | 39.8/69.3% | 1 | yes
Feature-based matching | 40.0% | 26.3/55.4% | 87.5% | 73.9/94.5% | 7.5% | 2.6/19.9% | 77.5% | 62.5/87.7% | 1 | yes
Object-specific classifier | 80.0% | 65.2/89.5% | 100.0% | 91.2/100.0% | 87.5% | 73.9/94.5% | 75.0% | 59.8/85.8% | 2.5K | yes
OCGM identifier (Ours) | 82.5% | 68.0/91.3% | 95.0% | 83.5/98.6% | 95.0% | 83.5/98.6% | 92.5% | 80.1/97.4% | 1 | no

TABLE I: Accuracy of target object identification. We evaluate our method and
two baselines on 160 test scenes (40 scenes per connector), and report the
accuracy, Lower Limit (LL) and Upper Limit (UL) of the Wilson score interval
(WSI) with confidence interval of $95\%$. While template matching and object-
specific classification requires human intervention, such as cropping a
reference target object image and labelling training data, our OCGM identifier
successfully identifies the target object from only a single demonstration
without such human intervention.

###  IV-A Experimental Setup

Several industrial insertion tasks, inspired by the NIST assembly boards [38],
are used within our experiments (see Figure 4). To verify the robustness of
the target object identification using an OCGM, mounts for each socket and
arbitrary obstacles with different colour and size are used. In our
experiments, we use a Franka Panda robot (7-DOF robot arm) and rigidly attach
each connector to the robot’s end-effector. Each phase of our framework uses
different controllers: a joint position controller to follow a trajectory
planned by the motion planner, a Cartesian pose controller for the skill
transition network, and a Cartesian velocity impedance controller for the RL
policy. For each evaluation trial, the socket, robot arm, and one or two
obstacles are randomly placed in the robot’s operational space. While the
rotation around the z-axis of the sockets is consistent across all of the
evaluation trials, we can readily extend our system to accommodate such cases
by training a transition network that also predicts the displacement of the
z-axis orientation to move the end-effector to the initiation set of a low-
level manipulation policy. We leave this extension for future work.

Given the pre-trained OCGM, for each new manipulation skill, our modular
framework requires a total of $10$ minutes of human-supervised demonstrations
and a maximum of $130$ minutes of unsupervised training comprising of: up to
$90$ minutes for RL policy training and $40$ minutes for data collection and
training of the skill transition network.

###  IV-B Efficient and Versatile Target Object Identification

First, the OCGM identifier is evaluated against several baselines on 160 test
scenes (40 for each of the four sockets) with the target locations hand-
labelled with bounding boxes for quantitative comparison. During testing, if
the intersection of union (IoU) between a ground truth and the returned
bounding box from the tested algorithm is greater than $0.5$, we count it as
successful [39].

We evaluate our proposed goal identification approach against three baselines.
Template matching finds the target object in the current scene by calculating
a correlation coefficient given a manually cropped target object reference
image. Feature-based matching finds a pair of the best matched keypoints
between the manually cropped target object reference image and the current
scene using FLANN-based matching [40] and SIFT descriptor [41]. Object-
specific classifier trained on a dataset of manually cropped object images
with binary labels, is queried with cropped images found using a region
proposal method [42]. Lastly, we evaluate our method by retrieving the minimum
bounding box of the target object mask predicted by the OCGM.

#### Results.

We report the accuracy of target object identification in Table I. Our method
achieves commensurate or better performance compared to other baselines,
whilst not requiring human intervention or an object-specific dataset which
requires laborious manual data labelling. This result motivates the use of
OCGMs for efficient goal acquisition for MP. Template matching often struggles
to find a target object with high confidence, potentially due to the
projection of a slanted camera and small objects. Feature-based matching also
shows a lower success rate for several objects due to a lack of distinguishing
features, especially for small objects in a scene. Object-specific classifier,
on the other hand, generally performs well because it is tailored to a single
object, and can be further improved by collecting more data. However, such
classifiers requires re-training on a new dataset manually labelled for each
new object, limiting the versatility and efficiency of goal specification.

###  IV-C Insertion Tasks in Obstructed Environments

We evaluate our proposed system on several industrial insertion tasks in
obstructed environments against a series of baselines composed of competing
methods. All baselines that utilise MP make use of the OCGM for target object
identification. For each task, we conduct 30 trials and report the success
rate in Table II. Figure 5 illustrates the execution of our method for each
insertion task in the obstructed environments.

We compare the performance of our approach against a state-of-the-art RL
algorithm and four comparable instantiations of our approach. _Soft Actor-
Critic (SAC)_ a state-of-the-art RL algorithm that predicts the desired
Cartesian velocity from sparse rewards, trained with 25 demonstrations using
FERM [8] (see Appendix B-C for training details). _MP+Demonstration Replay_
substitutes replaying a single expert demonstration for the learned RL policy
execution in our method, inspired by previous work [36]. _MP+BC_ replaces the
learned RL policy in our method with Behaviour Cloning (BC) [43, 44], trained
from $25$ demonstrations. _MP+Heuristic_ uses a manually designed heuristic
policy [9] instead of the learned RL policy in our method to solve the task
(see Appendix B-D for details of the heuristic policy). Lastly, we evaluate
our method without a skill transition network (_MP+RL w/o skill transition_)
for comparison.

#### Results.

As described in Table II, our method (MP+RL) as outlined in Section III
records the highest success rate for all tasks. The results for the SAC
baseline show that it is unable to solve any of the tasks, likely because it
requires a large number of samples to train the policies in the robot’s
operational space with obstructions. MP+Demonstration Replay is the most data-
efficient method, however, it mostly fails to solve any of the tasks because
it requires very accurate estimation of pose offsets for the demonstration
replay to be successful. MP+BC is another efficient skill acquisition method
because it does not require any additional interactions with the environments
other than the given demonstrations to learn manipulation skills. However, due
to the narrow state coverage, it struggles to solve the tasks. While
MP+Heuristic is able to solve some insertion tasks, such as USB-A and E-model,
almost one-third of the time, it fails to solve the tasks the majority of the
time due to the need for accurate pose offset (the same reason for failure as
MP+Demonstration Replay). While our method achieves high success rate over 4
industrial insertion tasks, the main failure case is caused by the
misidentification of the target object by the OCGM. These failure modes can be
readily eliminated by extended and/or augmenting the OCGM training.

Figure 5: Real-world industrial assembly tasks in obstructed environments. The
OCGM is used to specify a goal for MP, followed by a skill transition network
and a learned RL policy. Our method successfully solves complex manipulation
tasks with a high success rate.

| VGA | RJ45 | E-model | USB-A
---|---|---|---|---
Method | Success | WSI | Success | WSI | Success | WSI | Success | WSI
SAC | 0.0% | 0.0/16.1% | 0.0% | 0.0/16.1% | 0.0% | 0.0/16.1% | 0.0% | 0.0/16.1%
MP + Demonstration Replay | 3.3% | 1.0/16.7% | 0.0% | 0.0%/16.1% | 3.3% | 1.0/16.7% | 0.0% | 0.0/16.1%
MP + BC | 16.7% | 7.3%/33.6% | 16.7 | 7.3/33.6% | 23.3% | 11.8/40.9% | 26.7% | 14.2/44.5%
MP + Heuristic | 10.0% | 3.5/25.6% | 16.7 | 7.3/33.6% | 36.7% | 21.9/54.5% | 43.3% | 27.4/60.8%
MP + RL w/o skill transition | 73.3% | 55.6/85.8% | 46.7% | 30.2/63.8% | 80.0% | 62.7/90.5% | 70.0% | 52.1/83.3%
MP + RL (our method) | 86.7% | 70.3/94.7% | 83.3% | 66.4/92.7% | 93.3% | 78.7/98.2% | 96.7% | 83.3/99.4%

TABLE II: Real-world assembly results. We report the success rate, Lower Limit
(LL) and Upper Limit (UL) of the Wilson score interval (WSI) with confidence
interval of $95\%$ over 30 trials. Our method outperforms, by a significant
margin, all of the other methods including a state-of-the-art RL method and
several comparable instantiations of our method.

Examining whether the transition network is required to solve complex
manipulation tasks in obstructed environments (see Table II), the results
verify that using the skill transition network results in higher success rates
than without the skill transition policy. Due to errors caused by the OCGM,
camera extrinsics and estimation of the 3D goal poses, a terminal state of MP
can often be outside of the initiation set of the learned RL skill. Therefore,
by introducing the skill transition module to move the robot arm into the
initiation set of the skill, we can mitigate these issues and achieve better
performance.

##  V Conclusion

In this work, we propose a modular system that leverages an OCGM for one-shot
goal identification and re-identification as a vital component to combine MP
and RL to solve complex manipulation tasks in obstructed environments.
Specifically, the OCGM extracts a target object from only a single
demonstration and re-identifies the object to determine a goal pose for MP
without the need of fine-tuning on an object-specific dataset. The
experimental results show that our method for goal specification using an OCGM
achieves better versatility and comparable accuracy to other tested baselines.
In addition, our method successfully solves real-world industrial insertion
tasks in obstructed environments from few demonstrations. In future work we
will look to investigate more advanced settings, such as randomising the
socket orientation.

## ACKNOWLEDGMENT

This work was supported by a UKRI/EPSRC Programme Grant [EP/V000748/1], we
would also like to thank the University of Oxford for providing Advanced
Research Computing (ARC) (http://dx.doi.org/10.5281/zenodo.22558) and the SCAN
facility in carrying out this work.

## References

  * [1] N. M. Amato and Y. Wu, “A randomized roadmap method for path and manipulation planning,” in _IEEE International Conference on Robotics and Automation_ , 1996.
  * [2] S. M. LaValle, “Rapidly-exploring random trees: A new tool for path planning,” Computer Science Department, Iowa State University, Tech. Rep. TR 98-11, 1998.
  * [3] S. Karaman and E. Frazzoli, “Sampling-based algorithms for optimal motion planning,” _International Journal of Robotics Research_ , vol. 30, no. 7, pp. 846–894, 2011.
  * [4] M. Khodeir, B. Agro, and F. Shkurti, “Learning to search in task and motion planning with streams,” 2021.
  * [5] M.-T. Le, C.-H. G. Li, S.-M. Guo, and J.-J. J. Lien, “Embedded-based object matching and robot arm control,” in _2019 IEEE 15th International Conference on Automation Science and Engineering (CASE)_ , 2019, pp. 1296–1301.
  * [6] M. A. Lee, C. Florensa, J. Tremblay, N. Ratliff, A. Garg, F. Ramos, and D. Fox, “Guided uncertainty-aware policy optimization: Combining learning and model-based strategies for sample-efficient policy learning,” _IEEE International Conference on Robotics and Automation_ , 2020.
  * [7] D. Kalashnikov, A. Irpan, P. Pastor, J. Ibarz, A. Herzog, E. Jang, D. Quillen, E. Holly, M. Kalakrishnan, V. Vanhoucke _et al._ , “Scalable deep reinforcement learning for vision-based robotic manipulation,” in _Conference on Robot Learning_ , 2018, pp. 651–673.
  * [8] A. Zhan, P. Zhao, L. Pinto, P. Abbeel, and M. Laskin, “A framework for efficient robotic manipulation,” _arXiv preprint arXiv:2012.07975_ , 2020\.
  * [9] J. Luo, O. Sushkov, R. Pevceviciute, W. Lian, C. Su, M. Vecerik, N. Ye, S. Schaal, and J. Scholz, “Robust multi-modal policies for industrial assembly via reinforcement learning and demonstrations: A large-scale study,” _arXiv preprint arXiv:2103.11512_ , 2021.
  * [10] T. Haarnoja, A. Zhou, P. Abbeel, and S. Levine, “Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor,” in _International Conference on Machine Learning_ , 2018.
  * [11] J. Yamada, Y. Lee, G. Salhotra, K. Pertsch, M. Pflueger, G. S. Sukhatme, J. J. Lim, and P. Englert, “Motion planner augmented reinforcement learning for obstructed environments,” in _Conference on Robot Learning_ , 2020.
  * [12] Y. Wu, O. P. Jones, M. Engelcke, and I. Posner, “Apex: Unsupervised, object-centric scene segmentation and tracking for robot manipulation,” in _2021 IEEE/RSJ International Conference on Intelligent Robots and Systems (IROS)_. IEEE, 2021, pp. 3375–3382.
  * [13] S. Gu, E. Holly, T. Lillicrap, and S. Levine, “Deep reinforcement learning for robotic manipulation with asynchronous off-policy updates,” in _IEEE International Conference on Robotics and Automation_ , 2017.
  * [14] S. Levine, C. Finn, T. Darrell, and P. Abbeel, “End-to-end training of deep visuomotor policies,” _Journal of Machine Learning Research_ , 2016.
  * [15] M. Vecerik, O. Sushkov, D. Barker, T. Rothörl, T. Hester, and J. Scholz, “A practical approach to insertion with variable socket position using deep reinforcement learning,” 2018.
  * [16] M. A. Lee, Y. Zhu, K. Srinivasan, P. Shah, S. Savarese, L. Fei-Fei, A. Garg, and J. Bohg, “Making sense of vision and touch: Self-supervised learning of multimodal representations for contact-rich tasks,” 2018.
  * [17] T. Davchev, K. S. Luck, M. Burke, F. Meier, S. Schaal, and S. Ramamoorthy, “Residual learning from demonstration: Adapting DMPs for contact-rich manipulation,” _IEEE Robotics and Automation Letters_ , 2022.
  * [18] J. Carvalho, D. Koert, M. Daniv, and J. Peters, “Residual robot learning for object-centric probabilistic movement primitives,” 2022.
  * [19] M. Vecerik, T. Hester, J. Scholz, F. Wang, O. Pietquin, B. Piot, N. Heess, T. Rothörl, T. Lampe, and M. Riedmiller, “Leveraging demonstrations for deep reinforcement learning on robotics problems with sparse rewards,” 2017.
  * [20] S. Schaal, “Is imitation learning the route to humanoid robots?” _Trends in Cognitive Sciences_ , vol. 3, no. 6, pp. 233–242, 1999.
  * [21] A. Billard, S. Calinon, R. Dillmann, and S. Schaal, “Survey: Robot programming by demonstration,” _Springer Handbook of Robotics_ , pp. 1371–1394, 2008\.
  * [22] O. Groth, C.-M. Hung, A. Vedaldi, and I. Posner, “Goal-conditioned end-to-end visuomotor control for versatile skill primitives,” in _2021 IEEE International Conference on Robotics and Automation (ICRA)_ , 2021, pp. 1319–1325.
  * [23] O. Spector and D. Di Castro, “Insertionnet – a scalable solution for insertion,” 2021.
  * [24] B. Wen, W. Lian, K. Bekris, and S. Schaal, “You only demonstrate once: Category-level manipulation from single visual demonstration,” 2022.
  * [25] L. Kavraki and J.-C. Latombe, “Randomized preprocessing of configuration for fast path planning,” in _IEEE International Conference on Robotics and Automation_ , 1994.
  * [26] S. M. LaValle, J. J. Kuffner, B. Donald _et al._ , _Algorithmic and computational robotics: new directions_. AK Peters, 2001, ch. Rapidly-exploring random trees: Progress and prospects, pp. 293–308.
  * [27] F. Xia, C. Li, R. Martín-Martín, O. Litany, A. Toshev, and S. Savarese, “Relmogen: Leveraging motion generation in reinforcement learning for mobile manipulation,” _arXiv preprint arXiv:2008.07792_ , 2020\.
  * [28] C.-Y. Kuo, A. Schaarschmidt, Y. Cui, T. Asfour, and T. Matsubara, “Uncertainty-aware contact-safe model-based reinforcement learning,” _IEEE Robotics and Automation Letters_ , vol. 6, no. 2, pp. 3918–3925, 2021\.
  * [29] J. Scholz and M. Stilman, “Combining motion planning and optimization for flexible robot manipulation,” in _2010 10th IEEE-RAS International Conference on Humanoid Robots_ , 2010, pp. 80–85.
  * [30] F. Locatello, D. Weissenborn, T. Unterthiner, A. Mahendran, G. Heigold, J. Uszkoreit, A. Dosovitskiy, and T. Kipf, “Object-centric learning with slot attention,” in _Advances in Neural Information Processing Systems_ , vol. 33, 2020.
  * [31] Z. Lin, Y.-F. Wu, S. V. Peri, W. Sun, G. Singh, F. Deng, J. Jiang, and S. Ahn, “Space: Unsupervised object-oriented scene representation via spatial attention and decomposition,” in _International Conference on Learning Representations_ , 2020.
  * [32] J. Tobin, R. Fong, A. Ray, J. Schneider, W. Zaremba, and P. Abbeel, “Domain randomization for transferring deep neural networks from simulation to the real world,” in _2017 IEEE/RSJ international conference on intelligent robots and systems (IROS)_. IEEE, 2017, pp. 23–30.
  * [33] Z. Wu, Y. Xiong, S. X. Yu, and D. Lin, “Unsupervised feature learning via non-parametric instance discrimination,” in _2018 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)_. IEEE Computer Society, 2018, pp. 3733–3742.
  * [34] M. Laskin, A. Srinivas, and P. Abbeel, “Curl: Contrastive unsupervised representations for reinforcement learning,” _Proceedings of the 37th International Conference on Machine Learning, Vienna, Austria, PMLR 119_ , 2020, arXiv:2004.04136.
  * [35] M. Laskin, K. Lee, A. Stooke, L. Pinto, P. Abbeel, and A. Srinivas, “Reinforcement learning with augmented data,” _arXiv preprint arXiv:2004.14990_ , 2020.
  * [36] E. Johns, “Coarse-to-fine imitation learning: Robot manipulation from a single demonstration,” in _IEEE International Conference on Robotics and Automation (ICRA)_ , 2021.
  * [37] A. Hornung, K. M. Wurm, M. Bennewitz, C. Stachniss, and W. Burgard, “OctoMap: An efficient probabilistic 3D mapping framework based on octrees,” _Autonomous Robots_ , 2013.
  * [38] K. Kimble, J. Falco, E. Messina, K. V. Wyk, Y. Sun, M. Shibata, W. Uemura, and Y. Yokokohji, “Benchmarking protocols for evaluating small parts robotic assembly systems,” no. 5, 2020.
  * [39] M. Everingham, S. M. A. Eslami, L. Van Gool, C. K. I. Williams, J. Winn, and A. Zisserman, “The pascal visual object classes challenge: A retrospective,” _International Journal of Computer Vision_ , vol. 111, no. 1, pp. 98–136, Jan. 2015.
  * [40] M. Muja and D. G. Lowe, “Fast approximate nearest neighbors with automatic algorithm configuration.” in _VISAPP (1)_. INSTICC Press, 2009, pp. 331–340.
  * [41] D. G. Lowe, “Distinctive image features from scale-invariant keypoints,” _Int. J. Comput. Vision_ , vol. 60, no. 2, pp. 91–110, Nov. 2004.
  * [42] J. Uijlings, K. van de Sande, T. Gevers, and A. Smeulders, “Selective search for object recognition,” _International Journal of Computer Vision_ , 2013\.
  * [43] T. Zhang, Z. McCarthy, O. Jow, D. Lee, X. Chen, K. Goldberg, and P. Abbeel, “Deep imitation learning for complex manipulation tasks from virtual reality teleoperation,” in _IEEE International Conference on Robotics and Automation_ , 2018, pp. 5628–5635.
  * [44] M. Bojarski, D. D. Testa, D. Dworakowski, B. Firner, B. Flepp, P. Goyal, L. D. Jackel, M. Monfort, U. Muller, J. Zhang, X. Zhang, J. Zhao, and K. Zieba, “End to end learning for self-driving cars,” _arXiv preprint arXiv:1604.07316_ , 2016.
  * [45] S. James, Z. Ma, D. R. Arrojo, and A. J. Davison, “Rlbench: The robot learning benchmark & learning environment,” _IEEE Robotics and Automation Letters_ , vol. 5, no. 2, pp. 3019–3026, 2020.

##  Appendix A Pre-training Details

###  A-A APEX Training Details and Hyperparameters

We collect a synthetic dataset consisting of $30K$ trajectories each of which
has $10$ to $15$ frames simulated using RLBench [45]. In each trajectory, a
robot randomly pushes objects on a table. To apply the pre-trained APEX to the
real world, we add a small amount of noise, ranging from $1$cm to $-1$cm, to
the Cartesion camera position in the simulation so that the background texture
changes slightly for each trajectory and prevents APEX from overfitting to the
specific background texture. Table III shows hyperparameters of APEX. For
further details, see prior work [12].

TABLE III: APEX hyperparameter Parameter | Value
---|---
Observation Rendering | (128, 128), RGB
Optimizer | Adam
Learning rate | 1e-4
horizon | 10
Image size | 128
Foreground std | 0.11
Background std | 0.04
KL divergence for $z_{what}$ | 3e-4
KL divergence for $z_{where}$ | 15
KL divergence for $z_{pre}$ (discovery) | 32
KL divergence for $z_{pre}$ (tracking) | 1
Background reconstruction loss weight | 10
foreground reconstruction loss weight | 1

##  Appendix B In-Situ Training Details

###  B-A RL Training Details and Hyperparameters

We define a sparse reward function such that a policy receives $1$ when the L2
distance between robot’s end-effector pose and the pre-defined target pose is
within a small tolerance of $0.8$cm for VGA and $1$cm for the other insertion
tasks. The policy and critic take as input a wrist camera image, Cartesian
velocity, and F/T sensor values of the end-effector. The critic also takes as
input an action sampled from the policy. The architecture of the policy and
the critic includes a CNN for processing the wrist camera image followed by a
concatenation with state information. The concatenated features are fed to the
policy and critic feedforward neural networks with the outputs being actions
and values respectively. Similar to FERM [8], we apply random crops and
brightness changes to the image observations to acquire a robust policy
efficiently. The policy runs at $10$Hz and we find that this is sufficient for
completing the complex insertion tasks. Table IV shows hyperparameters.

TABLE IV: FERM hyperparameter Parameter | Value
---|---
Optimizer | Adam
Learning rate | 1e-3
Discount factor ($\gamma$) | 0.99
Replay buffer size | $10^{4}$
Latent dimension | 50
Convolution filters | $[8,16,32,64]$
Convolution strides | $[2,2,2,2]$
Convolution filter size | 3
Hidden Units (MLP) | [128]
Nonlinearity | Leaky ReLU
Target smoothing coefficient ($\tau$) | 0.005
Target update interval | 2
Actor update interval | 2
Network update per environment step | 10

###  B-B Transition Network Details and Hyperparameters

The dataset for the transition network is collected by uniformally sampling a
Cartesian offset from within $\pm 5$cm of the initial position used for RL
training. The robot moves linearly between the initial position used for RL
training and the sampled offset, recording both the wrist camera image and the
current offset at $15$ steps throughout. After collecting $100$ of these
trajectories, in less than $30$ minutes, the skill transition network can be
trained.

###  B-C Baseline Details - SAC

To train a policy with SAC [10] from scratch using FERM [8], we randomly place
a socket and one or two obstacles from our set of obstacles (see Figure 6) on
a table. Then, we define a reward function for collision avoidance in addition
to a sparse reward for solving the insertion task within the entire
operational space of the robot:

$R=I[s\in S_{g}]-0.005\cdot I[s\in S_{coll}]$ (3)

where $S_{g}$ and $S_{coll}$ are a set of goal states and states involving
collisions and $I$ is an indicator function. Furthermore, the policy also
takes as input an image from the external camera in addition to a wrist camera
image. We use the same hyperparameters described in Table IV.

Figure 6:  The set of obstacles used in our experiments. In the evaluation, we
randomly sample one or two obstacles from this set of obstacles that vary in
size and colour.

###  B-D Baseline Details - MP + Heuristic

We design a heuristic policy for insertion tasks, inspired by previous work
[9] that utilises the F/T snesor data. The heuristic policy first moves the
end-effector downwards until the arm contacts an object. Then, the arm follows
an outwards spiral pattern on a 2D grid to find the socket hole. Finally,
after the connector is aligned with the hole, the arm moves downwards again
with a small circular movement in the X-Y plane to insert the connector into
the socket.


11institutetext:  CNRS & Sorbonne Université, Institut d’Astrophysique de
Paris (IAP), UMR 7095, 98 bis bd Arago, F-75014 Paris, France22institutetext:
Heuristic and Evolutionary Algorithms Laboratory, University of Applied
Sciences Upper Austria, Hagenberg, Austria33institutetext: Institute of
Cosmology & Gravitation, University of Portsmouth, Dennis Sciama Building,
Portsmouth, PO1 3FX, UK44institutetext: Astrophysics, University of Oxford,
Denys Wilkinson Building, Keble Road, Oxford OX1 3RH, UK55institutetext:
Center for Computational Astrophysics, Flatiron Institute, 162 5th Avenue, New
York, NY 10010, USA

# A precise symbolic emulator of the linear matter power spectrum

Deaglan J. Bartlett<EMAIL_ADDRESS>Lukas Kammerer 22 Gabriel
Kronberger 22 Harry Desmond 33 Pedro G. Ferreira 44 Benjamin D. Wandelt 1155
Bogdan Burlacu 22 David Alonso 44 Matteo Zennaro 44

(Received XXX; accepted YYY)

###### Abstract

Context. Computing the matter power spectrum, $P(k)$, as a function of
cosmological parameters can be prohibitively slow in cosmological analyses,
hence emulating this calculation is desirable. Previous analytic
approximations are insufficiently accurate for modern applications, so black-
box, uninterpretable emulators are often used.

Aims. To construct an efficient, differentiable, interpretable, symbolic
emulator for the redshift zero linear matter power spectrum which achieves
sub-percent level accuracy. We also wish to obtain a simple analytic
expression to convert $A_{\rm s}$ to $\sigma_{8}$ given the other cosmological
parameters.

Methods. We utilise an efficient genetic programming based symbolic regression
framework to explore the space of potential mathematical expressions which can
approximate the power spectrum and $\sigma_{8}$. We learn the ratio between an
existing low-accuracy fitting function for $P(k)$ and that obtained by solving
the Boltzmann equations and thus still incorporate the physics which motivated
this earlier approximation.

Results. We obtain an analytic approximation to the linear power spectrum with
a root mean squared fractional error of 0.2% between $k=9\times
10^{-3}-9\,h{\rm\,Mpc^{-1}}$ and across a wide range of cosmological
parameters, and we provide physical interpretations for various terms in the
expression. We also provide a simple analytic approximation for $\sigma_{8}$
with a similar accuracy, with a root mean squared fractional error of just
0.4% when evaluated across the same range of cosmologies. This function is
easily invertible to obtain $A_{\rm s}$ as a function of $\sigma_{8}$ and the
other cosmological parameters, if preferred.

Conclusions. It is possible to obtain symbolic approximations to a seemingly
complex function at a precision required for current and future cosmological
analyses without resorting to deep-learning techniques, thus avoiding their
black-box nature and large number of parameters. Our emulator will be usable
long after the codes on which numerical approximations are built become
outdated.

###### Key Words.:

Cosmology: theory, Cosmology: cosmological parameters, Cosmology: large-scale
structure of Universe, Methods: numerical

##  1 Introduction

Machine learning (ML) methods have great potential for simplifying and
accelerating the analysis of astrophysical data sets. The primary focus has
been on what one might dub “advanced numerical methods” using, for example,
Gaussian processes or neural networks. In these cases, one tries to construct
efficient algorithms which can be used to either infer specific physical
properties from complex data sets or emulate complex processes which can then
be extrapolated to new situations. Typically these methods involve
constructing a set of pre-established basis functions and then inferring their
weights, or building complex, expressible functions, with parameters that can
be optimised via efficient gradient descent methods. These methods can be
easily incorporated in Bayesian inference frameworks that have achieved
significant success, becoming the standard practice in astrostatistics.

The drawback of the more traditional, numerical ML techniques is their
opaqueness; it is not always clear what information is being used and how
methods trained on (necessarily imperfect) simulations will perform when
applied to real-world data. A somewhat overlooked branch of machine learning
which has tremendous promise for the types of problems being considered in
astrophysics is Symbolic Regression (SR). With SR one tries to infer the
mathematical expressions that best capture the properties of the physical
system one is trying to study. The process is an attempt to mimic and
systematise the practice that physicists have always used: to infer simple
physical laws (i.e. formulae) from data. The field of SR has developed over
the years into a vibrant and active field of research in ML, typically
associated with evolutionary methods such as Genetic Programming. It has been
shown that it can be used to infer some well-established laws of physics from
data and infer new ones (e.g. Lemos et al., 2023; Bartlett et al., 2022, 2023;
Desmond et al., 2023; Kamerkar et al., 2023; Sousa et al., 2023; Delgado et
al., 2022; Miniati & Gregori, 2022; Wadekar et al., 2022).

Within the field of cosmology, one often compresses observations from galaxy
surveys into two-point correlation functions (or their Fourier transforms,
power spectra), which are compared to theory through Markov Chain Monte Carlo
methods to constrain cosmological parameters. As cosmological surveys become
increasingly vast and precise, a fundamental limitation to the feasibility of
such inferences has been the speed at which one can make this theoretical
prediction, since it involves solving a complex set of coupled, highly
nonlinear differential equations. Recently, instead of directly solving these
equations (Lewis et al., 2000; Blas et al., 2011; Hahn et al., 2023) and
adding non-linear corrections (Smith et al., 2003), emulation techniques such
as Neural Networks or Gaussian Processes have been used to accelerate these
calculations to directly output the matter power spectrum as a function of
cosmological parameters (Heitmann et al., 2014; Winther et al., 2019; Angulo
et al., 2021; Aricò et al., 2021; Euclid Collaboration et al., 2021; Spurio
Mancini et al., 2022; Mootoovaloo et al., 2022; Zennaro et al., 2023). These
methods act as black boxes and require up to several hundreds of parameters to
be optimised. However, through perturbation theory, one knows analytic limits
of the power spectrum and, through visual inspection, it does not appear to be
an extremely complex function. As such, one wonders whether an analytic
approximation exists. Indeed, for many years, the leading method of
accelerating this calculation has been an analytic approximation (Eisenstein &
Hu, 1998, 1999), however it is insufficiently precise for modern experiments.
Analytic approximations to beyond $\Lambda$CDM power spectra have been
proposed in the context of modified gravity (Bayron Orjuela-Quintana et al.,
2023), although these still only achieve a precision of between 1 and 2%.

Such an emulator has the advantage that it will not become deprecated when the
codes on which current numerical methods are built become outdated, whereas
other methods require the transfer of the inferred weights and biases as well
as the model architecture, hindering longevity. Even in the short term, an
analytic expression using standard operators is more portable, since it can be
more easily be incorporated into the user’s favourite programming language
without the need to install or write wrappers for the model. Moreover, having
an analytic expression allows one to interpret such a fit, and potentially
identify physical processes which could lead to certain terms, contrary to the
black-box numerical methods. Additionally, such expressions often contain
fewer free parameters to optimise than numerical ML methods.

In Section 2 we briefly describe the matter power spectrum and the Eisenstein
& Hu approximation, and in Section 3 we detail the SR method we use in this
work. We present an analytic emulator for $\sigma_{8}$ as a function of other
cosmological parameters in Section 4 (which is easily invertible to obtain
$A_{\rm s}$ as a function of cosmological parameters), and in Section 5 we
give our emulator for the linear matter power spectrum. The main results of
this paper are given in Eq. 4 and Eq. 6. We conclude and discuss future work
in Section 6.

Throughout this paper “$\log$” denotes the natural logarithm.

##  2 The matter power spectrum

###  2.1 Definition

We would like to construct an efficient, differentiable and (if at all
possible) interpretable emulator for the power spectrum of the matter
distribution in the Universe, $P(k;\bm{\theta})$, for wavenumber $k$ and
cosmological parameters $\bm{\theta}$.

The power spectrum is defined as follows: the matter density of the Universe,
$\rho(\bm{x})$ can be decomposed into a constant (in space) background
density, ${\bar{\rho}}$, and a density contrast, $\delta(\bm{x})$ such that
$\rho(\bm{x})={\bar{\rho}}[1+\delta(\bm{x})]$. If $\tilde{\delta}(\bm{k})$ is
the Fourier Transform of $\delta(\bm{x})$, and the matter distribution is
statistically homogeneous and isotropic, we have that

$(2\pi)^{3}P(k;\bm{\theta})\delta^{\rm
D}\left(\bm{k}-\bm{k}^{\prime}\right)\equiv\langle\tilde{\delta}(\bm{k})\tilde{\delta}^{\ast}(\bm{k}^{\prime})\rangle,$
(1)

where $\langle\cdots\rangle$ denotes an ensemble average and $\delta^{\rm D}$
is the Dirac delta function.

From observations of the Cosmic Microwave Background (CMB) (Planck
Collaboration, 2020), it is known that the density fluctuations at early times
were approximately Gaussian and thus fully described by $P(k;\bm{\theta})$. At
these early times, the power spectrum of the comoving curvature perturbations
is proportional to $A_{\rm s}k^{n_{\rm s}-4}$, where $n_{\rm s}\approx 0.9665$
(Planck Collaboration, 2020). Although structure formation through gravity
makes the present-day density field non-Gaussian (e.g. the intricate structure
of the cosmic web is typically associated with higher order statistics), the
power spectrum still holds a central role in modern cosmological analyses.

The current cosmological model is described by only six parameters: the
baryonic, $\Omega_{\rm b}$, and total matter, $\Omega_{\rm m}$, density
parameters, the Hubble constant, $H_{0}=100h{\rm\,km\,s^{-1}\,Mpc^{-1}}$, the
scalar spectral index, $n_{\rm s}$, the curvature fluctuation amplitude,
$A_{\rm s}$, and the reionisation optical depth, $\tau$. All other parameters
can be derived from these six, and thus sometimes a different set of
parameters is chosen. For example, instead of $A_{\rm s}$, one often quotes
$\sigma_{8}$ which is the root-mean-square density fluctuation when the
linearly evolved field is smoothed with a top-hat filter of radius
$8\,h^{-1}\text{Mpc}$. Specifically, one defines for a top-hat of radius $R$

$\sigma_{R}^{2}=\int{\rm
d}k\,\frac{k^{2}}{2\pi^{2}}P(k;\bm{\theta})\left|W(k,R)\right|^{2},$ (2)

where $\bm{\theta}$ is the set of cosmological parameters and the Fourier
transfer of the top-hat filter is

$W(k,R)=\frac{3}{(kR)^{3}}\left(\sin(kR)-kR\cos(kR)\right),$ (3)

and $\sigma_{8}$ is simply $\sigma_{R}$ for $R=8h^{-1}\text{Mpc}$. Throughout
this paper we will ignore the small dependence of the power spectrum on the
reionisation optical depth parameter, and focus on the remaining five
parameters. We set the neutrino mass to zero in all calculations.

###  2.2 Eisenstein & Hu Approximation

Since each evaluation of a Boltzmann solver to compute $P(k;\bm{\theta})$ can
be expensive, the ability to emulate this procedure and replace this solver
with a surrogate model has long been desirable. The most notable attempt to do
this in an analytic manner is given in a series of papers by Eisenstein & Hu
(1998, 1999). In these works, an approximation is constructed based on
physical arguments including baryonic acoustic oscillations (BAO), Compton
drag, velocity overshoot, baryon infall, adiabatic damping, Silk damping, and
cold dark matter growth suppression. Rather than repeat their findings, we
refer the reader to these papers to inspect the structure of the equations.
Such a model is accurate to a few percent which, although invaluable at the
time of writing, is insufficiently accurate for modern cosmological analyses.
It is thus the goal of this work to build upon this analytic emulator to
provide sub-percent level predictions. We note that alternative symbolic
approximations also exist to $P(k)$, such as the earlier, less accurate
approximation by Bardeen et al. (1986) (BBKS). More recently, Bayron Orjuela-
Quintana et al. (2022) found simple expressions using genetic programming
which can achieve similar accuracy to the Eisenstein & Hu expression, but we
choose to use Eisenstein & Hu’s approximation due to its physical motivation
and widespread use.

##  3 Symbolic regression

To extract analytic approximations from sampled data, we use the symbolic
regression package operon111https://github.com/heal-research/operon (Burlacu
et al., 2020). This package leverages the most popular (e.g. Lemos et al.,
2023; Cranmer et al., 2020; Cranmer, 2020, 2023; Schmidt & Lipson, 2009, 2011;
Virgolin et al., 2021a; de Franca & Aldeia, 2021; La Cava et al., 2019;
Kommenda et al., 2020; Virgolin et al., 2021b; Arnaldo et al., 2014) approach
to SR, namely genetic programming (Turing, 1950; David, 1989; Haupt & Haupt,
2004). Genetic programming describes the evolution of “computer programs”, in
our case mathematical expressions encoded as expression trees. Following the
principle of natural selection, over several iterations the worst performing
equations (given some fitness metric) are discarded and new equations are
produced by combining sub-expressions of the current population (crossover) or
by randomly inserting, replacing or deleting a subtree in an expression
(mutation). Over the course of several generations, the expectation is that
the population of equations evolve to become fitter and thus we obtain
increasingly accurate analytic expressions.

We note that many other techniques exist for SR, such as supervised or
reinforcement learning with neural networks (Petersen et al., 2021; Landajuela
et al., 2022; Tenachi et al., 2023; Biggio et al., 2021), deterministic
approaches (Worm & Chiu, 2013; Kammerer et al., 2021; Rivero et al., 2022;
McConaghy, 2011), Markov chain Monte Carlo (Jin et al., 2019), physics-
inspired searches (Udrescu & Tegmark, 2020; Udrescu et al., 2020; René Broløs
et al., 2021), and exhaustive searches (Bartlett et al., 2022). However, we
choose operon and thus genetic programming due to its speed, high memory
efficiency and its strong performance in benchmark studies (Cava et al., 2021;
Burlacu, 2023).

To improve the search, every time a terminal node appears in an expression
tree (i.e. $k$ or one of the cosmological parameters), a scaling parameter is
introduced, which is then optimised (Kommenda et al., 2020) using the
Levenberg–Marquardt algorithm (Levenberg, 1944; Marquardt, 1963). We denote
the total number of nodes in the expression excluding the scaling as the
“length” of the model, and the “complexity” refers to the total number of
nodes, including these.

When comparing objective values during non-dominated sorting (NSGA2), operon
implements the concept of $\epsilon$-dominance (Laumanns et al., 2002), where
the parameter $\epsilon$ is defined such that two objective values which are
within $\epsilon$ of each other are considered equal. This parameter therefore
affects the number of duplicate equations in the population and is designed to
promote convergence to a representative well distributed approximation of the
global Pareto front. We choose different values for this parameter when
searching for our two emulators, and these were found after some
experimentation with different values to find settings which produced accurate
yet compact models.

Model selection is an essential part of any SR search. Since one optimises
both accuracy and simplicity during the search, SR if often a Pareto-
optimisation problem. Although principled methods exist to combine these
competing requirements in the presence of statistical errors (Bartlett et al.,
2022, 2023), we do not have this ability for our current application since
there is no noise in our data, and thus have to rely on more heuristic
methods. We visually inspect the most accurate solution found for each model
length and make a qualitative judgement as to the function which is
sufficiently compact to be interpretable yet is accurate enough for our
applications. Further details are given in Sections 4 and 5.

##  4 Analytic emulator for $\sigma_{8}$

Parameter | Minimum | Maximum
---|---|---
$10^{9}\,A_{\rm s}$ | 1.7 | 2.5
$\Omega_{\rm m}$ | 0.24 | 0.40
$\Omega_{\rm b}$ | 0.04 | 0.06
$h$ | 0.61 | 0.73
$n_{\rm s}$ | 0.92 | 1.00
Table 1: Cosmological parameters used for analytic emulators. We sample all
parameters independently and uniformly in the range between the minimum and
maximum values given. Figure 1: Pareto front of solutions obtained using
operon when fitting $\sigma_{8}$ as a function of $A_{\rm s}$, $\Omega_{\rm
b}$, $\Omega_{\rm m}$, $h$ and $n_{\rm s}$. We plot the root mean squared
error as a function of model length from the training and validation sets
separately. The model given in Eq. 4 has a model length of 15.

We begin by considering the simplest emulator one may want for power spectrum
related quantities: an emulator for $\sigma_{8}$ as a function of other
cosmological parameters ($A_{\rm s},\Omega_{\rm b},\Omega_{\rm m},h,n_{\rm
s}$) or, equivalently, an emulator for $A_{\rm s}$ given $\sigma_{8}$ and the
other cosmological parameters. Although the set of neural network emulators
BACCO contains a function to do this (Aricò et al., 2021), to the best of the
authors’ knowledge an analytic approximation is not currently in common use.
The standard approach is to compute the linear matter power spectrum with a
Boltzmann code assuming some initial guess of $A_{\rm s}$, then compute the
integral in Eq. 2 to obtain $\sigma_{8}$. For a target $\sigma_{8}$ of
$\sigma_{8}^{\prime}$, one should then use $A_{\rm
s}^{\prime}=(\sigma_{8}^{\prime}/\sigma_{8})^{2}A_{\rm s}$.

We wish to accelerate this process with a symbolic emulator. To compute this,
we construct a Latin hypercube (LH) of 100 sets of cosmological parameters,
using uniform priors in the ranges given in Table 1, which are the same as
those used in Euclid Collaboration et al. (2021). We construct a second LH of
100 points to be used for validation. For these parameters, we compute
$\sigma_{8}$ using camb (Lewis et al., 2000) and attempt to learn this mapping
using a mean squared error loss function with operon.

For the equation search, we use a population size of 1000 with a brood size of
10 and tournament size of 5, optimising both the mean squared error and the
length of the expression simultaneously, with $\epsilon=10^{-5}$ (see Section
3). For numerical stability, we fit using $10^{9}A_{\rm s}$ instead of $A_{\rm
s}$ so that all cosmological parameters are $\mathcal{O}(1)$. Parameters are
optimised during the search using a nonlinear least squared optimiser with up
to 1000 iterations per optimisation attempt. We set the maximum allowed model
length to 200 and maximum number of iterations to $10^{7}$, although we find
that both of these are much larger than the required values needed to converge
to a desirable solution. The candidate expressions are comprised of standard
arithmetic operations (addition, subtraction, multiplication, division), as
well as the natural logarithm, cosine, power and analytic quotient operator
($\operatorname{aq}(x,y)\equiv x/\sqrt{1+y^{2}}$).

After 30 minutes of operation on one node of 128 cores, we find the Pareto
front of expressions given in Fig. 1. We see that the training and validation
losses are comparable at all model lengths, reaching a root mean squared error
of around than $3\times 10^{-3}$ by a length of 14. The best model found (of
length 15) is given by

$\sigma_{8}\approx\left(a_{0}A_{\rm s}+a_{1}n_{\rm
s}\right)\left(a_{2}\Omega_{\rm b}+\log{\left(a_{3}\Omega_{\rm
m}\right)}\right)\log{\left(a_{4}h\right)}+a_{5},$ (4)

where the optimised parameters are $\bm{a}=[1.61320734729\times
10^{8}\mathchar 44\relax\penalty 00.343134609906\mathchar 44\relax\penalty
0-7.859274\mathchar 44\relax\penalty 018.200232\mathchar 44\relax\penalty
03.666163\mathchar 44\relax\penalty 00.003359]$. We note several important
features of this equation which make it desirable. First, we find that it is a
highly accurate approximation, with a root mean squared fractional error on
the validation set of only 0.4%, which is far smaller than the precision to
which one can measure this number with cosmological experiments. Second, one
sees the $A_{\rm s}$ only appears once in this equation and as a linear term.
Thus, it is trivial to invert this equation to obtain $A_{\rm s}$ as a
function of the other cosmological parameters, as is often needed. Finally, we
note that the final parameter, $a_{5}=0.003359$, is much smaller than the
value of $\sigma_{8}$ and thus this additive constant could be neglected in
most applications where such precision is not necessary.

##  5 Analytic emulator for the linear power spectrum

We now move on to the more challenging task of producing an analytic emulator
for the linear matter power spectrum. Given the previous success of Eisenstein
& Hu (1998, 1999), we believe it is sensible to build upon this work, not
least due to the physically-motivated terms included in their fit and so that
we must only have to fit a small residual (of the order of a few percent).
Thus, instead of directly fitting for $P(k,\bm{\theta})$, we define

$P(k;\bm{\theta})\equiv P_{\rm EH}(k;\bm{\theta})F(k;\bm{\theta}),$ (5)

where $P_{\rm EH}(k;\bm{\theta})$ is the zero-baryon fit of Eisenstein & Hu
(1998), which does not include an attempt to fit the BAO. We plot both
$P(k;\bm{\theta})$ and $\log F(k;\bm{\theta})$ in Fig. 2 for the best-fit
cosmology obtained by Planck (Planck Collaboration, 2020), where we see that
dividing out the Eisenstein & Hu term retains the BAO part of the power
spectrum and reduces the dynamic range required for the fit.

As before, we obtain 100 sets of cosmological parameters on a LH using the
priors in Table 1 and compute both $P(k;\bm{\theta})$ with camb and $P_{\rm
EH}(k;\bm{\theta})$ with the colossus (Diemer, 2018) implementation, using 200
logarithmically spaced values of $k$ in the range $9\times
10^{-3}-9\,h{\rm\,Mpc^{-1}}$. We note that this is an extremely small training
set compared to many power spectrum emulators, but we find that it is
sufficient to obtain sub-percent level fits.

We choose to symbolically regress $\log F(k;\bm{\theta})$ using a mean squared
error loss function, and thus wish to minimise the fractional error on this
residual. We choose to fit for $\log F$ as this ensures that our final
estimate of $P(k;\bm{\theta})$ is positive, as guaranteed by exponentiation,
which is physically required. Additionally, we first multiply $\log F$ by 100
so that the target is $\mathcal{O}(1)$. We use a further 100 sets of
cosmological parameters, also arranged on a LH, for validation. We choose to
fit using the cosmological parameters $\sigma_{8}$, $\Omega_{\rm b}$,
$\Omega_{\rm m}$, $h$ and $n_{\rm s}$. We use the same settings for operon as
in Section 4, except we choose $\epsilon=10^{-3}$ and terminate our search
after $10^{8}$ function evaluations.

The root mean squared error for the best function found at each model length
is given in Fig. 3, where we see that we are able to achieve values of
$\mathcal{O}(10^{-3})$ for $\log F$. Unlike for the $\sigma_{8}$ emulator, we
obtain slightly worse losses for the validation set compared to training,
however always by less than a factor of two.

Given this set of candidate solutions, we wish to choose one which is
sufficiently accurate for current applications yet is sufficiently compact to
be interpretable. In Fig. 3, one observes a plateau in accuracy between model
lengths $\sim 65-80$ and thus it seems reasonable to choose a solution in this
regime, since doubling the model length only achieves approximately a factor
of two improvement in fit beyond this point. Moreover, beyond this point the
training and validation curves begin to deviate, suggesting a degree of
overfitting.

We choose to report the model of length 77, as indicated by the dotted line in
Fig. 3, since this provided one of the most interpretable solutions, and
achieved a sub-percent error for 95% ($2\sigma$) of the cosmological
parameters considered, for both the training and validation set. After some
simplification, this can be written as

$\begin{split}\log F&\approx b_{0}h-b_{1}\\\ &+\left(\frac{b_{2}\Omega_{\rm
b}}{\sqrt{h^{2}+b_{3}}}\right)^{b_{4}\Omega_{\rm
m}}\left[\frac{b_{5}k-\Omega_{\rm b}}{\sqrt{b_{6}+(\Omega_{\rm
b}-b_{7}k)^{2}}}b_{8}(b_{9}k)^{-b_{10}k}\cos\left(b_{11}\Omega_{\rm
m}-\frac{b_{12}k}{\sqrt{b_{13}+\Omega_{\rm
b}^{2}}}\right)-b_{14}\left(\frac{b_{15}k}{\sqrt{1+b_{16}k^{2}}}-\Omega_{\rm
m}\right)\cos\left(\frac{b_{17}h}{\sqrt{1+b_{18}k^{2}}}\right)\right]\\\
&+b_{19}(b_{20}\Omega_{\rm
m}+b_{21}h-\log(b_{22}k)+(b_{23}k)^{-b_{24}k})\cos\left(\frac{b_{25}}{\sqrt{1+b_{26}k^{2}}}\right)\\\
&+(b_{27}k)^{-b_{28}k}\left(b_{29}k-\frac{b_{30}\log(b_{31}k)}{\sqrt{b_{32}+(\Omega_{\rm
m}-b_{33}h)^{2}}}\right)\cos\left(b_{34}\Omega_{\rm
m}-\frac{b_{35}k}{\sqrt{b_{36}+\Omega_{\rm b}^{2}}}\right),\end{split}$ (6)

Figure 2: Linear matter power spectrum (upper), the residuals (Eq. 5) from the Eisenstein & Hu fit without baryons (middle), and the fractional residuals on $P(k)$ compared to the truth for the Planck 2018 (Planck Collaboration, 2020) cosmology. In all panels we plot the truth computed with camb with solid red lines, and the analytic fit (Eq. 6) obtained in this paper with dashed blue lines. We see that the fit is accurate within 0.3% across all $k$ considered.  Parameter | Value | Parameter | Value
---|---|---|---
$b_{0}$ | 0.0545 | $b_{19}$ | 0.0111
$b_{1}$ | 0.0038 | $b_{20}$ | 5.35
$b_{2}$ | 0.0397 | $b_{21}$ | 6.421
$b_{3}$ | 0.1277 | $b_{22}$ | 134.309
$b_{4}$ | 1.35 | $b_{23}$ | 5.324
$b_{5}$ | 4.0535 | $b_{24}$ | 21.532
$b_{6}$ | 0.0008 | $b_{25}$ | 4.742
$b_{7}$ | 1.8852 | $b_{26}$ | 16.6872
$b_{8}$ | 0.1142 | $b_{27}$ | 3.078
$b_{9}$ | 3.798 | $b_{28}$ | 16.987
$b_{10}$ | 14.909 | $b_{29}$ | 0.0588
$b_{11}$ | 5.56 | $b_{30}$ | 0.0007
$b_{12}$ | 15.8274 | $b_{31}$ | 195.498
$b_{13}$ | 0.0231 | $b_{32}$ | 0.0038
$b_{14}$ | 0.8653 | $b_{33}$ | 0.2767
$b_{15}$ | 0.8425 | $b_{34}$ | 7.385
$b_{16}$ | 4.554 | $b_{35}$ | 12.3961
$b_{17}$ | 5.117 | $b_{36}$ | 0.0134
$b_{18}$ | 70.0234 |  |
Table 2: Best-fit parameters for the linear matter power spectrum emulator
given in Eq. 6. Although units are excluded in this table, the units for each
parameter are easily obtained by noting that these are defined assuming that
$k$ is measured in $h\,{\rm Mpc^{-1}}$ in Eq. 6. Figure 3: Pareto front of
solutions obtained using operon when fitting the linear matter power spectrum
as a function of $\sigma_{8}$, $\Omega_{\rm b}$, $\Omega_{\rm m}$, $h$ and
$n_{\rm s}$. We plot the root mean squared error on $\log F$ as a function of
model length for the training and validation sets separately. The model given
in Eq. 6 has a model length of 77, as indicated by the dotted line. Figure 4:
The distribution of fractional errors as a function of $k$ on the linear
matter power spectrum across all cosmologies in the training and validation
sets, as compared to the predictions of camb. The bands give the 1 and
$2\sigma$ values. The dotted line corresponds to a 1% error, and we see that
our expression achieves this for all cosmologies and values of $k$ considered,
with a root mean squared fractional error of 0.2%. Figure 5: Contributions to
$\log F$ from our emulator as a function of $k$ for the Planck 2018 cosmology.
The line numbers indicated in the legend correspond to the line in Eq. 6. One
sees that the first term provides an overall offset, the second and fourth
capture the BAO signal, and the third term contains a broad oscillation and
then matches on to the decaying residual at high $k$.

where the best-fit parameters for this function are given in Table 2. We find
that there are 37 different parameters required for this fit, far fewer than
would be used if one were to emulate this with a neural network.

We plot this fit and the residuals compared to camb for the Planck 2018
cosmology in Fig. 2, which we note is not included in either our training or
validation sets. One can see that the difference between the true power
spectrum and our analytic fit is almost imperceivable, and in the residuals
plot we see that for all $k$ considered, the fractional error does not exceed
0.3%. This is smaller than the error on $\log F$ given in Fig. 3, since we
compare at the level of the full $P(k;\bm{\theta})$, such that a moderate
error on $\log F$ becomes very small once substituted into Eq. 5. This is
shown in Fig. 4, where we plot the distribution of fractional residuals in
$P(k;\bm{\theta})$ for all the cosmologies in the training and validations
sets. We obtain sub-percent level predictions for all cosmologies and values
of $k$ considered, with a root mean squared fractional error of 0.2%.

Part of the appeal of a symbolic emulator is the possibility for
interpretability and to easily identify what information used in the input is
used to make the prediction. To begin, we note that, although we obtained our
emulator by varying $\sigma_{8}$, $\Omega_{\rm b}$, $\Omega_{\rm m}$, $h$ and
$n_{\rm s}$, we see that Eq. 6 contains neither $\sigma_{8}$ nor $n_{\rm s}$.
For the linear matter power spectrum, one expects that $A_{\rm s}$ and $n_{\rm
s}$ only appear as a multiplicative factor of $A_{\rm s}k^{n_{\rm s}-1}$, with
all other terms independent of these parameters. Given that the Eisenstein &
Hu term already contains this expression, it is unsurprising that $\log F$ is
independent of $n_{\rm s}$. Indeed, if it did appear, this would indicate a
degree of overfitting. Since $\sigma_{8}$ is not proportional to $A_{\rm s}$
(see Eq. 4), we cannot use the same argument to explain the lack of its
appearance in our expression, but can conclude that a combination of the
Eisenstein & Hu term and the first line of Eq. 6 can sufficiently approximate
$A_{\rm s}$, since this line is $k$ independent and thus contributes to an
overall offset for the emulator.

Turning to the remaining lines of Eq. 6, we observe that each term contains an
oscillation modulated by a $k$\- and cosmology-dependent damping. Despite
there being four such terms across the remaining three lines, we find that we
can split these into two pairs with the same structure of the oscillations.
Firstly, we have cosines with an argument proportional to $1/\sqrt{1+bk^{2}}$,
for some constant $b$. This functional form ($x/\sqrt{1+y^{2}}$) arises due to
the inclusion of the analytic quotient operator, which also explains why the
constant $1$ appears multiple times in Eq. 6. These terms give oscillations
which vary slowly as a function of $k$. In particular, as plotted in Fig. 5,
the third line of Eq. 6 contains approximately one cycle of oscillation across
the range of $k$ considered, with a minimum during the BAO part of the power
spectrum, and a maximum just afterwards. Beyond this point, this term fits the
non-oscillatory, decaying part of the residual beyond $k\sim
1\,h{\rm\,Mpc}^{-1}$ (compare the middle panel of Fig. 2 to the third term
plotted in Fig. 5).

The remaining oscillatory terms are of the form $\cos(\omega k+\phi)$. The
phase, $\phi$, of these oscillations is proportional to the total matter
density, $\Omega_{\rm m}$, such that changing this parameter at fixed
$\Omega_{\rm b}$ shifts the BAOs to peak at different values of $k$. The
frequency of these oscillations is $\omega\propto 1/\sqrt{b+\Omega_{\rm
b}^{2}}$ for some parameter $b$, such that cosmologies with a higher fraction
of baryons have many more cycles of BAO in a given range of $k$, as one would
physically expect. From Fig. 5, one can see how the second and fourth lines of
Eq. 6 capture the BAO signal with opposite signs, such that they combine to
give the familiar damped oscillatory feature. Using $\Omega_{\rm
b}h^{2}=0.02242$ and $h=0.6766$, as appropriate for the Planck 2018 cosmology
(Planck Collaboration, 2020), the frequency of the oscillations are
$b_{12}/\left(h\sqrt{b_{13}+\Omega_{\rm b}^{2}}\right)=146.5{\rm\,Mpc}$ and
$b_{35}/\left(h\sqrt{b_{36}+\Omega_{\rm b}^{2}}\right)=145.8{\rm\,Mpc}$, and
are thus approximately equal to the sound horizon, which is
$r_{\ast}=144.6{\rm\,Mpc}$ for this cosmology. One can therefore view these
frequencies as symbolic approximations to the sound horizon, although we refer
the reader to Aizpuru et al. (2021) for alternative SR fits.

Thus, although we did not enforce physically motivated terms in the equation
search, we see that simple oscillatory contributions for the BAOs have emerged
and thus our symbolic emulator is not merely a high order series expansion,
but contains terms which are both compact and interpretable. We find that such
terms exist in many functions given in Fig. 3, however we find that using
shorter run times for operon of only 2-4 hours (compared to approximately 24
hours on a single node of 128 cores for our fiducial analysis) do not provide
as interpretable expressions as Eq. 6.

As a note of caution, one can identify a few terms in Eq. 6 which will become
problematic if extrapolated to values of $k$ much smaller than those used to
train the emulator, namely those containing $\log(k)$ and $k$ raised to a
power proportional to $k$. For $k\lesssim 10^{-3}\,h{\rm\,Mpc}^{-1}$ this can
lead to an error on $P(k)$ of more than one percent. Although one is likely
cosmic-variance dominated in this regime so such errors should not be
problematic, we know that the Eisenstein & Hu provides a very good
approximation, and thus we suggest that Eq. 6 is included in a piece-wise fit,
such that it is only used approximately in the range of $k$ which were used to
obtain it. A similar effect is seen if one extrapolates to higher $k$ than
considered here. Although this is far beyond the validity of the linear
approximation, we caution that applying any parameterisation of the non-linear
power spectrum that depends on the linear one may suffer from potentially
catastrophic extrapolation failures at high $k$ if used beyond the $k$ range
considered here. Again, it is potentially advisable to just use the Eisenstein
& Hu fit in this regime.

If the reader wishes to use a more accurate, yet less interpretable, emulator,
we provide the most accurate equation found in Appendix A, which has a model
length of 142, with 73 parameters and yields a root mean of squared fractional
errors on $P(k)$ of 0.1% for both the training and validation sets.

##  6 Discussion and Conclusion

In this paper we have found analytic approximations to $\sigma_{8}$ (Eq. 4)
and the linear matter power spectrum (Eq. 6) as a function of cosmological
parameters which are accurate to sub-percent levels. In the case of
$\sigma_{8}$, the simple yet accurate expression we have identified can be
easily inverted to obtain $A_{\rm s}$ as a function of $\sigma_{8}$ and the
other cosmological parameters. Our approximation to $P(k)$ is built by fitting
the residuals between the output of a Boltzmann solved (camb) and the physics-
inspired approximation of Eisenstein & Hu (1998). As such, unlike neural
network or Gaussian process based approaches, our expression explicitly
captures many physical processes (and is thus interpretable) whilst still
achieving sub-percent accuracy.

This work is the first step in a programme of work dedicated to obtaining
analytic approximations to $P(k)$ which can be used in current and future
cosmological analyses. In this paper we have focused on the linear $P(k)$,
i.e. the power spectrum of the linearly evolved density fluctuations. Although
this approach is valid on large scales, the real Universe is non-linear, such
that non-linear corrections are required at $k\gtrsim
10^{-1}\,h{\rm\,Mpc}^{-1}$ to accurately model the observed matter power
spectrum across a wider range of scales. Future work will be dedicated to
extending our framework to capture such non-linear physics and to include
redshift dependence in our emulator. Finally, in our emulator we have
considered a $\Lambda$CDM Universe with massless neutrinos. In the future we
will add corrections to the expressions found in this work to incorporate the
effects of massive neutrinos and include beyond $\Lambda$CDM effects, such as
a $w_{0}-w_{a}$ parametrisation of dark energy.

We have demonstrated that, despite the temptation to blindly apply black-box
methods such as neural networks to approximate physically useful functions,
even in ostensibly challenging situations such as the matter power spectrum,
one can achieve the required precision with relatively simple analytic fits.
Given the unknown lifetime of current codes upon which numerical ML
approximations are built and the ease of copying a few mathematical functions
into your favourite programming language, finding analytic expressions allows
one to more easily future-proof such emulators and should therefore be
encouraged wherever possible.

## Acknowledgements

DJB is supported by the Simons Collaboration on “Learning the Universe.” LK
was supported by a Balzan Fellowship. HD is supported by a Royal Society
University Research Fellowship (grant no. 211046). PGF acknowledges support
from STFC and the Beecroft Trust. BDW acknowledges support from the Simons
Foundation. DA acknowledges support from the Beecroft Trust, and from the
Science and Technology Facilities Council through an Ernest Rutherford
Fellowship, grant reference ST/P004474/1. MZ is supported by STFC. We made
extensive use of computational resources at the University of Oxford
Department of Physics, funded by the John Fell Oxford University Press
Research Fund, and at the Institut d’Astrophysique de Paris.

For the purposes of open access, the authors have applied a Creative Commons
Attribution (CC BY) licence to any Author Accepted Manuscript version arising.

## Data availability

The data underlying this article will be shared on reasonable request to the
corresponding author. We provide python implementations of Eqs. 4, 5, 6 and 7
at https://github.com/DeaglanBartlett/symbolic_pofk.

## References

  * Aizpuru et al. (2021) Aizpuru, A., Arjona, R., & Nesseris, S. 2021, Phys. Rev. D, 104, 043521
  * Angulo et al. (2021) Angulo, R. E., Zennaro, M., Contreras, S., et al. 2021, MNRAS, 507, 5869
  * Aricò et al. (2021) Aricò, G., Angulo, R. E., & Zennaro, M. 2021, arXiv e-prints, arXiv:2104.14568
  * Arnaldo et al. (2014) Arnaldo, I., Krawiec, K., & O’Reilly, U.-M. 2014, in Proceedings of the 2014 Annual Conference on Genetic and Evolutionary Computation, GECCO ’14 (New York, NY, USA: Association for Computing Machinery), 879–886
  * Bardeen et al. (1986) Bardeen, J. M., Bond, J. R., Kaiser, N., & Szalay, A. S. 1986, ApJ, 304, 15
  * Bartlett et al. (2022) Bartlett, D. J., Desmond, H., & Ferreira, P. G. 2022, IEEE Transactions on Evolutionary Computation [arXiv:2211.11461]
  * Bartlett et al. (2023) Bartlett, D. J., Desmond, H., & Ferreira, P. G. 2023, in The Genetic and Evolutionary Computation Conference 2023
  * Bayron Orjuela-Quintana et al. (2022) Bayron Orjuela-Quintana, J., Nesseris, S., & Cardona, W. 2022, arXiv e-prints, arXiv:2211.06393
  * Bayron Orjuela-Quintana et al. (2023) Bayron Orjuela-Quintana, J., Nesseris, S., & Sapone, D. 2023, arXiv e-prints, arXiv:2307.03643
  * Biggio et al. (2021) Biggio, L., Bendinelli, T., Neitz, A., Lucchi, A., & Parascandolo, G. 2021, in Proceedings of Machine Learning Research, Vol. 139, Proceedings of the 38th International Conference on Machine Learning, ed. M. Meila & T. Zhang (PMLR), 936–945
  * Blas et al. (2011) Blas, D., Lesgourgues, J., & Tram, T. 2011, J. Cosmology Astropart. Phys., 2011, 034
  * Burlacu (2023) Burlacu, B. 2023, in Proceedings of the Companion Conference on Genetic and Evolutionary Computation, GECCO ’23 Companion (New York, NY, USA: Association for Computing Machinery), 2412–2419
  * Burlacu et al. (2020) Burlacu, B., Kronberger, G., & Kommenda, M. 2020, in Proceedings of the 2020 Genetic and Evolutionary Computation Conference Companion, GECCO ’20 (New York, NY, USA: Association for Computing Machinery), 1562–1570
  * Cava et al. (2021) Cava, W. G. L., Orzechowski, P., Burlacu, B., et al. 2021, CoRR, abs/2107.14351 [2107.14351]
  * Cranmer (2020) Cranmer, M. 2020, PySR: Fast & Parallelized Symbolic Regression in Python/Julia
  * Cranmer (2023) Cranmer, M. 2023, arXiv e-prints, arXiv:2305.01582
  * Cranmer et al. (2020) Cranmer, M., Sanchez Gonzalez, A., Battaglia, P., et al. 2020, Discovering Symbolic Models from Deep Learning with Inductive Biases
  * David (1989) David, E. 1989, Genetic Algorithms in Search, Optimization and Machine Learning (Addison-Wesley)
  * de Franca & Aldeia (2021) de Franca, F. O. & Aldeia, G. S. I. 2021, Evolutionary Computation, 29, 367
  * de Franca & Kronberger (2023) de Franca, F. O. & Kronberger, G. 2023, in Proceedings of the Genetic and Evolutionary Computation Conference, GECCO ’23 (New York, NY, USA: Association for Computing Machinery), 1064–1072
  * Delgado et al. (2022) Delgado, A. M., Wadekar, D., Hadzhiyska, B., et al. 2022, MNRAS, 515, 2733
  * Desmond et al. (2023) Desmond, H., Bartlett, D. J., & Ferreira, P. G. 2023, MNRAS, 521, 1817
  * Diemer (2018) Diemer, B. 2018, ApJS, 239, 35
  * Eisenstein & Hu (1998) Eisenstein, D. J. & Hu, W. 1998, ApJ, 496, 605
  * Eisenstein & Hu (1999) Eisenstein, D. J. & Hu, W. 1999, ApJ, 511, 5
  * Euclid Collaboration et al. (2021) Euclid Collaboration et al. 2021, MNRAS, 505, 2840
  * Hahn et al. (2023) Hahn, O., List, F., & Porqueres, N. 2023, arXiv e-prints, arXiv:2311.03291
  * Haupt & Haupt (2004) Haupt, R. & Haupt, S. 2004, Practical genetic algorithms, 2nd edn. (Wyley)
  * Heitmann et al. (2014) Heitmann, K., Lawrence, E., Kwan, J., Habib, S., & Higdon, D. 2014, ApJ, 780, 111
  * Jin et al. (2019) Jin, Y., Fu, W., Kang, J., Guo, J., & Guo, J. 2019, arXiv e-prints, arXiv:1910.08892
  * Kamerkar et al. (2023) Kamerkar, A., Nesseris, S., & Pinol, L. 2023, Phys. Rev. D, 108, 043509
  * Kammerer et al. (2021) Kammerer, L., Kronberger, G., Burlacu, B., et al. 2021, arXiv e-prints, arXiv:2109.13895
  * Kommenda et al. (2020) Kommenda, M., Burlacu, B., Kronberger, G., & Affenzeller, M. 2020, Genetic Programming and Evolvable Machines, 21, 471–501
  * La Cava et al. (2019) La Cava, W., Helmuth, T., Spector, L., & Moore, J. H. 2019, Evolutionary Computation, 27, 377
  * Landajuela et al. (2022) Landajuela, M., Lee, C. S., Yang, J., et al. 2022, A Unified Framework for Deep Symbolic Regression
  * Laumanns et al. (2002) Laumanns, M., Thiele, L., Deb, K., & Zitzler, E. 2002, Evol. Comput., 10, 263–282
  * Lemos et al. (2023) Lemos, P., Jeffrey, N., Cranmer, M., Ho, S., & Battaglia, P. 2023, Machine Learning: Science and Technology, 4, 045002
  * Levenberg (1944) Levenberg, K. 1944, Quarterly of Applied Mathematics, 2, 164
  * Lewis et al. (2000) Lewis, A., Challinor, A., & Lasenby, A. 2000, ApJ, 538, 473
  * Marquardt (1963) Marquardt, D. W. 1963, Journal of the Society for Industrial and Applied Mathematics, 11, 431
  * McConaghy (2011) McConaghy, T. 2011, FFX: Fast, Scalable, Deterministic Symbolic Regression Technology, ed. R. Riolo, E. Vladislavleva, & J. H. Moore (New York, NY: Springer New York), 235–260
  * Miniati & Gregori (2022) Miniati, F. & Gregori, G. 2022, Scientific Reports, 12, 11709
  * Mootoovaloo et al. (2022) Mootoovaloo, A., Jaffe, A. H., Heavens, A. F., & Leclercq, F. 2022, Astronomy and Computing, 38, 100508
  * Petersen et al. (2021) Petersen, B. K., Larma, M. L., Mundhenk, T. N., et al. 2021, in International Conference on Learning Representations
  * Planck Collaboration (2020) Planck Collaboration. 2020, A&A, 641, A6
  * René Broløs et al. (2021) René Broløs, K., Vieira Machado, M., Cave, C., et al. 2021, arXiv e-prints, arXiv:2104.05417
  * Rivero et al. (2022) Rivero, D., Fernandez-Blanco, E., & Pazos, A. 2022, Expert Systems with Applications, 198, 116712
  * Schmidt & Lipson (2009) Schmidt, M. & Lipson, H. 2009, Science, 324, 81
  * Schmidt & Lipson (2011) Schmidt, M. & Lipson, H. 2011, Age-Fitness Pareto Optimization, ed. R. Riolo, T. McConaghy, & E. Vladislavleva (New York, NY: Springer New York), 129–146
  * Smith et al. (2003) Smith, R. E., Peacock, J. A., Jenkins, A., et al. 2003, MNRAS, 341, 1311
  * Sousa et al. (2023) Sousa, T., Bartlett, D. J., Desmond, H., & Ferreira, P. G. 2023, arXiv e-prints, arXiv:2310.16786
  * Spurio Mancini et al. (2022) Spurio Mancini, A., Piras, D., Alsing, J., Joachimi, B., & Hobson, M. P. 2022, MNRAS, 511, 1771
  * Tenachi et al. (2023) Tenachi, W., Ibata, R., & Diakogiannis, F. I. 2023, arXiv e-prints, arXiv:2303.03192
  * Turing (1950) Turing, A. M. 1950, Mind, LIX, 433
  * Udrescu et al. (2020) Udrescu, S.-M., Tan, A., Feng, J., et al. 2020, AI Feynman 2.0: Pareto-optimal symbolic regression exploiting graph modularity
  * Udrescu & Tegmark (2020) Udrescu, S.-M. & Tegmark, M. 2020, Science Advances, 6, eaay2631
  * Virgolin et al. (2021a) Virgolin, M., Alderliesten, T., Witteveen, C., & Bosman, P. A. N. 2021a, Evolutionary Computation, 29, 211
  * Virgolin et al. (2021b) Virgolin, M., Alderliesten, T., Witteveen, C., & Bosman, P. A. N. 2021b, Evolutionary Computation, 29, 211
  * Wadekar et al. (2022) Wadekar, D., Thiele, L., Hill, J. C., et al. 2022, arXiv e-prints, arXiv:2209.02075
  * Winther et al. (2019) Winther, H. A., Casas, S., Baldi, M., et al. 2019, Phys. Rev. D, 100, 123540
  * Worm & Chiu (2013) Worm, T. & Chiu, K. 2013, in Proceedings of the 15th Annual Conference on Genetic and Evolutionary Computation, GECCO ’13 (New York, NY, USA: Association for Computing Machinery), 1021–1028
  * Zennaro et al. (2023) Zennaro, M., Angulo, R. E., Pellejero-Ibáñez, M., et al. 2023, MNRAS, 524, 2407

##  Appendix A Most accurate analytic expression found for linear power
spectrum

The expression we report for an analytic approximation for the linear matter
power spectrum (Eq. 6) is not the most accurate one found, but the one which
we deemed to appropriately balance accuracy, simplicity, and interpretability.
It may be desirable to have a more accurate symbolic expression if
interpretability is not a concern. In this case one may wish to use the most
accurate equation found, which is

$\begin{split}100\log F&\approx c_{0}k+c_{1}\left(\Omega_{\rm
b}c_{2}-\frac{c_{3}k}{\sqrt{c_{4}+k^{2}}}\right)\left(\frac{c_{34}\left(c_{35}k\right)^{-c_{36}k}}{\sqrt{c_{39}+\left(-\Omega_{\rm
b}+\Omega_{\rm m}c_{37}-c_{38}h\right)^{2}}}-\cos{\left(\Omega_{\rm
m}c_{32}-c_{33}k\right)}\right)\\\
&\times\Bigg{(}\frac{c_{17}\left(c_{25}k\right)^{-c_{26}k}\left(\left(\Omega_{\rm
b}c_{18}+\Omega_{\rm m}c_{19}-c_{20}h\right)\cos{\left(\Omega_{\rm
m}c_{21}-c_{22}k\right)}+\cos{\left(c_{23}k-c_{24}\right)}\right)}{\sqrt{c_{31}+\left(\frac{c_{27}\left(-\Omega_{\rm
m}c_{28}+c_{29}k\right)}{\sqrt{c_{30}+k^{2}}}-k\right)^{2}}}\\\
&-\frac{c_{5}\left(\Omega_{\rm
m}c_{12}+c_{13}k\right)^{-c_{14}k}\left(\Omega_{\rm
m}c_{6}-c_{7}k+\left(\Omega_{\rm b}c_{8}-c_{9}k\right)\cos{\left(\Omega_{\rm
m}c_{10}-c_{11}k\right)}\right)}{\sqrt{c_{16}+\left(\Omega_{\rm
b}c_{15}+k\right)^{2}}}\Bigg{)}\\\ &-c_{40}\left(\Omega_{\rm
m}c_{41}-c_{42}h+c_{43}k+\frac{c_{44}k}{\sqrt{c_{45}+k^{2}}\sqrt{c_{47}+\left(-\Omega_{\rm
m}-c_{46}h\right)^{2}}}-\frac{c_{48}\left(\Omega_{\rm
m}c_{49}+c_{50}k\right)}{\sqrt{c_{51}+k^{2}}}\right)\cos{\left(\frac{c_{52}k}{\sqrt{c_{53}+k^{2}}\sqrt{c_{55}+\left(\Omega_{\rm
m}c_{54}-k\right)^{2}}}\right)}\\\ &-c_{56}-\frac{c_{57}\left(\Omega_{\rm
m}c_{67}+c_{68}k\right)^{-c_{69}k}\left(\Omega_{\rm
m}c_{58}-c_{59}k+\left(-\Omega_{\rm b}c_{60}-\Omega_{\rm
m}c_{61}+c_{62}h\right)\cos{\left(\Omega_{\rm
m}c_{63}-c_{64}k\right)}+\cos{\left(c_{65}k-c_{66}\right)}\right)}{\sqrt{\frac{c_{70}\left(\Omega_{\rm
b}+\frac{c_{71}h}{\left(c_{72}+k^{2}\right)^{0.5}}\right)^{2}}{c_{73}+k^{2}}+1.0}}.\end{split}$
(7)

Note that this equation has 73 parameters, which is approximately twice as
many as Eq. 6, yet one only gains a factor of two in the fractional root mean
squared error. The best-fit parameter values are reported in Table 3. We note
that this function is the direct output of operon and is thus over-
parameterised so that some simplification could be applied. For example, one
only needs two of $c_{1}$, $c_{2}$ and $c_{3}$ as these only appear as
$c_{1}c_{2}$ and $c_{1}c_{3}$. Since we only provide this expression as a
precise emulator and do not attempt to interpret its terms, we choose not to
apply any simplifications (although see de Franca & Kronberger (2023) for an
automated method to do this).

Parameter | Value | Parameter | Value | Parameter | Value | Parameter | Value
---|---|---|---|---|---|---|---
$c_{0}$ | 5.1439 | $c_{19}$ | 19.855 | $c_{38}$ | 0.0177 | $c_{57}$ | 0.867
$c_{1}$ | 0.867 | $c_{20}$ | 15.939 | $c_{39}$ | 0.0146 | $c_{58}$ | 2.618
$c_{2}$ | 8.52 | $c_{21}$ | 9.547 | $c_{40}$ | 0.867 | $c_{59}$ | 2.1
$c_{3}$ | 0.2920 | $c_{22}$ | 97.34 | $c_{41}$ | 32.371 | $c_{60}$ | 114.391
$c_{4}$ | 0.0310 | $c_{23}$ | 94.83 | $c_{42}$ | 7.058 | $c_{61}$ | 13.968
$c_{5}$ | 0.0033 | $c_{24}$ | 1.881 | $c_{43}$ | 6.075 | $c_{62}$ | 11.133
$c_{6}$ | 240.234 | $c_{25}$ | 3.945 | $c_{44}$ | 16.311 | $c_{63}$ | 4.205
$c_{7}$ | 682.449 | $c_{26}$ | 11.151 | $c_{45}$ | 0.0025 | $c_{64}$ | 100.376
$c_{8}$ | 2061.023 | $c_{27}$ | 0.0004 | $c_{46}$ | 0.1632 | $c_{65}$ | 106.993
$c_{9}$ | 6769.493 | $c_{28}$ | 26.822 | $c_{47}$ | 0.0771 | $c_{66}$ | 3.359
$c_{10}$ | 7.125 | $c_{29}$ | 230.12 | $c_{48}$ | 0.0522 | $c_{67}$ | 1.539
$c_{11}$ | 108.136 | $c_{30}$ | 0.0009 | $c_{49}$ | 22.722 | $c_{68}$ | 1.773
$c_{12}$ | 6.2 | $c_{31}$ | 1.0796$\times 10^{-5}$ | $c_{50}$ | 774.688 | $c_{69}$ | 18.983
$c_{13}$ | 2.882 | $c_{32}$ | 3.162 | $c_{51}$ | 0.0027 | $c_{70}$ | 0.3838
$c_{14}$ | 59.585 | $c_{33}$ | 99.918 | $c_{52}$ | 1.0337 | $c_{71}$ | 0.0024
$c_{15}$ | 0.1384 | $c_{34}$ | 0.1210 | $c_{53}$ | 0.0058 | $c_{72}$ | 1.2865$\times 10^{-7}$
$c_{16}$ | 1.0825$\times 10^{-5}$ | $c_{35}$ | 0.495 | $c_{54}$ | 0.2472 | $c_{73}$ | 2.0482$\times 10^{-8}$
$c_{17}$ | 0.0033 | $c_{36}$ | 12.091 | $c_{55}$ | 0.29 |  |
$c_{18}$ | 85.791 | $c_{37}$ | 0.6070 | $c_{56}$ | 0.241 |  |
Table 3: Best-fit parameters for the most accurate linear matter power
spectrum emulator found, reported in Appendix A. Although units are excluded
in this table, the units for each parameter are easily obtained by noting that
these are defined assuming that $k$ is measured in $h\,{\rm Mpc^{-1}}$.