\section{Introduction}
Deep reinforcement learning (RL) enables autonomous learning of diverse and complex tasks with rich sensory inputs, temporally extended goals, and challenging dynamics, such as discrete game-playing domains~\citep{mnih2013playing,silver2016mastering}, and continuous control domains including locomotion~\citep{schulman2015trust,heess2017emergence} and manipulation~\citep{rajeswaran2017learning,kalashnikov2018qt,gu2017deep}.
Most of the deep RL approaches learn a Q-function or a policy that are directly optimized for the training task, which limits their generalization to new scenarios.
In contrast, MBRL methods~\citep{li2004iterative,deisenroth2011pilco,watter2015embed} can acquire dynamics models that may be utilized to perform unseen tasks at test time.
While this capability has been demonstrated in some of the recent works~\citep{levine2016end,nagabandi2018neural,chua2018deep,kurutach2018model,ha2018recurrent}, learning an accurate global model that works for all state-action pairs can be exceedingly challenging, especially for high-dimensional system with complex and discontinuous dynamics.
The problem is further exacerbated as the learned global model has limited generalization outside of the state distribution it was trained on and exploring the whole state space is generally infeasible. Can we retain the flexibility of model-based RL, while using model-free RL to acquire proficient low-level behaviors under complex dynamics?

While learning a global dynamics model that captures all the different behaviors for the entire state-space can be extremely challenging, learning a model for a specific behavior that acts only in a small part of the state-space can be much easier.
For example, consider learning a model for dynamics of all gaits of a quadruped versus a model which only works for a specific gait.
If we can learn many such behaviors and their corresponding dynamics, we can leverage model-predictive control to plan in the \textit{behavior space}, as opposed to planning in the action space. The question then becomes: how do we acquire such behaviors, considering that behaviors could be random and unpredictable? To this end, we propose \textit{Dynamics-Aware Discovery of Skills} (DADS), an unsupervised RL framework for learning low-level skills using model-free RL with the explicit aim of making model-based control easy. Skills obtained using DADS are directly optimized for \emph{predictability}, providing a better representation on top of which predictive models can be learned. Crucially, the skills do not require any supervision to learn, and are acquired entirely through autonomous exploration. This means that the repertoire of skills and their predictive model are learned before the agent has been tasked with any goal or reward function. When a task is provided at test-time, the agent utilizes the previously learned skills and model to immediately perform the task without any further training.

The key contribution of our work is an unsupervised reinforcement learning algorithm, DADS, grounded in mutual-information-based exploration. We demonstrate that our objective can embed learned primitives in continuous spaces, which allows us to learn a large, diverse set of skills.
Crucially, our algorithm also learns to model the dynamics of the skills, which enables the use of model-based planning algorithms for downstream tasks. We adapt the conventional model predictive control algorithms to plan in the space of primitives, and demonstrate that we can compose the learned primitives to solve downstream tasks without any additional training.
\section{Preliminaries}
Mutual information can been used as an objective to encourage exploration in reinforcement learning \citep{DBLP:journals/corr/HouthooftCDSTA16, mohamed2015variational}. According to its definition, $\mathcal{I}(X;Y) = \mathcal{H}(X) - \mathcal{H}(X \mid Y)$, maximizing mutual information $\mathcal{I}$ with respect to $Y$ amounts to maximizing the entropy $\mathcal{H}$ of $X$ while minimizing the conditional entropy $\mathcal{H}(X \mid Y)$.
In the context of RL, $X$ is usually a function of the state and $Y$ a function of actions. Maximizing this objective encourages the state entropy to be high, making the underlying policy to be exploratory.
Recently, multiple works \citep{eysenbach2018diversity, gregor2016variational, achiam2018variational} apply this idea to learn diverse skills which maximally cover the state space.

To leverage planning-based control, MBRL estimates the true dynamics of the environment by learning a model $\hat{p}(s' \mid s, a)$.
This allows it to predict a trajectory of states $\hat{\tau}_H = (s_t ,\hat{s}_{t+1}, \ldots \hat{s}_{t+H})$  resulting from a sequence of actions without any additional interaction with the environment. While model-based RL methods have been demonstrated to be sample efficient compared to their model-free counterparts, learning an effective model for the whole state-space is challenging. An open-problem in model-based RL is to incorporate temporal abstraction in model-based control, to enable high-level planning and move-away from planning at the granular level of actions.

These seemingly unrelated ideas can be combined into a single optimization scheme, where we first discover skills (and their models) without any extrinsic reward and then compose these skills to optimize for the task defined at test time using model-based planning.
At train time, we assume a Markov Decision Process (MDP) $\mathcal{M}_1 \equiv (\mathcal{S}, \mathcal{A}, p)$. The state space $\mathcal{S}$ and action space $\mathcal{A}$ are assumed to be continuous, and the $\mathcal{A}$ bounded. We assume the transition dynamics $p$ to be stochastic, such that $p: \mathcal{S} \times \mathcal{A} \times \mathcal{S} \mapsto [0, \infty)$. We learn a skill-conditioned policy $\pi(a \mid s, z)$, where the skills $z$ belongs to the space $\mathcal{Z}$, detailed in Section~\ref{sec:DADS}. We assume that the skills are sampled from a prior $p(z)$ over $\mathcal{Z}$. We simultaneously learn a skill-conditioned transition function $q(s' \mid s, z)$, coined as \textit{skill-dynamics}, which predicts the transition to the next state $s'$ from the current state $s$ for the skill $z$ under the given dynamics $p$. At test time, we assume an MDP $\mathcal{M}_2 \equiv (\mathcal{S}, \mathcal{A}, p, r)$, where $\mathcal{S}, \mathcal{A}, p$ match those defined in $\mathcal{M}_1$, and the reward function $r: \mathcal{S} \times \mathcal{A} \mapsto (-\infty, \infty)$. We plan in $\mathcal{Z}$ using $q(s' \mid s, z)$ to compose the learned skills $z$ for optimizing $r$ in $\mathcal{M}_2$, which we detail in Section \ref{sec:planning}.
\section{Dynamics-Aware Discovery of Skills (DADS)}
\label{sec:DADS}
\begin{figure}[!hb]
    \centering
    \begin{minipage}{0.5\textwidth}
    \includegraphics[width=.9\textwidth]{images/skill_learning.png}
    \end{minipage}
    \begin{minipage}{0.4\textwidth}
    \small
    \begin{algorithm}[H]
    \label{algorithm:usl}
    \SetAlgoLined
     Initialize $\pi, q_\phi$\;
     \While{\textit{not converged}}{
      Sample a skill $z \sim p(z)$ every episode\;
      Collect new $M$ on-policy samples\;
      Update $q_\phi$ using $K_1$ steps of gradient descent on $M$ transitions\;
      Compute $r_z(s, a, s')$ for $M$ transitions\;
      Update $\pi$ using any RL algorithm\;
     }
     \caption{\small Dynamics-Aware Discovery of Skills (DADS)}
    \end{algorithm}
    \end{minipage}

    \caption{\small The agent $\pi$ interacts with the environment to produce a transition $s \rightarrow s'$. Intrinsic reward is computed  by computing the transition probability under $q$ for the current skill $z$, compared to random samples from the prior $p(z)$. The agent maximizes the intrinsic reward computed for a batch of episodes, while $q$ maximizes the log-probability of the actual transitions of $(s, z) \rightarrow s'$.}
    \label{Fig:skill_discovery_summary}
\end{figure}
We use the information theoretic paradigm of mutual information to obtain our unsupervised skill discovery algorithm. In particular, we propose to maximize the mutual information between the next state $s'$ and current skill $z$ conditioned on the current state $s$.
\begin{align}
\label{eq:mi_objective}
\mathcal{I}(s';z \mid s) &= \mathcal{H}(z\mid s) - \mathcal{H}(z \mid s', s)\\
\label{eq:mi_mbrl_decomposition}
&= \mathcal{H}(s' \mid s) - \mathcal{H}(s' \mid s, z)
\end{align}
Mutual information in Equation~\ref{eq:mi_objective} quantifies how much can be known about $s'$ given  $z$ and $s$, or symmetrically, $z$ given the transition from $s \rightarrow s'$. From Equation~\ref{eq:mi_mbrl_decomposition}, maximizing this objective corresponds to maximizing the diversity of transitions produced in the environment, that is denoted by the entropy $\mathcal{H}(s' \mid s)$, while making $z$ informative about the next state $s'$ by minimizing the entropy $\mathcal{H}(s' \mid s, z)$. Intuitively, skills $z$ can be interpreted as abstracted action sequences which are identifiable by the transitions generated in the environment (and not just by the current state). Thus, optimizing this mutual information can be understood as encoding a diverse set of skills in the latent space $\mathcal{Z}$, while making the transitions for a given $z \in \mathcal{Z}$ predictable. We use the entropy-decomposition in Equation~\ref{eq:mi_mbrl_decomposition}
 to connect this objective with model-based control.

We want to optimize the our skill-conditioned controller $\pi(a \mid s, z)$ such that the latent space ${z \sim p(z)}$ is maximally informative about the transitions $s \rightarrow s'$. Using the definition of conditional mutual information, we can rewrite Equation~\ref{eq:mi_mbrl_decomposition} as:
\begin{align}
\mathcal{I}(s'; z \mid s) = \int p(z, s, s') \log\frac{p(s' \mid s, z)}{p(s' \mid s)} ds' ds dz
\end{align}
We assume the following generative model: $p(z, s, s') = p(z) p(s \mid z) p(s' \mid s, z)$, where $p(z)$ is user specified prior over $\mathcal{Z}$, $p(s | z)$ denotes the stationary state-distribution induced by $\pi(a \mid s, z)$ for a skill $z$ and $p(s' \mid s, z)$ denotes the transition distribution under skill $z$. Note, $p(s' \mid s, z) = \int p(s' \mid s, a) \pi(a \mid s, z) da$ is intractable to compute because the underlying dynamics are unknown. However, we can variationally lower bound the objective as follows:
\begin{align}
\mathcal{I}(s';z \mid s) &= \mathbb{E}_{z, s, s' \sim p} \Big[\log \frac{p(s' \mid s, z)}{p(s' \mid s)}\Big]\nonumber\\
&= \mathbb{E}_{z, s, s' \sim p}\Big[\log\frac{q_\phi(s' \mid s, z)}{p(s' \mid s)}\Big] + \mathbb{E}_{s, z \sim p}\Big[\mathcal{D}_{KL}(p(s' \mid s, z) \mid\mid q_\phi(s' \mid s, z))\Big]\label{eq:alternate_opt}\nonumber\\
&\geq \mathbb{E}_{z, s, s' \sim p}\Big[\log\frac{q_\phi(s' \mid s, z)}{p(s' \mid s)}\Big]
\end{align}
where we have used the non-negativity of KL-divergence, that is $\mathcal{D}_{KL} \geq 0$. Note, skill-dynamics $q_\phi$ represents the variational approximation for the transition function $p(s' \mid s, z)$, which enables model-based control as described in Section~\ref{sec:planning}. Equation~\ref{eq:alternate_opt} suggests an alternating optimization between $q_\phi$ and $\pi$, summarized in Algorithm~\ref{algorithm:usl}.
In every iteration: \\
(\textit{Tighten variational lower bound}) We minimize $D_{KL}(p(s' \mid s, z) \mid\mid q_\phi(s' \mid s, z))$ with respect to the parameters $\phi$ on $z, s \sim p$ to tighten the lower bound. For general function approximators like neural networks, we can write the gradient for $\phi$ as follows:
\begin{align}
\nabla_\phi \mathbb{E}_{s, z} [\mathcal{D}_{KL}(p(s' \mid s, z )\mid\mid q_\phi(s' \mid s, z))] &= \nabla_\phi \mathbb{E}_{z, s, s'} \Big[\log \frac{p(s' \mid s, z)}{q_\phi(s' \mid s, z)}\Big] \nonumber\\
&= -\mathbb{E}_{z, s, s'} \Big[ \nabla_\phi \log q_\phi(s' \mid s, z)\Big]
\end{align}
which corresponds to maximizing the likelihood of the samples from $p$ under $q_\phi$.

(\textit{Maximize approximate lower bound}) After fitting $q_\phi$, we can optimize $\pi$ to maximize $\mathbb{E}_{z, s, s'} [ \log q_\phi(s' \mid s, z) - \log p(s' \mid s)]$. Note, this is a reinforcement-learning style optimization with a reward function $\log q_\phi(s' \mid s, z) - \log p(s' \mid s)$. However, $\log p(s' \mid s)$ is intractable to compute, so we approximate the reward function for $\pi$:
\begin{align}
r_z(s, a, s') = \log{\frac{q_\phi(s' \mid s, z)}{\sum_{i=1}^L q_\phi(s' \mid s, z_i)}} + \log{L}, \quad z_i \sim p(z). \label{eq:final_reward}
\end{align}
The approximation is motivated as follows: $p(s' \mid s) = \int p(s' \mid s, z) p(z | s) dz \approx \int q_\phi(s' \mid s, z) p(z) dz \approx \frac{1}{L}\sum_{i=1}^L q_\phi(s' \mid s, z_i)$ for $z_i \sim p(z)$, where $L$ denotes the number of samples from the prior. We are using the marginal of variational approximation $q_\phi$ over the prior $p(z)$ to approximate the marginal distribution of transitions. We discuss this approximation in Appendix~\ref{appendix:reward_approximation}. Note, the final reward function $r_z$ encourages the policy $\pi$ to produce transitions that are (a) predictable under $q_\phi$ (\textit{predictability}) and (b) different from the transitions produced under $z_i \sim p(z)$ (\textit{diversity}).

To generate samples from $p(z, s, s')$, we use the rollouts from the current policy $\pi$ for multiple samples $z \sim p(z)$ in an episodic setting for a fixed horizon $T$. We also introduce entropy regularization for $\pi(a \mid s, z)$, which encourages the policy to discover action-sequences with similar state-transitions and to be clustered under the same skill $z$, making the policy robust besides encouraging exploration \citep{haarnoja2018soft}. The use of entropy regularization can be justified from an information bottleneck perspective as discussed for Information Maximization algorithm in \citep{mohamed2015variational}. This is even more extensively discussed from the graphical model perspective in Appendix~\ref{appendix:infobot}, which connects unsupervised skill discovery and information bottleneck literature, while also revealing the temporal nature of skills $z$. Details corresponding to implementation and hyperparameters are discussed in Appendix~\ref{appendix:implementation}.

\section{Planning using Skill Dynamics}
\label{sec:planning}
\begin{figure}[!htb]

    \centering
    \begin{minipage}{0.55\textwidth}
    \centering
    \includegraphics[width=\textwidth]{images/model_predictive_control.png}
    \end{minipage}
    \begin{minipage}{0.4\textwidth}
    \begin{algorithm}[H]
    \label{algorithm:planning}
    \small
    \SetAlgoLined
    $s\gets s_0$\;
    Initialize parameters $\mu_1, \ldots \mu_{H_P}$\;
    \For{$i\gets1$ \KwTo $H_E / H_Z$}{
        \For{$j\gets1$ \KwTo $R$}{
            $\{z_{i}, \ldots z_{i + H_P - 1}\}_{k=1}^K \sim \mathcal{N}_i, \ldots \mathcal{N}_{i+ H_P - 1}$ \;
            Compute $r_{env}$ for $\{z_{i}, \ldots z_{i + H_P - 1}\}_{k=1}^K$\;
            Update $\mu_i, \ldots, \mu_{i + H_P - 1}$\;
        }
        Sample $z_i$ from $\mathcal{N}_i$\;
        Execute $\pi(a | s, z_i)$ for $H_Z$ steps\;
        Initialize $\mu_{i+H_P}$\;
    }
    \caption{\small Latent Space Planner}
    \end{algorithm}
    \end{minipage}

    \caption{\small At test time, the planner executes simulates the transitions in environment using skill-dynamics $q$, and updates the distribution of plans according to the computed reward on the simulated trajectories. After a few updates to the plan, the first primitive is executed in the environment using the learned agent $\pi$.}
    \label{Fig:planning_overview}
\end{figure}

Given the learned skills $\pi(a \mid s, z)$ and their respective skill-transition dynamics $q_\phi(s' \mid s, z)$, we can perform model-based planning in the latent space $\mathcal{Z}$ to optimize for a reward $r$ that is given to the agent at test time.
Note, that this essentially allows us to perform zero-shot planning given the unsupervised pre-training procedure described in Section~\ref{sec:DADS}.

In order to perform planning, we employ the model-predictive-control (MPC) paradigm~\cite{garcia1989model}, which in a standard setting generates a set of action plans $P_k = (a_{k,1}, \ldots a_{k,H}) \sim P$ for a planning horizon~$H$.
The MPC plans can be generated due to the fact that the planner is able to simulate the trajectory $\hat{\tau}_k = (s_{k, 1}, a_{k, 1} \ldots s_{k, H+1})$ assuming access to the transition dynamics $\hat{p}(s' \mid s, a)$.
In addition, each plan computes the reward $r(\hat{\tau}_k)$ for its trajectory according to the reward function $r$ that is provided for the test-time task.
Following the MPC principle, the planner selects the best plan according to the reward function $r$ and executes its first action $a_1$.
The planning algorithm repeats this procedure for the next state iteratively until it achieves its goal.

We use a similar strategy to design an MPC planner to exploit previously-learned skill-transition dynamics $q_\phi(s' \mid s,z)$.
Note that unlike conventional model-based RL, we generate a plan $P_k = (z_{k,1}, \ldots z_{k, H_P})$ in the latent space $\mathcal{Z}$ as opposed to the action space $\mathcal{A}$ that would be used by a standard planner.
Since the primitives are temporally meaningful, it is beneficial to hold a primitive for a horizon $H_Z > 1$, unlike actions which are usually held for a single step.
Thus, effectively, the planning horizon for our latent space planner is $H = H_P \times H_Z$, enabling longer-horizon planning using fewer primitives.
Similar to the standard MPC setting, the latent space planner simulates the trajectory $\hat{\tau}_k = (s_{k, 1}, z_{k, 1}, a_{k, 1}, s_{k, 2}, z_{k, 2}, a_{k, 2}, \ldots s_{k, H+1})$ and computes the reward $r(\hat{\tau}_k)$.
After a small number of trajectory samples, the planner selects the first latent action $z_1$ of the best plan, executes it for $H_Z$ steps in the environment, and the repeats the process until goal completion.

The latent planner $P$ maintains a distribution of latent plans, each of length $H_P$.
Each element in the sequence represents the distribution of the primitive to be executed at that time step. For continuous spaces, each element of the sequence can be modelled using a normal distribution, $\mathcal{N}(\mu_1, \Sigma), \ldots \mathcal{N}(\mu_{H_P}, \Sigma)$.
We refine the planning distributions for $R$ steps, using $K$ samples of latent plans $P_k$, and compute the $r_k$ for the simulated trajectory $\hat{\tau}_k$. The update for the parameters follows that in Model Predictive Path Integral (MPPI) controller~\cite{williams2016aggressive}:
\begin{align}
    \mu_i = \sum_{k=1}^K\frac{\exp(\gamma r_k)}{\sum_{p=1}^K \exp(\gamma r_p)} z_{k, i} \quad \forall i = 1, \ldots H_P
\end{align}
While we keep the covariance matrix of the distributions fixed, it is possible to update that as well as shown in~\cite{williams2016aggressive}. We show an overview of the planning algorithm in Figure \ref{Fig:planning_overview}, and provide more implementation details in Appendix \ref{appendix:implementation}.
\section{Related Work}
Central to our method is the concept of skill discovery via mutual information maximization. This principle, proposed in prior work that utilized purely model-free unsupervised RL methods~\citep{daniel2012hierarchical,florensa2017stochastic,eysenbach2018diversity,gregor2016variational, warde2018unsupervised, thomas2018disentangling}, aims to learn diverse skills via a discriminability objective: a good set of skills is one where it is easy to distinguish the skills from each other, which means they perform distinct tasks and cover the space of possible behaviors.
Building on this prior work, we distinguish our skills based on how they modify the original uncontrolled dynamics of the system.
This simultaneously encourages the skills to be both \emph{diverse} and \emph{predictable}. We also demonstrate that constraining the skills to be predictable makes them more amenable for hierarchical composition and thus, more useful on downstream tasks.

Another line of work that is conceptually close to our method copes with intrinsic motivation~\citep{oudeyer2009intrinsic, oudeyer2007intrinsic, schmidhuber2010formal} which is used to drive the agent's exploration. Examples of such works include empowerment~\cite{klyubin2005empowerment,mohamed2015variational}, count-based exploration~\cite{bellemare2016unifying,oh2015action,tang2017exploration, NIPS2017_6851}, information gain about agent's dynamics~\cite{StadieLA15} and forward-inverse dynamics models~\cite{pathakICMl17curiosity}. While our method uses an information-theoretic objective that is similar to these approaches, it is used to learn a variety of skills that can be directly used for model-based planning, which is in contrast to learning a better exploration policy for a single skill.

The skills discovered using our approach can also provide extended actions and temporal abstraction, which enable more efficient exploration for the agent to solve various tasks, reminiscent of hierarchical RL (HRL) approaches. This ranges from the classic option-critic architecture~\citep{sutton1999between,stolle2002learning, perkins1999using} to some of the more recent work~\citep{bacon2017option,vezhnevets2017feudal,nachum2018data, hausman2018learning}.
However, in contrast to end-to-end HRL approaches~\citep{heess2016learning,peng2017deeploco}, we can leverage a stable, two-phase learning setup.
The primitives learned through our method provide action and temporal abstraction, while planning with skill-dynamics enables hierarchical composition of these primitives, bypassing many problems of end-to-end HRL.

In the second phase of our approach, we use the learned skill-transition dynamics models to perform model-based planning - an idea that has been explored numerous times in the literature.
Model-based reinforcement learning has been traditionally approached with methods that are well-suited for low-data regimes such as Gaussian Processes~\citep{rasmussen2003gaussian} showing significant data-efficiency gains over model-free approaches~\citep{deisenroth2013gaussian,kamthe2017data,kocijan2004gaussian,ko2007gaussian}.
More recently, due to the challenges of applying these methods to high-dimensional state spaces, MBRL approaches employs Bayesian deep neural networks~\citep{nagabandi2018neural,chua2018deep,gal2016improving,fu2016one,lenz2015deepmpc} to learn dynamics models.
In our approach, we take advantage of the deep dynamics models that are conditioned on the skill being executed, simplifying the modelling problem.
In addition, the skills themselves are being learned with the objective of being predictable, further assists with the learning of the dynamics model.
There also have been multiple approaches addressing the planning component of MBRL including linear controllers for local models~\citep{levine2016end, kumar2016optimal, chebotar2017combining}, uncertainty-aware~\citep{chua2018deep,gal2016improving} or deterministic planners~\citep{nagabandi2018neural} and stochastic optimization methods~\citep{williams2016aggressive}.
The main contribution of our work lies in discovering model-based skill primitives that can be further combined by a standard model-based planner, therefore we take advantage of an existing planning approach - Model Predictive Path Integral~\citep{williams2016aggressive} that can leverage our pre-trained setting.
\vspace{-4.3mm}
\section{Experiments}
Through our experiments, we aim to demonstrate that: (a) DADS as a general purpose skill discovery algorithm can scale to high-dimensional problems; (b) discovered skills are amenable to hierarchical composition and; (c) not only is planning in the learned latent space feasible, but it is competitive to strong baselines. In Section~\ref{sec:qualitative}, we provide visualizations and qualitative analysis of the skills learned using DADS. We demonstrate in Section~\ref{sec:variance_primitives} and Section~\ref{sec:hrl_on_skills} that optimizing the primitives for predictability renders skills more amenable to temporal composition that can be used for Hierarchical RL
We benchmark against state-of-the-art model-based RL baseline in Section~\ref{sec:mbrl}, and against goal-conditioned RL in Section~\ref{sec:gcrl}.
\subsection{Qualitative Analysis}
\label{sec:qualitative}

\begin{figure}[!htb]
    \begin{minipage}{0.32\textwidth}
    \centering
    \includegraphics[width=\textwidth]{images/humanoid_fused.png}
    \end{minipage}
    \hfill
    \begin{minipage}{0.33\textwidth}
    \centering
    \includegraphics[width=\textwidth]{images/ant_skills.png}
    \end{minipage}
    \hfill
    \begin{minipage}{0.32\textwidth}
    \centering
    \includegraphics[width=\textwidth]{images/walker_fused.png}
    \end{minipage}

    \caption{\small Skills learned on different MuJoCo environments in the OpenAI gym. DADS can discover diverse skills without any extrinsic rewards, even for problems with high-dimensional state and action spaces.}
    \label{Fig:skills}
\end{figure}

In this section, we provide a qualitative discussion of the unsupervised skills learned using DADS. We use the MuJoCo environments \citep{todorov2012mujoco} from the OpenAI gym as our test-bed \citep{DBLP:journals/corr/BrockmanCPSSTZ16}. We find that our proposed algorithm can learn diverse skills without any reward, even in problems with high-dimensional state and actuation, as illustrated in Figure \ref{Fig:skills}. DADS can discover primitives for Half-Cheetah to run forward and backward with multiple different gaits, for Ant to navigate the environment using diverse locomotion primitives and for Humanoid to walk using stable locomotion primitives with diverse gaits and direction. The videos of the discovered primitives are available at: \url{https://sites.google.com/view/dads-skill}

Qualitatively, we find the skills discovered by DADS to be predictable and stable, in line with implicit constraints of the proposed objective. While the Half-Cheetah will learn to run in both backward and forward directions, DADS will disincentivize skills which make Half-Cheetah flip owing to the reduced predictability on landing. Similarly, skills discovered for Ant rarely flip over, and tend to provide stable navigation primitives in the environment. This also incentivizes the Humanoid, which is characteristically prone to collapsing and extremely unstable by design, to discover gaits which are stable for sustainable locomotion.

One of the significant advantages of the proposed objective is that it is compatible with continuous skill spaces, which has not been shown in prior work on skill discovery~\citep{eysenbach2018diversity}. Not only does this allow us to embed a large and diverse set of skills into a compact latent space, but also the smoothness of the learned space allows us to interpolate between behaviors generated in the environment.
We demonstrate this on the Ant environment (Figure \ref{Fig:latent_space_visualization}), where we learn two-dimensional continuous skill space with a uniform prior over $(-1, 1)$ in each dimension, and compare it to a discrete skill space with a uniform prior over $20$ skills. Similar to~\cite{eysenbach2018diversity}, we restrict the observation space of the skill-dynamics $q$ to the cartesian coordinates $(x, y)$. We hereby call this the \emph{x-y prior}, and discuss its role in Section \ref{sec:variance_primitives}.

\begin{figure}[!htb]
\centering
\begin{minipage}{0.31\textwidth}
\centering
\scriptsize Trajectories in Discrete Skill Space\par
\includegraphics[width=\textwidth]{images/traj_plot_discrete_fix1.png}
\end{minipage}
\begin{minipage}{0.31\textwidth}
\centering
\scriptsize \hspace{0.5mm} Trajectories in Continuous Skill Space\par
\includegraphics[width=\textwidth]{images/traj_plot_continuous_fix1.png}
\end{minipage}
\begin{minipage}{0.36\textwidth}
\centering
\scriptsize  \hspace{-5mm} Orientation of Ant Trajectory\par
\vspace{-0.5mm}
\includegraphics[width=\textwidth]{images/heatmap_fix4.png}
\end{minipage}
\caption{\small (Left, Centre) X-Y traces of Ant skills and (Right) Heatmap to visualize the learned continuous skill space. Traces demonstrate that the continuous space offers far greater diversity of skills, while the heatmap demonstrates that the learned space is smooth, as the orientation of the X-Y trace varies smoothly as a function of the skill.}
\label{Fig:latent_space_visualization}
\end{figure}

In Figure~\ref{Fig:latent_space_visualization}, we project the trajectories of the learned Ant skills from both discrete and continuous spaces onto the Cartesian plane. From the traces of the skills, it is clear that the continuous latent space can generate more diverse trajectories. We demonstrate in Section~\ref{sec:mbrl}, that continuous primitives are more amenable to hierarchical composition and generally perform better on downstream tasks. More importantly, we observe that the learned skill space is semantically meaningful. The heatmap in Figure~\ref{Fig:latent_space_visualization} shows the orientation of the trajectory (with respect to the $x$-axis) as a function of the skill $z \in \mathcal{Z}$, which varies smoothly as $z$ is varied, with explicit interpolations shown in Appendix~\ref{Appendix:interpolation}.
\subsection{Skill Variance Analysis}
\label{sec:variance_primitives}

\begin{figure}[!htb]
\centering
\begin{minipage}{0.48\textwidth}
\centering
\small \hspace{4mm}Standard Deviation of Trajectories\par
\vspace{-0.5mm}
\includegraphics[width=0.98\textwidth]{images/primitive_variance_y_fix3.png}
\small DADS without x-y prior\par
\includegraphics[width=0.9\textwidth]{images/multi_executions_full_obs_fix1.png}
\end{minipage}
\begin{minipage}{0.48\textwidth}
\centering
\small DIAYN with x-y prior\par
\includegraphics[width=0.9\textwidth]{images/multi_executions_diayn_xy_fix1.png}
\small DADS with x-y prior\par
\includegraphics[width=0.9\textwidth]{images/multi_executions_dads_xy_fix1.png}
\end{minipage}
\caption{\small (Top-Left) Standard deviation of Ant's position as a function of steps in the environment, averaged over multiple skills and normalized by the norm of the position. (Top-Right to Bottom-Left Clockwise) X-Y traces of skills learned using DIAYN with $x$-$y$ prior, DADS with $x$-$y$ prior and DADS without x-y prior, where the same color represents trajectories resulting from the execution of the same skill $z$ in the environment. High variance skills from DIAYN offer limited utility for hierarchical control.}
\label{Fig:variance_analysis}
\end{figure}

In an unsupervised skill learning setup, it is important to optimize the primitives to be diverse. However, we argue that diversity is not sufficient for the learned primitives to be useful for downstream tasks. Primitives must exhibit low-variance behavior, which enables long-horizon composition of the learned skills in a hierarchical setup. We analyze the variance of the $x$-$y$ trajectories in the environment, where we also benchmark the variance of the primitives learned by DIAYN~\citep{eysenbach2018diversity}. For DIAYN, we use the $x$-$y$ prior for the skill-discriminator, which biases the discovered skills to diversify in the $x$-$y$ space.
This step was necessary for that baseline to obtain a competitive set of navigation skills.
Figure~\ref{Fig:variance_analysis} (Top-Left) demonstrates that DADS, which optimizes the primitives for predictability and diversity, yields significantly lower-variance primitives when compared to DIAYN, which only optimizes for diversity. This is starkly demonstrated in the plots of X-Y traces of skills learned in different setups. Skills learned by DADS show significant control over the trajectories generated in the environment, while skills from DIAYN exhibit high variance in the environment, which limits their utility for hierarchical control. This is further demonstrated quantitatively in Section~\ref{sec:hrl_on_skills}.

While optimizing for predictability already significantly reduces the variance of the trajectories generated by a primitive, we find that using the $x$-$y$ prior with DADS brings down the skill variance even further. For quantitative benchmarks in the next sections, we assume that the Ant skills are learned using an $x$-$y$ prior on the observation space, for both DADS and DIAYN.

\subsection{Model-Based Reinforcement Learning}
\vspace{-2mm}
\label{sec:mbrl}
The key utility of learning a parametric model $q_\phi(s' | s, z)$ is to take advantage of planning algorithms for downstream tasks, which can be extremely sample-efficient. In our setup, we can solve test-time tasks in zero-shot, that is \textit{without any learning on the downstream task}. We compare with the state-of-the-art model-based RL method \citep{DBLP:journals/corr/abs-1805-12114}, which learns a dynamics model parameterized as $p(s' | s, a)$, on the task of the Ant navigating to a specified goal with a dense reward. Given a goal $g$, reward at any position $u$ is given by $r(u) = -\lVert g - u\rVert_2$. We benchmark our method against the following variants:
\begin{itemize}
\itemsep0em

    \item Random-MBRL (\textit{rMBRL}): We train the model $p(s' | s, a)$ on randomly collected trajectories, and test the zero-shot generalization of the model on a distribution of goals.
    \item Weak-oracle MBRL (\textit{WO-MBRL}): We train the model $p(s' | s, a)$ on trajectories generated by the planner to navigate to a goal, randomly sampled in every episode. The distribution of goals during training matches the distribution at test time.
    \item Strong-oracle MBRL (\textit{SO-MBRL}): We train the model $p(s' | s, a)$ on a trajectories generated by the planner to navigate to a specific goal, which is fixed for both training and test time.
\end{itemize}
Amongst the variants, only the rMBRL matches our assumptions of having an unsupervised task-agnostic training. Both WO-MBRL and SO-MBRL benefit from goal-directed exploration during training, a significant advantage over DADS, which only uses mutual-information-based exploration.

We use $\Delta = \sum_{t=1}^H\frac{-r(u)}{H\lVert g\rVert_2}$ as the metric, which represents the distance to the goal $g$ averaged over the episode (with the same fixed horizon $H$ for all models and experiments), normalized by the initial distance to the goal $g$. Therefore, lower $\Delta$ indicates better performance and $0 < \Delta \leq 1$ (assuming the agent goes closer to the goal). The test set of goals is fixed for all the methods, sampled from $[-15, 15]^2$.

Figure~\ref{Fig:mbrl_benchmark} demonstrates that the zero-shot planning significantly outperforms all model-based RL baselines, despite the advantage of the baselines being trained on the test goal(s). For the experiment depicted in Figure~\ref{Fig:mbrl_benchmark} (Right), DADS has an unsupervised pre-training phase, unlike SO-MBRL which is training directly for the task.
A comparison with Random-MBRL shows the significance of mutual-information-based exploration, especially with the right parameterization and priors. This experiment also demonstrates the advantage of learning a continuous space of primitives, which outperforms planning on discrete primitives.

\begin{figure}[!htb]
\centering
\begin{minipage}{0.48\textwidth}
\includegraphics[width=\textwidth]{images/mbrl_bar_fix4.png}
\end{minipage}
\hfill
\begin{minipage}{0.48\textwidth}
\centering
\includegraphics[width=\textwidth]{images/so_mbrl_dads_fix4.png}
\end{minipage}

\caption{\small (Left) The results of the MPPI controller on skills learned using DADS-c (continuous primitives) and DADS-d (discrete primitives) significantly outperforms state-of-the-art model-based RL. (Right) Planning for a new task does not require any additional training and outperforms model-based RL being trained for the specific task. }
\label{Fig:mbrl_benchmark}
\end{figure}

\vspace{-4mm}
\subsection{Hierarchical Control with Unsupervised Primitives}
\label{sec:hrl_on_skills}
\vspace{-2mm}
\begin{figure}[!htb]
\centering
\begin{minipage}{0.48\textwidth}
\includegraphics[width=\textwidth]{images/hierarchical_control_usl_fix4.png}
\end{minipage}
\hfill
\begin{minipage}{0.48\textwidth}
\centering
\includegraphics[width=\textwidth]{images/gcrl_benchmark_30_fix4.png}
\end{minipage}

\caption{\small (Left) A RL-trained meta-controller is unable to compose primitive learned by DIAYN to navigate Ant to a goal, while it succeeds to do so using the primitives learned by DADS. (Right) Goal-Conditioned RL (GCRL-dense/sparse) does not generalize outside its training distribution, while MPPI controller on learned skills (DADS-dense/sparse) experiences significantly smaller degrade in performance.}
\label{Fig:mfrl_benchmark}
\end{figure}

We benchmark hierarchical control for primitives learned without supervision, against our proposed scheme using an MPPI based planner on top of DADS-learned skills. We persist with the task of Ant-navigation as described in \ref{sec:mbrl}. We benchmark against Hierarchical DIAYN \citep{eysenbach2018diversity}, which learns the skills using the DIAYN objective, freezes the low-level policy and learns a meta-controller that outputs the skill to be executed for the next $H_Z$ steps. We provide the $x$-$y$ prior to the DIAYN's disciminator while learning the skills for the Ant agent. The performance of the meta-controller is constrained by the low-level policy, however, this hierarchical scheme is agnostic to the algorithm used to learn the low-level policy. To contrast the quality of primitives learned by the DADS and DIAYN, we also benchmark against Hierarchical DADS, which learns a meta-controller the same way as Hierarchical DIAYN, but learns the skills using DADS.

From Figure \ref{Fig:mfrl_benchmark} (Left) We find that the meta-controller is unable to compose the skills learned by DIAYN, while the same meta-controller can learn to compose skills by DADS to navigate the Ant to different goals.
This result seems to confirm our intuition described in Section~\ref{sec:variance_primitives}, that the high variance of the DIAYN skills limits their temporal compositionality. Interestingly, learning a RL meta-controller reaches similar performance to the MPPI controller, taking an additional $200,000$ samples per goal.

\subsection{Goal-conditioned RL}
\label{sec:gcrl}
\vspace{-3mm}
To demonstrate the benefits of our approach over model-free RL, we benchmark against goal-conditioned RL on two versions of Ant-navigation: (a) with a dense reward $r(u)$ and (b) with a sparse reward $r(u) = 1$ if $\lVert u-g \rVert_2 \leq \epsilon$, else 0. We train the goal-conditioned RL agent using soft actor-critic, where the state variable of the agent is augmented with $u-g$, the position delta to the goal. The agent gets a randomly sampled goal from $[-10, 10]^2$ at the beginning of the episode.

In Figure \ref{Fig:mfrl_benchmark} (Right), we measure the average performance of the all the methods as a function of the initial distance of the goal, ranging from 5 to 30 metres. For dense reward navigation, we observe that while model-based planning on DADS-learned skills degrades smoothly as the initial distance to goal to increases, goal-conditioned RL experiences a sudden deterioration outside the goal distribution it was trained on. Even within the goal distribution observed during training of goal-conditioned RL model, skill-space planning performs competitively to it. With sparse reward navigation, goal-conditioned RL is unable to navigate, while MPPI demonstrates comparable performance to the dense reward up to about 20 metres. This highlights the utility of learning task-agnostic skills, which makes them more general while showing that latent space planning can be leveraged for tasks requiring long-horizon planning.
\vspace{-1.5mm}
\section{Conclusion}
\vspace{-3.5mm}
We have proposed a novel unsupervised skill learning algorithm that is amenable to model-based planning for hierarchical control on downstream tasks. We show that our skill learning method can scale to high-dimensional state-spaces, while discovering a diverse set of low-variance skills. In addition, we demonstrated that, without any training on the specified task, we can compose the learned skills to outperform competitive model-based baselines that were trained with the knowledge of the test tasks. We plan to extend the algorithm to work with off-policy data, potentially using relabelling tricks \citep{DBLP:journals/corr/AndrychowiczWRS17, nachum2018data} and explore more nuanced planning algorithms. We plan to apply the hereby-introduced method to different domains, such as manipulation and enable skill/model discovery directly from images.

\vspace{-1.5mm}
\section{Acknowledgements}
\vspace{-3mm}
We would like to thank Evan Liu, Ben Eysenbach, Anusha Nagabandi for their help in reproducing the baselines for this work. We are thankful to Ben Eysenbach for their comments and discussion on the initial drafts. We would also like to acknowledge Ofir Nachum, Alex Alemi, Daniel Freeman, Yiding Jiang, Allan Zhou and other colleagues at Google Brain for their helpful feedback and discussions at various stages of this work. We are also thankful to Michael Ahn and others in Adept team for their support, especially with the infrastructure setup and scaling up the experiments.


\section{Introduction}
The basic procedure of demapping received symbols back into their embedded soft bits, feeding a following stage of error correction decoding, is a crucial component in any modern communication system. The soft bit, quantifying the level of confidence in the quality of the symbol's demodulation process, is typically expressed in terms of the \emph{log-likelihood ratio} (LLR).

An exact evaluation of the LLR is achieved via a scheme, known as the \emph{log-MAP}~\cite{bib:Erfanian1994}, computing the logarithm of the ratio between the maximum a-posteriori (MAP) probabilities of the latent bit's two hypotheses given the observed symbol. Though statistically optimal, the computational complexity of the log-MAP algorithm scales with the size of the symbol constellation, making its direct implementation impractical in realistic systems.

A popular approximation of the optimal log-MAP rule, serving as a feasible golden target in designing practical systems, is the well-known \emph{max-log-MAP} algorithm~\cite{bib:Robertson1995}. However, although eliminating the need to compute complex exponential and logarithmic functions, in general the approximated max-log-MAP algorithm still consumes, on a symbol-rate basis, an extensive number of operations which scales with the modulation order.

There are certain constellation schemes, like the Gray-coded quadrature amplitude modulation (QAM), for which the max-log-MAP boils down to a reduced-complexity piecewise linear function and can be thus implemented, for instance, via a lookup table (LUT)~\cite{bib:Tosato2002}. This simplification is attributed to the QAM's inherent constellation mapping symmetries and its separability into two independent pulse amplitude modulations (PAM) in the real and imaginary parts of the symbol. As shall be exemplified in the sequel, although the computational ease in the operation of the max-log-MAP algorithm on QAM constellations, it still suffers from significant performance degradation in the estimation of the LLRs compared to the exact log-MAP, especially in the low and intermediate signal-to-noise ratio (SNR) regimes. This degradation is especially detrimental in the (almost) inevitable pragmatic case of non-ideal link adaptation. Further computationally simplified versions of the max-log-MAP for either QAM~\cite{bib:Tosato2002} or other modulation schemes, like phase-shift keying (PSK)~(\eg, \cite{bib:Wang2014} and references therein), suffer from an additional performance penalty w.r.t. the original "full-blown" max-log-MAP algorithm.

In this contribution, a machine-learning architecture for efficient universal soft demodulation, dubbed \emph{"LLRnet"}, is proposed. One should bear in mind that the demapping procedure can be simply abstracted as nothing but a symbol-to-bit LLRs function. Now, neural networks are well-known as a great tool in effectively approximating functions (viz. Cybenko's universal approximation theorem~\cite{bib:Cybenko1989}). Hence it seems very natural and beneficial to train a neural network to directly learn the functionality of either the impractical exact log-MAP, the cumbersome max-log-MAP, or any other expert-based target demapping rule. As shall be shown, the LLRnet provides an excellent mechanism for achieving the target demodulation rule performance. For the family of QAM constellations, LLRnet is shown to practically reproduce the exact log-MAP algorithm with a substantially reduced computational burden (\eg, ten times less the operations for 1024-QAM).

The paper is organized as follows. The LLR estimation problem is described in Section~\ref{sec_problem}. Section~\ref{sec_LLRnet}
introduces the proposed LLRnet architecture for soft demodulation. Section~\ref{sec_Examples} provides simulation results and discusses the performance of LLRnet in the context of two commercial system applications, namely the cellular 5G and the satellite DVB-S.2 standards. Finally, Section~\ref{sec_Conclusion} contains some concluding remarks.

\section{Problem Formulation}\label{sec_problem}
Consider a generic communication system (as depicted in Fig.~\ref{fig_system_model}) with a modulator at the transmitter mapping an incoming stream of encoded bits (where each bit is denoted by $c\in\{0,1\}$) to an outgoing stream of modulated symbols, $\vs\triangleq\{s_{1},\ldots,s_{N}\}^{T}\in\mathcal{C}$, chosen from an arbitrary finite set of constellation points in a (possibly hyper-dimensional) complex domain, $\mathcal{C}\subset\mathbb{C}^{N}$, $N\in\mathbb Z_{> 0}$. To this end, let $\vc\triangleq\{c_{1},\ldots,c_{M}\}^{T}\in\{0,1\}^{M}$, be a vector of some arbitrary $M$ consecutive bits in the stream being modulated to an $N$-dimensional complex symbol $\vs$.\comment{\footnote{Note that in most commercial systems $N=1$ is used.}}\comment{It is assumed that the modulator's output power is normalized, namely $\norm{\vs}_{2}=1$.} Hereinafter, it is assumed that all possible $M$-bit vectors, $\vc$, are equiprobable, thus all the modulated symbols in the constellation set $\mathcal{C}$ are equally likely to be transmitted. Incorporating knowledge, \eg, iteratively originated from the error correction decoding stage, on (unequal) prior probabilities of the modulated symbols in the discussed demodulation schemes is straightforward.

\begin{figure}[t!]
    \centering
    \includegraphics[width=\columnwidth,bb=0 0 380 96]{model}
    \caption{A communication system comprising transmit modulator, composite channel and receive demodulator. The composite channel encapsulates arbitrary processing stages of succeeding transmitter, physical channel, and receiver.}
    \label{fig_system_model}
\end{figure}

The symbol $\vs$ is transmitted through a composite channel comprising, in addition to the physical channel, an abstraction of other critical transceiver processing steps. On the transmitter side, these may consist of (but not limited to) precoding for transmission over multiple-input multiple-output (MIMO) systems, mapping to orthogonal frequency-division multiplexing (OFDM), digital-to-analog conversion (DAC) and analog mixing on a radio carrier. On the receiver side, the composite channel may encapsulate, \eg, analog-to-digital conversion (ADC), down-mixing to baseband, synchronization, filtering, OFDM demodulation, MIMO combining and channel estimation and equalization. These transceiver processing steps are standard and do not explicitly pertain to the problem under study of demapping, or soft demodulation, of the received symbol to bit LLRs.

The demodulator at the receiver observes the composite channel's complex-valued symbol estimate, $\hat{\vs}\in\mathbb{C}^{N}$, where the effect of the composite channel is captured by the relation between the transmitted and observed symbols, $\vs$ and $\hat{\vs}$, respectively. Although not necessarily accurate, for the purpose of soft demodulation to bit LLRs, the composite channel is typically modeled as an additive white Gaussian noise (AWGN) channel, to yield
\BE\label{eq_model}
    \hat{\vs}=\vs+\vn,
\EE
where the $N$-dimensional AWGN vector $\vn\sim\mathcal{CN}(\mathbf{0},\sigma^{2}\mI_{N})$ and $\sigma^{2}\geq\sigma_{0}^{2}$. The scalar  $\sigma_{0}$ denotes the standard deviation of the ambient noise in the physical channel, corresponding to some SNR. It is important to note that this work does not focus on optimizing the, so called, "composite channel" itself (\eg, via improving MIMO demodulation), but in streamlining symbol demapping by adopting a machine-learning approach. To this end, the observed symbol estimate, $\hat{\vs}$, can be viewed, for instance, as an equalized symbol following a linear detection stage.

Now mainly for the purpose of facilitating an improved decoding in a succeeding stage, the demodulator demaps the observed symbol, $\hat{\vs}$, into a vector of estimates of the $M$-bit LLRs, $\vl\in\mathbb{R}^{M}$, corresponding to the confidence in the inference of each of the original coded bits in $\vc$. The $i$'th entry of the LLRs vector, $l_{i}$, adheres to the logarithm of the MAP relation
\BE\label{eq_relation}
    l_{i}\triangleq\log\Bigg(\frac{\Pr{(c_{i}=0|\hat{\vs})}}{\Pr{(c_{i}=1|\hat{\vs})}}\Bigg),\quad i=1,\ldots,M.
\EE
An \emph{exact} computation of the log-MAP expression~\eqref{eq_relation}, under the examined model~\eqref{eq_model}, yields
\BE\label{eq_logMAP}
    l_{i}=\log\frac{{\sum_{\vs\in\mathcal{C}_{i}^{0}}\exp{\Big(-\frac{\norm{\hat{\vs}-\vs}_{2}^{2}}{\sigma^{2}}\Big)}}}{\sum_{\vs\in\mathcal{C}_{i}^{1}}\exp{\Big(-\frac{\norm{\hat{\vs}-\vs}_{2}^{2}}{\sigma^{2}}\Big)}},\quad i=1,\ldots,M,
\EE
where $\mathcal{C}_{i}^{\delta}\in\mathcal{C}$ is the subset of constellation points, known at the receiver side, for which the $i$'th bit is equal to $\delta\in\{0,1\}$.
Applying the approximation
\BE
    \log\Bigg(\sum_{j}\exp\big(-x_{j}^{2}\big)\Bigg)\approx\max_{j}\big(-x_{j}^{2}\big) \nonumber
\EE
on the exact log-MAP operation~\eqref{eq_logMAP}, \emph{approximated} LLR estimates can be derived from a simplified rule, well-known as the max-log-MAP
\BE\label{eq_maxlogMAP}
    l_{i}\approx\frac{1}{\sigma^{2}}\Bigg(\min_{\vs\in\mathcal{C}_{i}^{1}}\norm{\hat{\vs}-\vs}_{2}^{2}-\min_{\vs\in\mathcal{C}_{i}^{0}}\norm{\hat{\vs}-\vs}_{2}^{2}\Bigg),\quad i=1,\ldots,M.
\EE

The computational complexity per received symbol of a brute-force implementation of either the log-MAP~\eqref{eq_logMAP} or the max-log-MAP~\eqref{eq_maxlogMAP} algorithms scales with the size of the constellation, $|\mathcal{C}|$, to yield $\mathcal{O}(2^{MN})$ operations. However, the approximated max-log-MAP demapping rule has the advantageous property of eliminating the need for computing complex exponential and logarithmic functions. For dealing with the popular (Gray-coded) QAM constellations, although being severely sub-optimal the max-log-MAP is typically the common demapper of choice due to the fact that it can be reduced to a single LUT implementation of piecewise linear functions, which scales linearly with SNR. Note that the exact log-MAP algorithm can be only crudely approximated via multiple LUTs, essentially one for each SNR working point~\cite{bib:Yao2015}.

\section{Soft Demodulation with LLRnet}\label{sec_LLRnet}
\begin{figure}[t!]
    \centering
    \includegraphics[width=\columnwidth,bb=0 0 509 360]{LLRnet}
    \caption{LLRnet architecture composed of: An input of $N$-dimensional received symbol split into its real and imaginary parts; $K$-neurons hidden layer with (\eg) ReLU activation function; An $M$-bit LLR linear output layer.}
    \label{fig_model}
\end{figure}

LLRnet is a neural network designed to learn a desired soft demodulation scheme, demapping a complex symbol to its conveyed bits' real-valued LLRs. The architecture of LLRnet is illustrated in Fig.~\ref{fig_model}. The input layer of the LLRnet is fed by the received symbol estimate vector, $\hat{\vs}$, where it is divided into its real and imaginary parts. The output vectors of the two input nodes are injected to a hidden layer of $K$ neurons. The scalar output of the $k$'th neuron is determined by $y_{k}=f(\vw_{k}^{T}\vx_{k}+b_{k})$, where $\vx_{k},\vw_{k}\in\mathcal{R}^{2N}$ and $b_{k}\in\mathcal{R}$ are the $k$'th neuron's input vector, weights vector and bias, respectively. The operator $f(\cdot)$ denotes the neuron's transfer, or activation, function. Unless otherwise stated, in its hidden layer LLRnet uses a rectified linear unit (ReLU), namely $f(x)=x^{+}=\max(0,x)$. The outputs of the $K$ hidden neurons are then inserted to a linear output layer of $M$ nodes. The $m$'th output of the LLRnet, estimating the $m$'th bit LLR, is determined by $\hat{l}_{m}=\boldsymbol{\omega}_{m}^{T}\vy+\beta_{m}$, where $\vy\triangleq\{y_{1},\ldots,y_{K}\}^{T}\in\mathcal{R}^{K}$ is a concatenation of the hidden layer outputs, while $\boldsymbol{\omega}_{m}\in\mathcal{R}^{K}$ and $\beta_{m}\in\mathcal{R}$ are the $m$'th output node's weights vector and bias, respectively. The set of all trainable parameters of the LLRnet, namely all the weights $\vw_{k}$, $\boldsymbol{\omega}_{m}$ and biases $b_{k}$, $\beta_{m}$, are denoted by $\boldsymbol{\theta}$.

Fig.~\ref{fig_diagram} schematically describes the training process of the neural demodulator. Since the LLRnet's input-to-output function is fully differentiable w.r.t. the set of all trainable parameters, $\boldsymbol{\theta}$, a gradient-based training approach can be adopted. The trainable parameters are first randomly initialized. For a certain batch of $B$ received symbols, $\hat{\vs}^{(b)}$, $b=1,\ldots,B$, compute the corresponding bit LLRs via a desired demapping algorithm, \eg~log-MAP or max-log-MAP (for implementing these specific target demodulation algorithms, an estimate of $\sigma^{2}$ is also required).\footnote{Any other expert-based demapping rule can be targeted, \eg, in the case of colored noise due to equalization.} Then feed both the conventionally computed LLRs, $\vl^{(b)}$, along with the LLRs computed through LLRnet, $\hat{\vl}^{(b)}$, into a loss function which is evidently a function of the trainable set of parameters $\boldsymbol{\theta}$. Such a loss function could be, for instance, the mean-squared error (MSE)
\BE
    L^{\text{MSE}}(\boldsymbol{\theta})\triangleq\frac{1}{B}\sum_{b=1}^{B}\norm{\hat{\vl}^{(b)}-\vl^{(b)}}_{2}^{2}, \nonumber
\EE
or alternatively the cross-entropy function
\BE
    L^{\text{CE}}(\boldsymbol{\theta})\triangleq-\frac{1}{B}\sum_{b=1}^{B}\sum_{m=1}^{M}l_{m}^{(b)}\log(\hat{l}_{m}^{(b)})+(1-l_{m}^{(b)})\log(1-\hat{l}_{m}^{(b)}). \nonumber
\EE
Define a stop criterion which can be either a fixed number of iterations, a threshold on the loss or a number of iterations during which the loss
has not decreased. Unless the stop criterion is met, update the parameter set, $\boldsymbol{\theta}$, based on a learning algorithm using gradient descent,
\BE
    \boldsymbol{\theta}\leftarrow\boldsymbol{\theta}-\alpha\nabla_{\boldsymbol{\theta}}L(\boldsymbol{\theta}), \nonumber
\EE
where $\alpha>0$ is the learning rate. Compute the loss function again under the newly learned set of trained parameters. When the stop criterion is met and training is complete, one moves to the inference stage where the LLRnet serves as the sole demodulator, efficiently imitating the functionality of the desired soft demodulator. In the following section LLRnet is simulated and its performance is evaluated in two end-to-end system usecases.

\begin{figure}[t!]
    \centering
    \includegraphics[width=\columnwidth,bb=0 0 217 219]{diagram}
    \caption{Block diagram for the training of LLRnet.}
    \label{fig_diagram}
\end{figure}

\section{Simulation Examples}\label{sec_Examples}
\subsection{Throughput of PDSCH in 5G-NR}\label{sec_5G}

\begin{figure}[hp]
    \centering
    \begin{subfigure}[b]{\columnwidth}
        \centering
        \includegraphics[width=\columnwidth,bb=0 0 225 172]{LLR16}
        \caption{16-QAM}
    \end{subfigure}
    \vskip\baselineskip
    \begin{subfigure}[b]{\columnwidth}
        \centering
        \includegraphics[width=\columnwidth,bb=0 0 225 172]{LLR64}
        \caption{64-QAM (Legend is identical to that of (a).)}
    \end{subfigure}
    \vskip\baselineskip
    \begin{subfigure}[b]{\columnwidth}
        \centering
        \includegraphics[width=\columnwidth,bb=0 0 225 172]{LLR256}
        \caption{256-QAM (Legend is identical to that of (a).)}
    \end{subfigure}
    \caption{LLR, derived via log-MAP, max-log-MAP and LLRnet, as a function of the real part of the symbol for the odd bits in 16/64/256-QAM.}
    \label{fig_LLR_QAM}
\end{figure}

\begin{figure}[hp]
    \centering
    \begin{subfigure}[b]{\columnwidth}
        \centering
        \includegraphics[width=\columnwidth,bb=0 0 149 115]{tput16}
        \caption{16-QAM}
    \end{subfigure}
    \vskip\baselineskip
    \begin{subfigure}[b]{\columnwidth}
        \centering
        \includegraphics[width=\columnwidth,bb=0 0 149 115]{tput64}
        \caption{64-QAM}
    \end{subfigure}
    \vskip\baselineskip
    \begin{subfigure}[b]{\columnwidth}
        \centering
        \includegraphics[width=\columnwidth,bb=0 0 149 115]{tput256}
        \caption{256-QAM}
    \end{subfigure}
    \caption{PDSCH throughput, derived via log-MAP, max-log-MAP and LLRnet, as a function of SNR for 16/64/256-QAM.}
    \label{fig_tput}
\end{figure}

\setlength{\aboverulesep}{0pt}
\setlength{\belowrulesep}{0pt}
\setlength{\extrarowheight}{.75ex}

\begin{table*}[tp]
\captionsetup{size=normalsize,

    skip=5pt, position = bottom}
\renewcommand\thetable{II}
\caption{Computational complexity comparison of the demodulation schemes in terms of real operations per received symbol.}
\label{table_complexity}
    \centering
    \begin{threeparttable}
    \begin{tabular}{laaccccaa}
    \toprule
     & \multicolumn{2}{c}{16-QAM}& &\multicolumn{2}{c}{64-QAM}& &\multicolumn{2}{c}{256-QAM}\\
              \cmidrule{2-3} \cmidrule{5-6} \cmidrule{8-9}
            & {log-MAP} & {LLRnet}
            & & {log-MAP} & {LLRnet}
            & & {log-MAP} & {LLRnet\tnote{*}}\\
            & &  {(K=8)}&
            & &  {(K=16)}&
            & &  {(K=32)} \\
    \midrule
    Multiplication \& division             & 52  & 40&& 198  & 112 && 776 & 288 (320)\\
    Addition \& subtraction            & 104  & 40 && 564  & 112 && 2800 & 288 (352)\\
    Exponent \& logarithm      & 20 & 0 && 70  & 0  && 264 & 0 (64)\\
    Comparator      & 0 & 8 && 0  & 16  && 0 & 32 (0)\\
    \midrule\midrule
    Total & 176 & 88 && 832  & 240 && 3840 & 608 (736)\\
    \bottomrule
    \end{tabular}
    \begin{tablenotes}
        \item[*] Parentheses correspond to using $\tanh$, rather than ReLU, activation function at the hidden layer.
    \end{tablenotes}
    \end{threeparttable}
\end{table*}

\begin{table}[bp]
\renewcommand\thetable{III}
\caption{Computational complexity for 1024-QAM.}
\label{table_complexity_1024}
    \centering
    \begin{tabular}{lcc}
    \toprule
     & \multicolumn{2}{c}{1024-QAM} \\
              \cmidrule{2-3}
            & {log-MAP} & {LLRnet} \\
            & &  {(K=64)}\\
    \midrule
    Multiplication \& division             & 3082  & 704 \\
    Addition \& subtraction            & 13292  & 704 \\
    Exponent \& logarithm      & 1034 & 0 \\
    Comparator      & 0 & 64 \\
    \midrule\midrule
    Total & 17408 & 1472 \\
    \bottomrule
    \end{tabular}
\end{table}

\begin{table}[tp]
\renewcommand\thetable{I}
\caption{Simulated 5G-NR key configuration parameters.}
\label{tab_param}
    \centering
\begin{tabular}{|l|c|}
  \hline

  \# of 10ms frames & 200 \\
  Bandwidth & 20MHz \\
  \# of (12 sub-carriers) resource blocks & 51 \\
  Subcarrier spacing & 30KHz \\
  Cyclic prefix & Normal \\
  \# of Tx. antennas & 8 \\
  \# of Rx. antennas & 2 \\
  \# of layers & 2 \\
  Transport channel coding & LDPC \\
  Target code rate & 0.4785 \\
  Modulation & 16-QAM, 64-QAM, 256-QAM \\
  PDSCH precoding & SVD, single matrix \\
  Waveform & CP-OFDM \\
  Channel model & Clustered delay line (CDL) \\
  Channel estimation & Ideal \\
  MIMO equalization & Linear MMSE \\
  Synchronization & Perfect \\
  HARQ & Enabled, 16 processes \\
  \hline
\end{tabular}
\end{table}

The proposed LLRnet demodulation engine is utilized in the decoding of the physical downlink shared channel (PDSCH) in a 5G New Radio (NR) link, as defined by the 3GPP NR standard. The key configuration parameters of the simulated link are listed in Table~\ref{tab_param}.
In the simulated configuration, a frame (10ms) is divided into 10 subframes, where each subframe (1ms) is composed of 2 slots, and each slot (0.5ms) consists of 14 symbols. For each slot, out of the $51\times12\times14\times2=17136$ resource elements in the received resource grid, $14712$ convey PDSCH symbols.

The training of the LLRnet relies on only about $1\%$ ($148$ symbols) of the $14712$ received PDSCH symbols, taken evenly across a single slot which is, in the reported results, the first slot in a 2 seconds transmission per evaluated SNR point. These $148$ symbols are randomly divided to three sets: (1) about $70\%$ of them are used for training, (2) about $15\%$ used for validation that the trained network has low enough generalization error and for avoiding excessive training which may result in undesirable overfitting, and (3) lastly, about $15\%$ of the symbols are used for a completely independent testing of the LLRnet generalization. Note that the relatively small size of the data set, found sufficient for training the LLRnet, stems from the inherent smoothness of the demapping functions from symbols to bit LLRs.

The chosen training algorithm in this example is the Levenberg-Marquardt backpropagation algorithm~\cite{bib:Hagan1994} and the training process continues until the validation error fails to decrease for $6$ consecutive iterations. The chosen loss function is the MSE, where the desired, or target, demodulation algorithm the LLRnet is trained to reproduce is the exact, yet costly, log-MAP scheme (with $\sigma^{2}=\sigma_{0}^{2}$). The number of neurons in the hidden layer is set to $K=8,16,32$ for the constellations $16,64,256$-QAM (thus $N=1$), respectively. As mentioned previously, the neuron's activation function is a ReLU. Upon completion of the training stage, PDSCH symbols in the consecutive slots are demodulated exclusively by the LLRnet. Re-training, so the LLRnet can readapt and learn the new demapping functions, is required only when a (non-marginal) SNR working point change is identified. Alternatively, training for different SNR values can be performed offline or in a quasi-offline manner, with the obtained trained sets being stored in the receiver's memory. Another option is to provide the SNR level as an additional input to the LLRnet and train it within a wide range of SNRs. By doing so, the LLRnet can then learn to also generalize across SNR points.

Figs.~\ref{fig_LLR_QAM}(a)-(c) present, for 16-QAM, 64-QAM and 256-QAM (in which each symbol conveys $4$, $6$ and $8$ bits, respectively), in three different SNR levels, the explicit demapping functions of the real part of the (received) symbol into its carried odd-numbered bits' LLR for three different soft demodulation implementations: exact log-MAP~\eqref{eq_logMAP}, approximate max-log-MAP~\eqref{eq_maxlogMAP} and the proposed LLRnet. As expected, in the higher SNR regimes (lower row in each figure) for the examined three QAM modulations the max-log-MAP exhibits a good estimate of the optimal log-MAP rule. However, for the low and intermediate SNR regimes (two upper rows in each figure) the max-log-MAP only serves as a crude (low SNR) to reasonable (intermediate SNR) approximation to the exact demapper. On the other hand the soft bits inferred by the LLRnet practically coincide with the optimal ones in all three modulation cases across the entire relevant SNR range.

Next, the performance of LLRnet is evaluated in terms of measuring the PDSCH throughput. Figs.~\ref{fig_tput}(a)-(c) plot the PDSCH throughput as a function of SNR. The throughput is displayed both in terms of absolute values, in Mbps, and relative throughput, in percentage, w.r.t. the link's maximum possible throughput under the given configuration. It is evident that LLRnet (again, with only $K=8,16,32$ neurons populating the hidden layer, respectively) essentially provides the same throughput as the optimal log-MAP algorithm.\footnote{For high SNR points under $256$-QAM modulation, it was observed that using a hyperbolic tangent sigmoid transfer function yields slightly better performance than ReLU, at the expense of marginally more operations as depicted in Table~\ref{table_complexity}.} It is also observed that in the low and intermediate SNR levels the throughput corresponding to the tractable max-log-MAP substantially lags behind the throughput associated with the LLRnet (and log-MAP).

\begin{figure}[t!]
    \centering
    \includegraphics[width=\columnwidth, bb=0 0 420 315]{2D}
    \caption{Comparison of LLRnet vs. max-log-MAP in a performance-complexity chart w.r.t. the exact log-MAP.}
    \label{fig_2D}
\end{figure}

The capability of LLRnet to successfully imitate the operation of the log-MAP algorithm is even more remarkable when comparing their computational complexity. Table~\ref{table_complexity} lists the number of (real) operations, per symbol, required for the execution of log-MAP and LLRnet for the three different QAM constellations. Note that without any loss in throughput, LLRnet costs about $50\%$, $70\%$, and $85\%$ \emph{less} total operations than the log-MAP algorithm. Bear in mind again these algorithms are run on symbol rate. For completeness, and although currently not yet being an integral part of 5G-NR, but of Wi-Fi 6 (802.11ax), the enumeration of the complexity savings (more than 90\%) of LLRnet vs. log-MAP for the 1024-QAM are listed in Table~\ref{table_complexity_1024}. The overhead complexity of the training phase itself in this application is relatively small since the learning relies on only few symbols within a single received slot, thus it is omitted from the comparison in Tables~\ref{table_complexity} and~\ref{table_complexity_1024}. In addition, the number of epochs required for training the LLRnet was typically below $30$. In this particular case of QAM constellations which can be divided into two separate PAM constellations, it is interesting to observe how LLRnet successfully learns, as anticipated, to null half of the hidden layer's weights. Alternatively in such a case, one can simply use two smaller LLRnet architectures processing separately the real and imaginary parts of the symbol.

Fig~\ref{fig_2D} visually compares the three different demapping algorithms through the construction of a two-dimensional chart in which the horizontal axis measures the relative complexity (w.r.t. the brute-force exact log-MAP, in percentage), and the vertical axis measures the relative loss in throughput (again, w.r.t. log-MAP, in percentage). For three different working points, with three different QAM modulation schemes, it is illustrated how the implementation of LLRnet, rather than the popular max-log-MAP, yields an improved position in the performance-complexity domain: Diminishing to practically zero with LLRnet, original throughput losses between $15\%$ to $25\%$, originated from the max-log-MAP, yet still maintaining feasible computational complexity burden via LLRnet. Observe how for each tested case, the position in the two-dimensional chart, based on max-log-MAP, moves into a new cost-effective position due to LLRnet (namely, {\color{red}$\blacktriangle$}, connected through a dashed line, shift into {\color{blue}$\bullet$}, located in a better position). Also marked are the performance-complexity metric position for the case of LLRnet in 1024-QAM and the reference metric position for the exact log-MAP (\ie, $\blacksquare$ at $100\%$ complexity and zero throughput loss point).

\subsection{Packet Error Rate in DVB-S.2}
\begin{figure}[t!]
    \centering
    \includegraphics[width=\columnwidth, bb=0 0 149 115]{8PSK}
    \caption{Packet error rate vs. SNR for DVB-S.2 link with modulation 8-PSK and LDPC code rate 3/5.}
    \label{fig_per_8psk}
\end{figure}

In this example, LLRnet is incorporated in the second generation digital video broadcasting standard (DVB-S.2) for broadband satellite communications. In this application example, LLRnet, with $K=8$ neurons in the hidden layer, is trained to mimic the operation of the approximate max-log-MAP algorithm in the demapping of 8-PSK modulation (thus $N=1$). A $3/5$ LDPC (decoded with a maximum of 50 iterations) and BCH codes serve as the inner and outer codes, respectively, in an AWGN channel. $1000$ DVB-S.2 frames were simulated for each SNR point. The training stage for this application example follows the procedure described for the training in Section~\ref{sec_5G}.

Fig.~\ref{fig_per_8psk} plots the simulated packet error rate as a function of the SNR using two implementations of soft demodulation: max-log-MAP and LLRnet, this time learned to imitate the approximate max-log-MAP (with $\sigma^{2}=\sigma_{0}^{2}$), rather than the exact log-MAP as in Section~\ref{sec_5G}. It can be observed that a perfect alignment in the performance curves of the two algorithms is achieved. However, it should be noted, that in this particular case, using LLRnet with $K=8$ neurons, the two implementations require roughly the same total number of operations (about 64 per symbol). It may happen that deep, rather than shallow, learning architectures could demonstrate, in addition to accurate inference of the LLRs, also further computational savings for 8-PSK and higher PSK, or amplitude and PSK (APSK) constellations.

\section{Conclusion}\label{sec_Conclusion}
In this paper, a "Machine-LLRning" approach is revealed for which a simple neural network architecture is trained to efficiently soft demodulate symbols to their bit LLRs, thus utilizing an artificial intelligence paradigm straight into one of the most generic building blocks of the physical layer processing. The proposed concept of LLRnet can be extended not only for multi-layer deep learning architectures, but also be integrated in a more holistic trainable receiver structure, jointly carrying out the tasks of neural demodulator (as proposed in this contribution), along with trainable quantization of the LLRs, and a trainable channel decoding. An evaluation of the LLRnet under realistic dynamical link adaptation, rather than a fixed modulation and coding scheme (MCS), is also called for. Furthermore, by feeding the channel estimation itself into the LLRnet, the tasks of MIMO and multiuser detection could also be potentially tackled in a straightforward manner in such a framework.

\bibliographystyle{IEEEtran}