DR-PETS: Learning-Based Control With Planning in Adversarial Environments

Extracted source text · arXiv 2503.20660

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%2345678901234567890123456789012345678901234567890123456789012345678901234567890
%        1         2         3         4         5         6         7         8

\documentclass[letterpaper, 10 pt, conference]{ieeeconf}  % Comment this line out if you need a4paper

%\documentclass[a4paper, 10pt, conference]{ieeeconf}      % Use this line for a4 paper

\IEEEoverridecommandlockouts                              % This command is only needed if 
                                                          % you want to use the \thanks command

\overrideIEEEmargins                                      % Needed to meet printer requirements.

%In case you encounter the following error:
%Error 1010 The PDF file may be corrupt (unable to open PDF file) OR
%Error 1000 An error occurred while parsing a contents stream. Unable to analyze the PDF file.
%This is a known problem with pdfLaTeX conversion filter. The file cannot be opened with acrobat reader
%Please use one of the alternatives below to circumvent this error by uncommenting one or the other
%\pdfobjcompresslevel=0
%\pdfminorversion=4

% See the \addtolength command later in the file to balance the column lengths
% on the last page of the document

% The following packages can be found on http:\\www.ctan.org
\usepackage{graphics} % for pdf, bitmapped graphics files
\usepackage{epsfig} % for postscript graphics files
% \usepackage{mathptmx} % assumes new font selection scheme installed
\usepackage{mathrsfs}
\usepackage{times} % assumes new font selection scheme installed
\usepackage{amsmath} % assumes amsmath package installed
\usepackage{amssymb}  % assumes amsmath package installed
\usepackage{amsfonts}
\usepackage{amssymb}
\usepackage{pgf}
\usepackage{graphicx}
\usepackage{mathrsfs}
\usepackage{enumerate}
\usepackage{soul}
\usepackage{tikz}
\usepackage{pgfplots}
\usepackage{standalone}
\usepackage{url}
% \usepackage{subfigure}
% \usepackage{caption}
\usepackage{algorithm,algorithmicx}
% \usepackage{algorithm2e}
% \SetAlFnt{\small}
% \SetAlCapFnt{\large}
% \SetAlCapNameFnt{\large}
% \usepackage{algorithmic}
% \algsetup{linenosize=\tiny}

\usepackage[prependcaption,colorinlistoftodos]{todonotes}


\usepackage{algpseudocode}
\newtheorem{assumption}{Assumption}
\newtheorem{definition}{Definition}
\renewcommand{\algorithmicrequire}{\textbf{Input:}}
\renewcommand{\algorithmicensure}{\textbf{Output:}}
\newcommand{\norm}[1]{\left\lVert#1\right\rVert}

\newcommand{\sC}{\mathcal{C}}
\newcommand{\sD}{\mathcal{D}}
\newcommand{\sS}{\mathcal{S}}
\newcommand{\sA}{\mathcal{A}}
\newcommand{\sL}{\mathcal{L}}
\newcommand{\sM}{\mathcal{M}}
\newcommand{\sP}{\mathbb{P}}
\newcommand{\sN}{\mathcal{N}}
\newcommand{\sV}{\mathcal{V}}
\newcommand{\sE}{\mathcal{E}}
\newcommand{\sF}{\mathcal{F}}
\newcommand{\sQ}{\mathcal{Q}}
\newcommand{\sT}{\mathcal{T}}
%\newcommand{\Nh}{N_h}
\newcommand{\Nh}{n}

\newcommand{\R}{\mathbb{R}}
\newcommand{\E}{\mathbb{E}}
\newcommand{\Z}{\mathbb{Z}}

\newcommand{\bv}[1]{\mathbf{#1}}
\newcommand{\policy}[2]{\pi_{#1}^{#2}}

\newcommand{\AAC}[1]{\textcolor{blue}{#1}}
\newcommand{\HJ}[1]{\textcolor{red}{#1}}

%%

\definecolor{gnblue6}{RGB}{35,156,255} % 239cff

\newcommand{\todoinb}[1]{\todo[inline,color=gnblue6!10, linecolor=gnblue6!10]{\small#1}}
\newcommand{\grtodo}[1]{\todoinb{\textbf{GRTodo from GR:} #1}}

%%
\DeclareFontEncoding{LS1}{}{}
\DeclareFontSubstitution{LS1}{stix}{m}{n}
\DeclareSymbolFont{symbols2}{LS1}{stixfrak}{m}{n}
\DeclareSymbolFont{symbols3}{LS1}{stixbb}{m}{n}
\SetSymbolFont{symbols2}{bold}{LS1}{stixfrak}{b}{n}
\SetSymbolFont{symbols3}{bold}{LS1}{stixbb}{b}{n}
\DeclareSymbolFont{arrows1}{LS1}{stixsf}{m}{n}
\SetSymbolFont{arrows1}{bold}{LS1}{stixsf}{b}{n}
\DeclareMathSymbol{\varrightarrow}{\mathrel}{arrows1}{"99}
\DeclareMathSymbol{\varleftarrow}{\mathrel}{arrows1}{"7D}
\DeclareMathSymbol{\fourvdots}{\mathord}{symbols2}{"38}
\newcommand{\tran}{'}
\newcommand{\ttran}{''}
\newcommand{\tttran}{'''}
\newtheorem{thm}{Theorem}[section]
\newtheorem{lem}[thm]{Lemma}
\newtheorem{prop}[thm]{Proposition}
\newtheorem{cor}{Corollary}
\newtheorem{Problem}{Problem}
\newtheorem{Remark}{Remark}
\title{\LARGE \bf
DR-PETS: Learning-Based Control With Planning in Adversarial Environments
}


\author{Hozefa Jesawada$^{a,1}$, Antonio Acernese$^{b}$, Giovanni Russo$^{a}$, Carmen~Del~Vecchio$^{b}$% <-this % stops a space
\thanks{$^{a}$ Department of Information and Electrical Engineering and Applied Mathematics, University of Salerno, Italy.}
%\thanks{$^{b}$ Mitsubishi Electric Research Laboratories, Cambridge, MA, USA.}
\thanks{$^{b}$  Department of Engineering, University of Sannio, Italy. Corresponding author Carmen Del Vecchio {\tt\small c.delvecchio@unisannio.it}}%
%\thanks{$^{d}$ Department of Electrical and Computer Engineering Division of Systems Engineering, Boston University, USA.} 
\thanks{$^{1}$ Work done while at the University of Sannio} 
\thanks{The work of HJ and CDV is supported by  PRIN PNRR program of the Italian government CUP C53D23008320001}% <-this % stops a space
}


\begin{document}

\maketitle
\thispagestyle{empty}
\pagestyle{empty}


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\begin{abstract}

%Ensuring robustness against epistemic, possibly adversarial, perturbations is critical for reliable planning and decision-making in real-world systems. While the Probabilistic Ensembles with Trajectory Sampling (PETS) algorithm inherently handles certain types of uncertainty through its ensemble-based probabilistic model, its robustness is implicit and limited to "average-case" scenarios, lacking guarantees against structured adversarial perturbations or worst-case uncertainty distributions. To bridge this gap, we propose DR-PETS, a distributionally robust extension of PETS that explicitly equips the algorithm with certified robustness against adversarial uncertainty. By redefining the uncertainty characterization using an ambiguity set based on the $p$-Wasserstein distance, we mathematically formalize adversarial perturbations and enable worst-case-aware planning. While the original PETS framework implicitly accounts for stochasticity, our approach explicitly optimizes for robustness by solving a min-max problem over this Wasserstein ambiguity set. To make the problem tractable, we derive a convex approximation and integrate it into the PETS planning loop. Experiments in control tasks like pendulum stabilization and cart-pole balancing demonstrate that DR-PETS maintains performance under adversarial parameter perturbations, outperforming vanilla PETS in worst-case scenarios while retaining its computational efficiency. This work shows how explicit distributional robustness can augment inherently robust algorithms like PETS for safety-critical applications.


Ensuring robustness against epistemic, possibly adversarial, perturbations is essential for reliable real-world decision-making. While the Probabilistic Ensembles with Trajectory Sampling (PETS) algorithm inherently handles uncertainty via ensemble-based probabilistic models, it lacks guarantees against structured adversarial or worst-case uncertainty distributions. To address this, we propose DR-PETS, a distributionally robust extension of PETS that  certifies robustness against adversarial perturbations. We formalize uncertainty via a $p$-Wasserstein ambiguity set, enabling worst-case-aware planning through a min-max optimization framework. While PETS passively accounts for stochasticity, DR-PETS actively optimizes robustness via a tractable convex approximation integrated into PETS’ planning loop. Experiments on pendulum stabilization and cart-pole balancing show that DR-PETS certifies robustness against adversarial parameter perturbations, achieving consistent performance in worst-case scenarios where PETS deteriorates. %Experiments on pendulum stabilization and cart-pole balancing demonstrate DR-PETS’ superiority under adversarial parameter shifts, outperforming PETS in worst-case settings without sacrificing efficiency. 

%Ensuring robustness against epistemic, possibly adversarial,  perturbations is essential for effective planning and decision-making in real-world environments. In this work, we extend the popular Probabilistic Ensembles with Trajectory Sampling (PETS) algorithm to provide robustness guarantees against adversarial uncertainty distributions, and introduce our distributionally robust DR-PETS algorithm. We characterize the uncertainty using an ambiguity set defined by the $p$-Wasserstein distance. To address the resulting otherwise intractable problem, we propose a tractable approximation and leverage the PETS algorithm to solve the problem. We demonstrate the robust performance of our approach in environments such as the pendulum and cart-pole against parameter perturbations.
%\grtodo{Non ho toccato l'abstract in maniera sostanziale ma lo imposterei in maniera leggermente differente, motivando perche' stiamo estendendo PETS. Ha proprieta' di robustezza intrinseche e quindi lo stiamo equipaggiando in maniera esplicita questa proprieta'?}
\end{abstract}

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% PROPOSED INTRO
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

\section{INTRODUCTION}
Deep reinforcement learning (RL) provides a powerful framework for sequential decision-making. However, deploying RL in practical settings demands robustness to epistemic uncertainty (e.g., adversarial perturbations or model mis-match)~\cite{GARRABE202281}. These uncertainties are often heterogeneous, non-stationary, or even adversarial, making robustness guarantees critical for safe deployment. While recent work has focused on robust RL methods that account for worst-case transitions~\cite{Yu2015}, many state-of-the-art algorithms remain vulnerable to structured adversarial perturbations.

Model-based RL (MBRL) improves sample efficiency by learning an explicit environment model~\cite{pmlr-v168-lellis22a}, enabling planning strategies like model predictive Control (MPC)~\cite{Klink2021} to simulate and optimize long-horizon decisions. Among MBRL methods, the probabilistic ensembles with trajectory sampling (PETS) algorithm~\cite{chua2018} stands out: it combines probabilistic ensemble models with MPC to achieve high data efficiency and performance. However, PETS’ robustness stems implicitly from its ensemble-based uncertainty quantification, which assumes stochastic (not adversarial) deviations. This leaves it ill-equipped for scenarios where perturbations strategically exploit model inaccuracies~\cite{pmlr-v100-okada20a}.


Motivated by this, we propose DR-PETS, a distributionally robust version of PETS. In DR-PETS, planning is still achieved through MPC but, differently from PETS, the MPC cost is regularized and the regularizer arises from a suitable reformulation of the underlying  control problem with ambiguity sets defined by the $p$-Wasserstein distance to capture epistemic uncertainties. Since this leads to an otherwise intractable problem, DR-PETS employs a tractable approximation, resulting in a regularized version of the original  MPC cost. In what follows we briefly survey some closely related works on DR learning and MPC.

\subsubsection*{Related work}
A common approach to handle uncertainty in RL is to consider worst-case transition models~\cite{Iyengar2005}, often formalized via distributionally robust Markov decision processes (DR-MDPs)~\cite{Derman2023,Yu2015}, where ambiguity sets are modeled using the Wasserstein distance~\cite{Gao2023}. Building on this,~\cite{derman2020distributionalregularization} establishes a duality between robustness and regularization, a concept also leveraged in model-free DR RL, such as DR Q-learning~\cite{liu2022DRQ}, which recasts Q-values into a DR Bellman operator via strong duality. Convergence guarantees for Wasserstein-robust Q-learning in linear, time-invariant systems are provided in~\cite{zhao2023minimax}. In~\cite{queeney2023optimal}, a robust RL framework for constrained MDP  using optimal transport perturbations is proposed; however, unlike our method, it requires offline perturbation computation and application to training data. In DR-MPC,~\cite{mark2020stochastic} formulates DR chance constraints and proposes a data-driven approach for value-at-risk under Wasserstein ambiguity, while~\cite{coppens2021datadriven} derives DR-MPC for conic moment-based ambiguity sets, assuming set shrinkage with more data. Finally,~\cite{wu2022ambiguity} proposes an ambiguity tube MPC for nonlinear stochastic systems with cost-to-go concavity.

\subsubsection*{Contributions}
We propose DR-PETS, a DR extension of PETS that explicitly certifies robustness against adversarial perturbations. %To our knowledge, DR-PETS is the first model-based RL method to certify robustness using Wasserstein ambiguity sets while preserving PETS' efficiency. 
Our key contributions include:  
(i) integrating Wasserstein ambiguity sets into the MPC cost functional for worst-case-aware planning without perturbing training data~\cite{derman2020distributionalregularization, queeney2023optimal};  
(ii) reformulating the intractable max-min problem via Wasserstein duality~\cite{blanchet2016quantifying}, yielding a convex maximization with a regularized cost;  
(iii) empirical validation on adversarial pendulum and cart-pole environments, demonstrating superior robustness over PETS;  
(iv) open-source implementation at \url{https://tinyurl.com/595c76t6}.  
Our work advances prior robust RL methods in three key ways: (1) unlike \cite{queeney2023optimal}, which computes adversarial perturbations offline by corrupting training data and then deriving a policy, DR-PETS directly optimizes for a robust policy by maximizing a regularized reward function, eliminating the need for explicit perturbation injection; (2) avoids iterative forward simulations, unlike~\cite{hewing2019cautious}, preserving PETS' data efficiency; and (3) explicitly certifies robustness against adversarial perturbations, unlike PETS, which only accounts for stochastic uncertainty. 
While this  work focuses on theoretical certification of robustness against adversarial perturbations in an unconstrained setting, it establishes a foundational framework for future research to rigorously integrate  safety constraints; indeed, as showed in~\cite{queeney2023optimal}, the inclusion of safety constraints significantly amplifies the benefits of robust methods by reducing sensitivity to model perturbations. This is the objective of future researches.
% We propose DR-PETS, a DR RL algorithm that extends the Probabilistic Ensembles with Trajectory Sampling (PETS) framework to guarantee robustness against adversarial environmental perturbations. To our knowledge, DR-PETS is the first model-based RL method to explicitly certify robustness in adversarial settings using Wasserstein ambiguity sets while retaining the computational efficiency of PETS. Our key technical contributions include:
% (i) a distributionally robust formulation of PETS that integrates Wasserstein ambiguity sets \cite{derman2020distributionalregularization} and \cite{queeney2023optimal} into the Model Predictive Control (MPC) cost functional, enabling worst-case-aware planning without perturbing training data;
% (ii) a tractable reformulation of the resulting max-min optimization problem via Wasserstein duality \cite{blanchet2016quantifying}, recasting it into a convex maximization with a regularized cost that avoids computationally expensive bilevel optimization;
% (iii) empirical validation in adversarial pendulum and cart-pole environments, demonstrating DR-PETS’ superior robustness over PETS under non-stationary parameter perturbations; and
% (iv) an open-source implementation of DR-PETS and our adversarial benchmarks at \url{https://tinyurl.com/595c76t6}.

% Our work advances prior robust RL methods in three key ways. First, unlike \cite{queeney2023optimal}, which computes adversarial perturbations offline by corrupting training data and then deriving a policy, DR-PETS directly optimizes for a robust policy by maximizing a regularized reward function, eliminating the need for explicit perturbation injection. Second, compared to distributionally robust MPC approaches \cite{hewing2019cautious}, our Wasserstein duality-based reformulation avoids iterative forward simulations, retaining PETS’ data efficiency. Third, while PETS implicitly handles stochastic uncertainty through ensembles, DR-PETS explicitly certifies robustness against adversarial distributions via its regularized cost, addressing a critical gap for safety-critical systems.
\section{Preliminaries and problem setup}
%In this section we first provide some preliminaries on the probabilistic ensembles with trajectory sampling (PETS) algorithm by \cite{chua2018}, and explore how the MPC objective can be approximated using samples observed along generated trajectories. Next, we formulate the optimization problem in a distributionally robust setting.
% Finally, some preliminaries on the Wasserstein metric are shown.
% \grtodo{This is largely a copy ad paste of the notation sections we have in the probabilistic papers (except the last few lines). Is it sure that the notation is covered? Already above I found a case where there is undefined notation (see previous comment). We do not deal with pmfs!}
Sets are in {\em calligraphic} and vectors in {\bf bold}. A random variable is denoted by $\bv{V}$ and its realization is $\bv{v}$. We denote the \textit{probability density function}, or simply \textit{pdf}, of $\bv{v}$  by $p(\bv{v})$. Whenever we take the integral involving a pdf we always assume that the integral exists. The  expectation of a function $\mathbf{h}(\cdot)$ of $\bv{v}$ is $\E_{{p}}[\mathbf{h}(\bv{v})]:=\int\mathbf{h}(\bv{v})p(\bv{v})d\bv{\bv{v}}$, where the integration is over the support of $p(\bv{v})$. The {conditional} pdf of $\bv{v}_1$ with respect to $\bv{v}_2$ is $p\left( \bv{v}_1\mid   \bv{v}_2 \right)$. Countable sets are denoted by $w_{k_1:k_n}$, where $w_k$ is the generic set element, $k_1$ ($k_n$) is the index of the first (last) element and  $k_1:k_n$ is the set of consecutive integers between (including) $k_1$ and $k_n$. 
% A pmf of the form $p(\bv{v}_0,\ldots,\bv{v}_N)$ is compactly written as $p_{0:N}$ (by definition $p_{k:k} := p_k(\bv{v}_k)$). We use the shorthand notation $p_{k\mid k-1}$ to denote $p_k(\bv{v}_k\mid\bf{v}_{k-1})$. 
We use $<\cdot,\cdot>$ as the inner product and $\bigotimes$ as the Cartesian product. Let $\mathcal{M}(\mathcal{B})$ represent set of distributions over a Borel set $\mathcal{B}$. Whenever it is clear from the context, a function $g(\cdot)$ is denoted by $g$. The norm is denoted by $\norm{\cdot}$ and the dual norm by $\norm{\cdot}_{*}$.  The symbol $\approx$ denotes approximations.

\subsection{Probabilistic ensembles with trajectory sampling (PETS) algorithm}
Consider a MDP~\cite{chua2018}  defined by the tuple $(\mathcal{X}, \mathcal{U}, f, r, \gamma)$, where $\mathcal{X}$ is the set of states and $\mathcal{U}$ is the set of actions, such that $f: \mathcal{X} \times \mathcal{U} \varrightarrow P(\mathcal{X}) $ is the transition probability function where $P(\mathcal{X})$ denotes space of probability measures over $\mathcal{X}$, $r : \mathcal{X} \times \mathcal{U}  \varrightarrow \mathbb{R}$ is the reward function, and $\gamma \in [0,1)$ is the discount factor. At time $t$, given state $\bv{x}_{t} \in \mathcal{X}\subseteq\R^{x}$, the agent applies action $\bv{u}_{t} \in \mathcal{U}\subseteq\R^{u}$, then the subsequent state $\bv{x}_{t+1}$ is determined by the conditional distribution $f(\bv{x}_{t+1} | \bv{x}_{t}, \bv{u}_{t})$ i.e. $\bv{x}_{t+1}\sim f(\bv{x}_{t+1} | \bv{x}_{t}, \bv{u}_{t}),\forall \bv{x}_{t+1} \in \mathcal{X}$, and $r(\bv{x}_{t},\bv{u}_{t})$  returns the immediate reward received when the MDP is in state $\bv{x}_{t}$ and $\bv{u}_{t}$ is applied. The results in \cite{chua2018}  rely on the following:

\begin{assumption}\label{assm:MDP}
     Given an MDP, $(\mathcal{X}, \mathcal{U}, f, r, \gamma)$ the transition probability function $f$ is unknown and reward $r$ is known.
\end{assumption}

Under Assumption \ref{assm:MDP}, to estimate the transition probability function \( f \), PETS collects a dataset $\mathcal{D}$ of $N \in \mathbb{Z}_{+}$ samples from the system, defined as \(\{(\bv{x}_t, \bv{u}_t, \bv{x}_{t+1})\}_{t=0,1,\dots,N-1}\), where the actions are generated according to some criteria, e. g. randomly. Then an ensemble of \( B \in \mathbb{Z}_{+} \) neural networks (NNs) can be trained on the collected dataset; formally,  the $b$-th NN  is modelled as a Gaussian distribution by $\hat{f}_{\theta_b} \sim \mathcal{N}(m_{\theta_b}, \Sigma_{\theta_b})$, with $b=1,2,\dots,B$ and $\theta_{b}$ being the parameters of $b$-th NN. As in \cite{chua2018}, given a system state $\bv{x}_t$ and an action $\bv{u}_t$, the output of the ensemble, i.e. the probability to reach the new expected state $\hat{\bv{x}}_{t+1}$, is the average over the NNs, i.e., $\hat{f}_{\theta}(\hat{\bv{x}}_{t+1}|\bv{x}_t,\bv{u}_t) := \frac{1}{B} \sum\limits_{b=1}^B \hat{f}_{\theta_b}(\hat{\bv{x}}_{t+1}|\bv{x}_t,\bv{u}_t)$.

For planning and action selection, PETS uses an MPC-based policy. For a given current state $\hat{\bv{x}}_t$, system model $\hat{f}_{\theta}$ and a finite horizon $T\in\mathbb{Z}_{+}$, the RL planning through MPC maximizes the objective function 
\begin{equation}\label{eqn:Chua_obj}
\begin{split}
     \mathcal{J}_{\hat{f}_{\theta}}(\hat{\bv{x}}_t, \bv{u}_{t:t+T-1})=\mathbb{E}_{ \hat{f}_{\theta}}\left[\sum\limits_{i=t}^{t+T-1}\gamma^{i-t}{r}(\hat{\bv{x}}_i,\bv{u}_i) | \hat{\bv{x}}_t\right],
 \end{split}   
\end{equation}
where we denote a simulated state by $\hat{\bv{x}}_{i}$, and for $i=t$ we initialize $\hat{\bv{x}}_t$ to $\bv{x}_t$ as system state. In \eqref{eqn:Chua_obj} the expectation is over simulated future trajectories generated by $\hat{\bv{x}}_{i+1}\sim\hat{f}_{\theta}(\cdot|\hat{\bv{x}}_{i},\bv{u}_{i})$ given the initial state $\bv{x}_{t}$.

With this setting, PETS algorithm aims to find a finite sequence $\bv{u}_{t:t+T-1}^{*} := \{\bv{u}_{t}^{*}, \bv{u}_{t+1}^{*}, \dots, \bv{u}_{t+T-1}^{*}\}$ such that
\begin{equation}\label{eqn:MPC_prob}
\begin{split}
     &\bv{u}_{t:t+T-1}^{\ast} 
     \in \arg \max\limits_{\mathbf{u}_{t:t+T-1}}  \mathcal{J}_{\hat{f}_{\theta}}(\hat{\bv{x}}_t, \bv{u}_{t:t+T-1})\\
     &s.t. \quad \bv{u}_{t:t+T-1}\in\mathcal{U}.
\end{split}
\end{equation}

% In an infinite horizon MDP setting the objective is to find a policy $\pi:\mathcal{X} \varrightarrow \mathcal{U}$ that maximizes the expected cumulative sum of future discounted rewards, namely

% % \begin{equation}\label{eqn:MDP_obj}
% %     \sum\limits_{i=t}^{\infty}\mathbb{E}_{x_{i+1}\sim f}[\gamma^{i-t}r(\bv{x}_i,\pi(\bv{x}_i)) | \bv{x}_t],
% % \end{equation}

% \begin{equation}\label{eqn:MDP_obj}
%    \mathbb{E}_{f}\left[\sum\limits_{i=t}^{\infty}\gamma^{i-t}r(\bv{x}_i,\pi(\bv{x}_i)) | \bv{x}_t=\bv{x}\right],
% \end{equation}
% where the expectation is over trajectory generated by $\bv{x}_{i+1}\sim f(\cdot\mid\bv{x}_{i},\bv{u}_{i})$ given the initial state $\bv{x}\in\mathcal{X}$.
% % $\tau\sim(\pi,f)$ is the trajectory sampled using $\pi$ and $f(\bv{x}_{t+1} | \bv{x}_t, \bv{u}_t)$ \, 
% The objective function~\eqref{eqn:MDP_obj} can be expressed as the sum of the expected cumulative reward over a finite horizon $T \in \mathbb{Z}_{+}$ and a reward-to-go. With a proper choice of $\gamma$ and $T$, the reward-to-go can be neglected. 
% For planning and action selection PETS uses a model predictive control (MPC)-based policy, that is given the current state $\bv{x}_t$, the RL planning through MPC maximizes the objective function in \eqref{eqn:MPC_prob} over $T$ steps, and then the controller executes only the first action $\bv{u}_t^{\ast}$, observes the new state $\bv{x}_{t+1}$, and the process repeats. In what follows we denote the objective function of the problem \eqref{eqn:MPC_prob} as $\mathcal{J}_{f}(\bv{x}_{t},\bv{u}_{t:t+T-1})$.



% Then, the objective function in~\eqref{eqn:MPC_prob} can be rewritten as
% \begin{equation}\label{eqn:Chua_obj}
% \begin{split}
%      \mathcal{J}_{\hat{f}_{\theta}}(\hat{\bv{x}}_t, \bv{u}_{t:t+T-1})=\mathbb{E}_{ \hat{f}_{\theta}}\left[\sum\limits_{i=t}^{t+T-1}\gamma^{i-t}{r}(\hat{\bv{x}}_i,\bv{u}_i) | \hat{\bv{x}}_t\right],
%  \end{split}   
% \end{equation}


% \subsection{Problem setup}

% \grtodo{I moved this paragraph here, but it read confusing to me...I rewrote it, see if it is still correct/what you meant...is it sure you meant a "set" of transitions? Also, we never really said what is an ambiguity set...}

% \grtodo{PLEASE put here the assumption for the linearization, saying that this is standard and put the reference we have...}

\subsection{Distributionally robust MDPs}

% In Distributionally robust MDP (DRMDP) the 
% Let $\mathcal{F}_{\bv{x},\bv{u}}$ by a set of transition probability functions,  $f(\cdot|\bv{x},\bv{u})\in P(\mathcal{X})$, given a state-action pair. Also, let $\mathcal{F}:=\bigotimes_{(\bv{x},\bv{u})\in\mathcal{X}\times\mathcal{U}}\mathcal{F}_{\bv{x},\bv{u}}$.  Hence, $\mathcal{F}$ is a Borel set and we let any possible transition probability function $\tilde{f}$ be a random variable with support on $\mathcal{F}$, such that $\tilde{f}\sim\tilde{\mu}\in\mathcal{M}(\mathcal{F})$.

% %We consider distributionally robust MDP (DRMDP), where unknown transition probability function $f$ is assumed to be within an uncertainty set $\mathcal{F}$, i.e. $f\in\mathcal{F}$. As done in \cite{chen2019distributionally,wiesemann2013robust}, we assume that the uncertainty set $\mathcal{F}$ is $(\bv{x},\bv{u})$-rectangular i.e. $\mathcal{F}=\bigotimes_{(\bv{x},\bv{u})\in\mathcal{X}\times\mathcal{U}}\mathcal{F}_{\bv{x},\bv{u}}$ where $\mathcal{F}_{\bv{x},\bv{u}}$ is a set of transitions $f(\cdot|\bv{x},\bv{u})\in P(\mathcal{X})$ for a given state-action pair, hence, $\mathcal{F}$ is a Borel set and we let any possible transition probability function $\tilde{f}$ be a random variable with support on $\mathcal{F}$, such that $\tilde{f}\sim\tilde{\mu}\in\mathcal{M}(\mathcal{F})$. 
% Inspired by~\cite{wu2022ambiguity,shi2022distributionally},  we use the $p$-Wasserstein ball $\Omega_{\epsilon}(\hat{\mu})$ to formalize the ambiguity set over $\mathcal{F}$, centered in $\hat{\mu}$. This is defined as \[\Omega_{\epsilon}(\hat{\mu}) := \{\tilde{\mu}: \mathcal{W}_{p}(\tilde{\mu},\hat{\mu}) \leq \epsilon\},\] where $\epsilon>0$ is the radius of the ball, and $\hat{\mu} := \frac{1}{B} \sum\limits_{b=1}^B \delta_{\hat{f}_{\theta_b}}$, with $\delta_{\hat{f}_{\theta_b}}$ being the Dirac distribution with full mass on $\hat{f}_{\theta_{b}}$.  We also make the following assumption on the ambiguity set
We introduce distributionally robust MDPs (DR-MDPs) following \cite{derman2020distributionalregularization}, where the unknown transition probability function $f$ is within an uncertainty set $\mathcal{F}$, i.e. $f\in\mathcal{F}$.  As commonly assumed in the literature \cite{queeney2023optimal,derman2020distributionalregularization}, we let $\mathcal{F}$ to  be $(\bv{x},\bv{u})$-rectangular i.e. $\mathcal{F}=\bigotimes_{(\bv{x},\bv{u})\in\mathcal{X}\times\mathcal{U}}\mathcal{F}_{\bv{x},\bv{u}}$ where $\mathcal{F}_{\bv{x},\bv{u}}$ is a set of transitions $f(\cdot|\bv{x},\bv{u})\in P(\mathcal{X})$ for a given state-action pair. Hence, $\mathcal{F}$ is a Borel set and we let any possible transition probability function $\tilde{f}$ be a random variable with support on $\mathcal{F}$, such that $\tilde{f}\sim\tilde{\mu}\in\mathcal{M}(\mathcal{F})$. Here, the class of the distributions $\mathcal{M}$ is the ambuiguity set each of which have support on $\mathcal{F}$.

\begin{definition}\label{def:Wasserstein}($p$-Wasserstein distance \cite{blanchet2016quantifying})
     Let $p\in[1,+\infty]$ and $(\mathcal{F},d)$ be a metric space, where  $d:\mathcal{F}\times\mathcal{F}\rightarrow\mathbb{R}_{+}$ is a lower semicontinuous metric. The $p$-Wasserstein distance  between $\hat{\mu},\tilde{\mu}\in\mathcal{M}$($\mathcal{F}$) is
\begin{equation}\label{eqn:wasserstein}
        \mathcal{W}_{p}(\tilde{\mu},\hat{\mu}) = \inf\limits_{\psi\in\mathcal{Q}(\hat{\mu},\tilde{\mu})}\left(\int_{\mathcal{F}\times\mathcal{F}} d(\hat{f},\tilde{f})^{p} d\psi(\hat{f},\tilde{f})\right)^{\frac{1}{p}},
    \end{equation}
    where $\mathcal{Q}(\hat{\mu},\tilde{\mu})$ is the set of distributions over $\mathcal{F}\times\mathcal{F}$ with marginals $\hat{\mu}$ and $\tilde{\mu}$.
\end{definition}

Inspired by~\cite{derman2020distributionalregularization} and \cite{shi2022distributionally}, we choose a $p$-Wasserstein ball $\Omega_{\epsilon}(\hat{\mu})$ as the ambiguity set over $\mathcal{F}$, centered in $\hat{\mu}$ and defined as 
\begin{equation}\label{eqn:omega}
\Omega_{\epsilon}(\hat{\mu}) := \{\tilde{\mu}: \mathcal{W}_{p}(\tilde{\mu},\hat{\mu}) \leq \epsilon\}
\end{equation}
where $\epsilon>0$ is the radius of the ball, and $\hat{\mu} := \frac{1}{B} \sum\limits_{b=1}^B \delta_{\hat{f}_{\theta_b}}$, with $\delta_{\hat{f}_{\theta_b}}$ being the Dirac distribution with full mass on $\hat{f}_{\theta_{b}}$.

\subsection{Problem statement}
Following \cite{derman2020distributionalregularization}, we let the $p$-Wasserstein DR-MDP (WR-MDP) tuple as $(\mathcal{X},\mathcal{U}, r, \mathcal{F}, \Omega_{\epsilon}(\hat{\mu}))$, and we define corresponding DR formulation of \eqref{eqn:MPC_prob} as: 
\begin{Problem}\label{prob1}

Given a WR-MDP $(\mathcal{X},\mathcal{U}, r, \mathcal{F}, \Omega_{\epsilon}(\hat{\mu}))$, find the sequence of control inputs $\mathbf{u}_{t:t+T-1}^{*} := \{\bv{u}_{t}^{*}, \bv{u}_{t+1}^{*}, \dots, \bv{u}_{t+T-1}^{*}\}$ such that,
\begin{align}\label{eqn:DR_obj}
    &\mathbf{u}_{t:t+T-1}^{*}\in \arg \max\limits_{\mathbf{u}_{t:t+T-1}} \inf_{\tilde{\mu}\in \Omega_{\epsilon}(\hat{\mu})}\mathbb{E}_{\tilde{f}\sim\tilde{\mu}}\left[\mathcal{J}_{\tilde{f}}(\bv{x}_t, \bv{u}_{t:t+T-1})\right] \notag\\
    & s.t. \quad {\bv{u}_t, \bv{u}_{t+1}, \dots, \bv{u}_{t+T-1}} \in \mathcal{U},
\end{align}
    where $\mathcal{J}_{\tilde{f}}(\bv{x}_t, \bv{u}_{t:t+T-1})$ is defined as in \eqref{eqn:Chua_obj}.
    % \[\mathcal{J}_{\tilde{f}}(\bv{x}_t, \bv{u}_{t:t+T-1})=\mathbb{E}_{\tilde{f}}\left[\sum \limits_{i=t}^{t+T-1}\gamma^{i-t}r(\hat{\bv{x}}_i,\bv{u}_i)\mid \bv{x}_{t}\right].\]
\end{Problem}
% \textcolor{blue}{I do not understand the next sentence. What is the problem we are trying to solve? Are we giving a reformulation of (4)?}
By evaluating the infimum over the ambiguity set $\Omega_{\epsilon}(\hat{\mu})$ in \eqref{eqn:DR_obj} we consider a worst-case scenario for the objective function in \eqref{eqn:Chua_obj}, and this corresponds to a robust formulation of \eqref{eqn:MPC_prob}. 
% Thus, we aim to find a lower bound of~\eqref{eqn:Chua_obj}, namely~\eqref{eqn:DR_obj}, that holds with some probability depending on the radius of a Wasserstein ball centered in $\hat{\mu}$.
To lighten the notation, for any given transition probability function $\tilde{f}$ and time sequence $t+i:t+T-1$ starting from current time $t+i$, we denote $\mathcal{J}_{\tilde{f}}(\bv{x}_{t+i}, \bv{u}_{t+i:t+T-1})$ with $\mathcal{J}_{t+i}(\tilde{f})$. In the subsequent analysis we make the following  assumption:
\begin{assumption}\label{assm:linearity} 
Consider Problem \ref{prob1} with $\Omega_{\epsilon}(\hat{\mu})$ defined in~\eqref{eqn:omega}.  The radius of the ball, $\epsilon$, is such that
$
\mathcal{J}_{t}(\bar{f}) \approx \mathcal{J}_{t}(\hat{f}_{\theta_{b}})+\left<\nabla_{\hat{f}_{\theta_{b}}}\mathcal{J}_{t}(\hat{f}_{\theta_{b}}),v_{b}\right>,
$
for all $\bar{f} \in \Omega_{\epsilon}(\hat{\mu})$.
%\AAC{Given the $p$-Wasserstein distance and an adversarial perturbation $v_{b}$  such that $d(\tilde{f}_{b},\hat{f}_{\theta_b})=\norm{v_{b}}$ with $\tilde{f}_{b}=\hat{f}_{\theta_b}+v_{b}$ and $\frac{1}{B}\sum\limits_{b=1}^{B}\norm{v_{b}}^{p}\leq\epsilon^{p}$, 
%then  when $\epsilon$ is sufficiently small.}
\end{assumption}
\begin{Remark}
     Assumption \ref{assm:linearity} is standard in the literature on DR-RL (see e.g.~\cite{derman2020distributionalregularization},~\cite{queeney2023optimal} and~\cite{goodfellow2014explaining}). Intuitively, it implies that $\epsilon$ is sufficiently small, i.e.,  within the feasibility domain of Problem \ref{prob1} a first order approximation for the cost holds.
\end{Remark}
%\grtodo{Ass. 2 and the remark are the major things I changed.}

Authors in \cite{shi2022distributionally} and \cite{queeney2023optimal} showed that the objective function in Problem \ref{prob1} is intractable. Our main objective is to derive a tractable reformulation of cost function in \eqref{eqn:DR_obj}. We do this by leveraging the strong duality of \eqref{eqn:DR_obj}, as established in \cite{blanchet2016quantifying},
\begin{align}\label{eqn:DR_duality}
    &\inf_{\tilde{\mu}\in \Omega_{\epsilon}(\hat{\mu})}\mathbb{E}_{\tilde{f}\sim\tilde{\mu}}\mathcal{J}_{t}(\tilde{f})\notag\\ 
    &= \inf_{\tilde{\mu}} \sup\limits_{\lambda\geq0}\Biggl\{\mathbb{E}_{\tilde{f}\sim\tilde{\mu}}\left[\mathcal{J}_{t}(\tilde{f})\right]\\
    &\qquad \qquad \qquad-\lambda\Biggl(\epsilon^{p}-\frac{1}{B}\sum\limits_{b=1}^{B}\mathbb{E}_{{f_{b}}\sim{\mu_{b}}}\left[d(f_{b},\hat{f}_{\theta_{b}})^{p}\right]\Biggr)\Biggr\}\notag\\
     &= \sup\limits_{\lambda\geq0}\Biggl\{\frac{1}{B}\sum\limits_{b=1}^B\inf_{\tilde{f}_{b}}\left\{\mathcal{J}_{t}(\tilde{f}_{b})+\lambda(d(\tilde{f}_{b},\hat{f}_{\theta_{b}})^{p})\right\}-\lambda\epsilon^{p}\Biggr\}\notag.
\end{align}
Our approach aligns with the DR framework proposed in [11], although  developed for a finite discrete state space.  Indeed,  for our framework of a continuous state space we use of an ensemble of neural networks to model the transition dynamics $f$ and the subsequent derivation of a regularized reward function through a Wasserstein ambiguity set.
%In a similar context the similar approach has been followed by the authors in \cite{derman2020distributionalregularization}.
%\AAC{However, the difference between the approach in~\cite{derman2020distributionalregularization} is that we leverage an ensemble of neural network to model the transition probability function $f$, and formulate the regularized reward function. }
% \grtodo{This is a key point to be addressed. ...I would bring some of this to the contributions section.}
% \begin{algorithm}
% \caption{Probabilistic ensembles with trajectory sampling (PETS) \cite{chua2018}}
% \begin{algorithmic}[1]
% \State Initialize data $\mathcal{D}$ with a random controller for one trial.
% \For{Trial $k = 1$ to $K$}
%     \State Train a probabilistic ensemble dynamics model $\hat{f}_{\theta}$ given data-set $\mathcal{D}$.
%     \For{Time $t = 0$ to $T$}
%         \For{Actions sampled $\bv{u}_{t:t+T} \sim \mathbf{CEM}(\cdot), 1$ to $N$ samples} 
%         \State Propagate $p$ state particles $\bv{x}_t^p$ using trajectory sampling and $\hat{f}_{\theta}\mid\{D, \bv{u}_{t:t+T}\}$.
%         \State Evaluate actions as per \eqref{eqn:Chua_obj}
%         \State Update $CEM(\cdot)$ distribution.
%         \State Execute first action $\bv{u}_t^*$ (only) from optimal actions $\bv{u}_{t:t+T}^*$.
%         \State Record outcome: $\mathcal{D} \gets \mathcal{D} \cup \{\bv{x}_t, \bv{u}_t^*, \bv{x}_{t+1}\}$.
%         \EndFor
%     \EndFor
% \EndFor
% \end{algorithmic}
% \end{algorithm}


% To compute the expectation in \eqref{eqn:MPC_obj} we need knowledge about system's dynamics $f$, that in the general RL framework is not known.

% \subsection{Empirical MPC objective}

% Consider an MDP $(\mathbb{R}^{d_x}, \mathbb{R}^{d_u}, f, r)$ with $f$ unknown, and assume that the deterministic reward function $r$ is known. Under these assumptions, a way to maximize \eqref{eqn:MPC_obj} is to use samples to train a parametric function $\hat{f}_{\theta}(\bv{x}_{t+1}|\bv{x}_t,\bv{u}_t) \approx f(\bv{x}_{t+1}|\bv{x}_t,\bv{u}_t),\;\forall \bv{x}_t, \forall \bv{u}_t$, and use it to estimate the expectation. Specifically, assume that we collect a data set $\mathcal{D} := \{(\bv{x}_t,\bv{u}_t,\bv{x}_{t+1})\}_{t=0,1,\dots,N_{init}}$ using the true system, where the actions are generated according to some criteria, e.g., selecting random actions. We can then use $\mathcal{D}$ to train an approximation $\hat{f}_{\theta} \approx f$ of the true model. 

% However, because $\hat{f}_{\theta}$ is trained on a finite set $\mathcal{D}$, the model uncertainty can significantly affect the optimization objective, leading to non sense actions.

% A solution to this problem is to try to model the uncertainty about the model, i.e. the uncertainty about the parameters $\theta$. In the literature, the authors in \cite{chua2018} approximate the true dynamics with an ensemble of $B$ NNs. The intuition is that ideally if we diversify the $B$ NNs (e.g. with sample with replacement from $\mathcal{D}$), they will make different errors in the models in different directions. 

% Formally, we can define the general structure of the $b$-th NN as $\hat{f}_{\theta_b} \sim \mathcal{N}(m_{\theta_b}, \Sigma_{\theta_b})$, $b=1,2,\dots,B$. Given a state $\bv{x}_t$ from the true system and an action $\bv{u}_t$ is selected, the output of the ensemble, i.e. the probability to reach the new state $\hat{\bv{x}}_{t+1}$, is the average over the NNs, namely $\hat{f}_{\bar{\theta}}(\hat{\bv{x}}_{t+1}|\bv{x}_t,\bv{u}_t) := \frac{1}{B} \sum\limits_{b=1}^B \hat{f}_{\theta_b}(\hat{\bv{x}}_{t+1}|\bv{x}_t,\bv{u}_t)$. Then the model-based (MB) RL objective via MPC becomes

% \begin{equation}\label{eqn:Chua_obj}
%     \max\limits_{\bv{u}_t, \bv{u}_{t+1}, \dots, \bv{u}_{t+T-1}}  \sum\limits_{i=t}^{t+T-1}\mathbb{E}_{\tau\sim (\pi,\hat{f}_{\bar{\theta}})}[\gamma^{i-t}r(\hat{\bv{x}}_i,\bv{u}_i) | \bv{x}_t],
% \end{equation}

% where $\hat{\bv{x}}_t \equiv \bv{x}_t$ from the real system, and we denote a simulated state by $\hat{\bv{x}}$. 

% The main advantage of maximizing \eqref{eqn:Chua_obj} is that using $B$ NNs alleviates the effects about the uncertainty in the model. However, the MPC method itself is not very efficient (due to the fact that we solve \eqref{eqn:Chua_obj} at each $t$ but then we apply only $\bv{u}_t$). Moreover, even if the output from the $B$ NNs can be computed in parallel, it involves more computational load w.r.t. a single NN. Finally, for some applications in which the true $f$ is complex, or where $\mathcal{D}$ is small, the output of the ensemble $\hat{f}_{\bar{\theta}}$ can still exhibit some uncertainties.

% In the next section, we show how to apply the idea of the ensembles to robustly maximize the objective in \eqref{eqn:MPC_obj}, by maximizing a distributionally robust (DR) approximated lower bound of it.
% \subsection{Problem setup}

% In the DR settings, assume that $\mu : \mathcal{F} \varrightarrow [0, 1]$ is the distribution over the models, i.e., $\mu(f)$ is the probability of transition model $f$, where $\mathcal{F}$ is the space of all possible transition models.
% % $\mathcal{F} \varsubseteq \mathbb{R}^{d_x} \times \mathbb{R}^{d_u} \times \mathbb{R}^{d_x}$ is the support of the transition model $f$. 
% Moreover, consider $B$ empirical models sampled from $\mu$, i.e., $\hat{f}_{\theta_1}, \hat{f}_{\theta_2}, \dots, \hat{f}_{\theta_B} \sim \mu$, and let $\hat{\mu} := \frac{1}{B} \sum\limits_{b=1}^B \delta_{\hat{f}_{\theta_b}}$. 

% Consider a $p$-Wasserstein ball with radius $\epsilon$ in the space of distributions over models, centered around $\hat{\mu}$, namely $\Omega_{\epsilon}(\hat{\mu}) = \{\tilde{\mu}: \mathcal{W}(\tilde{\mu},\hat{\mu}) \leq \epsilon\}$. If $\epsilon$ is chosen such that $\mu \in \Omega_{\epsilon}(\hat{\mu})$, then for any function $\mathcal{J}_{f}$ we have

% \begin{equation}\label{eqn:DR_prob}
%     \mathbb{E}_{f\sim \mu}[\mathcal{J}_{f}] \geq \inf_{\tilde{\mu} \in \Omega_{\epsilon}(\hat{\mu})} \mathbb{E}_{f \sim \tilde{\mu}}[\mathcal{J}_{f}].
% \end{equation}

% Let the function $\mathcal{J}_f$ be defined as

% \begin{equation}\label{eqn:plan3}
%     \mathcal{J}_{f}(\bv{x}_t, \bv{u}_{t:t+T-1}) := \sum \limits_{i=t}^{t+T-1}\mathbb{E}_{{x}_{i+1 \sim }{f}}\bigr[\gamma^{i-t}r({x}_i,\bv{u}_i)\bigl],
% \end{equation}

% where the expectation is w.r.t. the probabilistic model. Thus, \eqref{eqn:plan3} is the performance index in \eqref{eqn:MPC_obj} that we aim to maximize.

% Plugging \eqref{eqn:plan3} into \eqref{eqn:DR_prob} we have 

% \begin{equation}\label{eqn:DR_start}
% \begin{aligned}
%     &\mathbb{E}_{{f}\sim {\mu}}\biggr[\mathcal{J}_{{f}}(\bv{x}_t, \bv{u}_{t:t+T-1})\biggl] \geq \mathcal{J}_{\Omega_{\epsilon}(\hat{\mu})}(\bv{x}_t, \bv{u}_{t:t+T-1})\\ 
%     &:= \inf_{\tilde{\mu}\in \Omega_{\epsilon}(\hat{\mu})}\mathbb{E}_{\tilde{f} \sim \tilde{\mu}}\biggr[\sum \limits_{i=t}^{t+T-1}\mathbb{E}_{\hat{\bv{x}}_{i+1} \sim \tilde{f}}\bigr(\gamma^{i-t}r(\hat{\bv{x}}_i,\bv{u}_i)\bigl)\biggl],\\
%     &= \inf_{\tilde{\mu}: \mathcal{W}(\tilde{\mu},\hat{\mu}) \leq \epsilon}\mathbb{E}_{\tilde{f} \sim \tilde{\mu}}\biggr[\sum \limits_{i=t}^{t+T-1}\mathbb{E}_{\hat{\bv{x}}_{i+1} \sim \tilde{f}}\bigr(\gamma^{i-t}r(\hat{\bv{x}}_i,\bv{u}_i)\bigl)\biggl].
% \end{aligned}
% \end{equation}
 
% Where, $\mathcal{W}(\tilde{\mu},\hat{\mu})$ is the Wasserstein metric defined as,
% \begin{equation*}
%     \mathcal{W}(\tilde{\mu},\hat{\mu}) = 
% \end{equation*}




% If $\epsilon$ is chosen such that $\mu \in \Omega_{\epsilon}(\hat{\mu})$, then for any function $\mathcal{J}_{f}$ we have

% \begin{equation}\label{eqn:DR_prob}
%     \mathbb{E}_{f\sim \mu}[\mathcal{J}_{f}] \geq \inf_{\tilde{\mu} \in \Omega_{\epsilon}(\hat{\mu})} \mathbb{E}_{f \sim \tilde{\mu}}[\mathcal{J}_{f}].
% \end{equation}

% Let the function $\mathcal{J}_f$ be defined as

% \begin{equation}\label{eqn:plan3}
%     \mathcal{J}_{f}(\bv{x}_t, \bv{u}_{t:t+T-1}) := \sum \limits_{i=t}^{t+T-1}\mathbb{E}_{\tau\sim (\pi,f)}\bigr[\gamma^{i-t}r({x}_i,\bv{u}_i)\bigl],
% \end{equation}

% where the expectation is w.r.t. the probabilistic model. Thus, \eqref{eqn:plan3} is the performance index in \eqref{eqn:MPC_obj} that we aim to maximize.

% Plugging \eqref{eqn:plan3} into \eqref{eqn:DR_prob} we have 

% \begin{equation}\label{eqn:DR_start}
% \begin{aligned}
%     \mathbb{E}_{{f}\sim {\mu}}\biggr[\mathcal{J}_{{f}}(\bv{x}_t, \bv{u}_{t:t+T-1})\biggl] \geq \mathcal{J}_{\Omega_{\epsilon}(\hat{\mu})}(\bv{x}_t, \bv{u}_{t:t+T-1}) :&= \inf_{\tilde{\mu}\in \Omega_{\epsilon}(\hat{\mu})}\mathbb{E}_{\tilde{f} \sim \tilde{\mu}}\biggr[\sum \limits_{i=t}^{t+T-1}\mathbb{E}_{\tau\sim(\pi,\tilde{f})}\bigr(\gamma^{i-t}r(\hat{\bv{x}}_i,\bv{u}_i)\bigl)\biggl],\\
%     &= \inf_{\tilde{\mu}: \mathcal{W}(\tilde{\mu},\hat{\mu}) \leq \epsilon}\mathbb{E}_{\tilde{f} \sim \tilde{\mu}}\biggr[\sum \limits_{i=t}^{t+T-1}\mathbb{E}_{\tau\sim (\pi,f)}\bigr(\gamma^{i-t}r(\hat{\bv{x}}_i,\bv{u}_i)\bigl)\biggl].
% \end{aligned}
% \end{equation}

\section{Main results}
To develop the DR-PETS algorithm we need a tractable reformulation of \eqref{eqn:DR_obj} for which we employ the last equality in \eqref{eqn:DR_duality}. In fact, in the following Theorem we show that \eqref{eqn:DR_duality} is formed of two terms: the empirical estimate of the performance index over the ensemble and a regularizer. 
% \begin{lem}\label{lem:Jimmy}
% Wasserstein DR policy optimization: derivation using first order approximation.\\
% Consider the $p$-Wasserstein distance, $p \in (1,\infty)$, and let $d(\mathbf{z} + v) = \norm{v}$ for some general norm. From Wasserstein duality, we are interested in developing a tractable reformulation of 
% \[\sup_{\lambda \geq 0}\biggr\{\frac{1}{n}\sum\limits_{i=1}^n[\inf_{\mathbf{z}_i}l(\mathbf{z}_i) + \lambda(d(\mathbf{z}_i, \hat{\mathbf{z}}_i)^p)] - \lambda \epsilon^p\biggl\}.\]
% Using a first order approximation $l(\hat{\mathbf{z}} + v_i) = l(\hat{\mathbf{z}}_i) + \langle g_i, v_i\rangle$, the closed-form
% solution for the Wasserstein DR objective is given by
% \[\frac{1}{n}\sum\limits_{i=1}^n l(\hat{\mathbf{z}}_i) - \epsilon\biggr( \frac{1}{n}\sum\limits_{i=1}^n \norm{g_i}_{\ast}^q\biggl)^{\frac{1}{q}}.\]
% \end{lem}
\begin{thm}\label{thm:Jimmy}
Let Assumption \ref{assm:linearity} hold. Then:
\begin{align}\label{eqn:thm_1_statement}
    &\sup_{\lambda \geq 0}\biggr\{\frac{1}{B}\sum\limits_{b=1}^B\inf_{\tilde{f}_b}
    \bigr[\mathcal{J}_{t}(\tilde{f}_{b}) + \lambda(d(\tilde{f}_b, \hat{f}_{\theta_{b}})^{p})\bigl] - \lambda \epsilon^p\biggl\}\notag
    \\
    &=\frac{1}{B}\sum\limits_{b=1}^B \mathcal{J}_{t}(\hat{f}_{\theta_{b}}) - \epsilon\biggr( \frac{1}{B}\sum\limits_{b=1}^B \norm{\nabla_{\hat{f}_{\theta_{b}}}\mathcal{J}_{t}(\hat{f}_{\theta_{b}})}_{\ast}^{\frac{p}{p-1}}\biggl)^{\frac{p-1}{p}}.
\end{align}

% where let $l(\hat{f}_{\theta_{b}}) = \mathbb{E}_{\hat{f}_{\theta_{b}}}\biggr[\sum \limits_{i=t}^{t+T-1}\gamma^{i-t}r(\hat{\bv{x}}_i,\bv{u}_i)\biggl]$ and $g_{b}(\hat{f}_{\theta_{b}}) = \nabla_{\hat{f}_{\theta_{b}}}\mathbb{E}_{\hat{f}_{\theta_{b}}}\bigr[\sum \limits_{i=t}^{t+T-1}\gamma^{i-t}r(\hat{\bv{x}}_i,\bv{u}_i)\bigl].$
\end{thm}

\begin{proof}
Motivated by the approach of \cite{goodfellow2014explaining} and relying on Assumption \ref{assm:linearity} we can write the last equality in \eqref{eqn:DR_duality} as
 \begin{align}\label{eqn:dual_problem_recast}
             &\sup_{\lambda \geq 0}\biggr\{\frac{1}{B}\sum\limits_{b=1}^B\inf_{\tilde{f}_b}
    \left[\mathcal{J}_{t}(\tilde{f}_{b}) + \lambda(d(\tilde{f}_b, \hat{f}_{\theta_{b}})^{p})\right] - \lambda \epsilon^p\biggl\}\notag\\
    % &=\sup_{\lambda \geq 0}\biggr\{\frac{1}{B}\sum\limits_{b=1}^B\inf_{v_{b}}\left[\mathcal{J}_{t}(\hat{f}_{\theta_{b}}+v_{b}) + \lambda(\norm{v_{b}}^{p})\right] - \lambda \epsilon^p\biggl\}\\
    &\approx\frac{1}{B}\sum\limits_{b=1}^{B}\mathcal{J}_{t}(\hat{f}_{\theta_{b}})+\sup_{\lambda \geq 0}\biggr\{\frac{1}{B}\sum\limits_{b=1}^B\inf_{v_{b}}\Bigl[\left<\nabla_{\hat{f}_{\theta_{b}}}\mathcal{J}_{t}(\hat{f}_{\theta_{b}}),v_{b}\right>\notag\\ 
    & \qquad \qquad + \lambda(\norm{v_{b}}^{p})\Bigr] - \lambda \epsilon^p\biggl\}.
 \end{align}
% \grtodo{Note that I changed to approximation the last line.}
% By plugging \eqref{eqn:first_order_approx} into \eqref{eqn:dual_problem}, we obtain
% \begin{equation}\label{eqn:dual_problem_recast}
%     \begin{split}
%     &\sup_{\lambda \geq 0}\biggr\{\frac{1}{B}\sum\limits_{b=1}^B\inf_{v_{b}}
%     \left[\mathcal{J}_{t}(\hat{f}_{\theta_{b}}+v_{b}) + \lambda(\norm{v_{b}}^{p})\right] - \lambda \epsilon^p\biggl\}\\
%     &=\frac{1}{B}\sum\limits_{b=1}^{B}\mathcal{J}_{t}(\hat{f}_{\theta_{b}})+\sup_{\lambda \geq 0}\biggr\{\frac{1}{B}\sum\limits_{b=1}^B\inf_{v_{b}}\Bigl[\left<\nabla_{\hat{f}_{\theta_{b}}}\mathcal{J}_{t}(\hat{f}_{\theta_{b}}),v_{b}\right>\\ 
%     & \qquad \qquad + \lambda(\norm{v_{b}}^{p})\Bigr] - \lambda \epsilon^p\biggl\}.
%     \end{split}
%  \end{equation}

We begin our proof by solving the inner minimization problem in \eqref{eqn:dual_problem_recast}; we define $j(v_b)$ as
\begin{equation}\label{eqn:inner_prob}
\inf_{v_{b}} j(v_{b}) := \inf\limits_{v_{b}}\left\{\left<\nabla_{\hat{f}_{\theta_{b}}}\mathcal{J}_{t}(\hat{f}_{\theta_{b}}),v_{b}\right>+\lambda\norm{v_{b}}^{p}\right\},
\end{equation}
% Note that value of $\lambda$ determined by the solving the outer maximization in \eqref{eqn:dual_problem_recast} (\AAC{what do you mean?}).
where we introduce the auxiliary variable $\delta_{b}$, such that $\norm{v_{b}}\geq\delta_{b}$, obtaining
\begin{equation}\label{eqn:inner_prob_aux}
\begin{split}
&\inf\limits_{v_{b},\delta_{b}}\left\{\left<\nabla_{\hat{f}_{\theta_{b}}}\mathcal{J}_{t}(\hat{f}_{\theta_{b}}),v_{b}\right>+\lambda\delta_{b}^{p}\mid\norm{v_{b}}\leq\delta_{b}\right\}\\
&=\inf_{v_{b},\delta_{b}}\left\{h(v_{b},\delta_{b})\mid\norm{v_{b}}\leq\delta_{b}\right\}.
\end{split}
\end{equation}
First, note that for any $v_{b}$, $j(v_{b})\leq h(v_{b},\delta_{b})$ for all feasible $\delta_{b}$ in \eqref{eqn:inner_prob_aux}. This implies $j(v_{b})\leq\inf_{\delta_{b}}\left\{h(v_{b},\delta_{b}\mid\norm{v_{b}}\leq\delta_{b}) \right\}$, and by minimizing both sides w.r.t $v_{b}$ we see that \[\inf_{v_{b}} j(v_{b})\leq\inf_{v_{b},\delta_{b}}\left\{h(v_{b},\delta_{b})\mid\norm{v_{b}}\leq\delta_{b}\right\}.\] Let $\delta_{b}^{*}=\norm{v_{b}^{*}}$, where $v_{b}^{*}$ is the optimal solution to \eqref{eqn:inner_prob}, and note that $(v_{b}^{*},\delta_{b}^{*})$ is a feasible solution to \eqref{eqn:inner_prob_aux} such that $h(v_{b}^{*},\delta_{b}^{*})=j(v_{b}^{*})$. Therefore, we can derive the optimal solution of \eqref{eqn:inner_prob} by finding the optimal solution of \eqref{eqn:inner_prob_aux}:
\begin{align}\label{eqn:delta_opt}
    &\inf\limits_{v_{b},\delta_{b}}\left\{\left<\nabla_{\hat{f}_{\theta_{b}}}\mathcal{J}_{t}(\hat{f}_{\theta_{b}}),v_{b}\right>+\lambda\delta_{b}^{p}\mid\norm{v_{b}}\leq\delta_{b}\right\}\notag\\
    &= \inf\limits_{\delta_{b}\geq0}\left\{\inf\limits_{v_{b}}\left\{\left<\nabla_{\hat{f}_{\theta_{b}}}\mathcal{J}_{t}(\hat{f}_{\theta_{b}}),v_{b}\right>+\lambda\delta_{b}^{p}\mid\norm{v_{b}}\leq\delta_{b}\right\}\right\}\notag\\
    &= \inf\limits_{\delta_{b}\geq0}\left\{\lambda\delta_{b}^{p}+\inf\limits_{\norm{v_{b}}\leq\delta_{b}}\left\{\left<\nabla_{\hat{f}_{\theta_{b}}}\mathcal{J}_{t}(\hat{f}_{\theta_{b}}),v_{b}\right>\right\}\right\}\notag\\
    % &= \inf\limits_{\delta_{b}\geq0}\left\{\lambda\delta_{b}^{p}-\sup\limits_{\norm{v_{b}}\leq\delta_{b}}\left\{\left<\nabla_{\hat{f}_{\theta_{b}}}\mathcal{J}_{t}(\hat{f}_{\theta_{b}}),v_{b}\right>\right\}\right\}\\
    &= \inf\limits_{\delta_{b}\geq0}\left\{\lambda\delta_{b}^{p}-\delta_{b}\norm{\nabla_{\hat{f}_{\theta_{b}}}\mathcal{J}_{t}(\hat{f}_{\theta_{b}})}_{*}\right\},
\end{align}
where the last step comes from the definition of the dual norm $\norm{\cdot}_{*}$. The minimization over $\delta_{b}$ in \eqref{eqn:delta_opt} is dependent on $\lambda$ and the value $\delta_{b}^{*}(\lambda)$ that achieves the minimum can be obtained by exploiting the convexity of \eqref{eqn:delta_opt} in $\delta_{b}$. To obtain $\delta_{b}^{*}(\lambda)$
 we set the derivative of $\lambda\delta_{b}^{p}-\delta_{b}\norm{\nabla_{\hat{f}_{\theta_{b}}}\mathcal{J}_{t}(\hat{f}_{\theta_{b}})}_{*}$ with respect to $\delta_{b}$ equal to zero, which leads to
% \begin{equation}\label{eqn:delta_div}
% p\lambda\delta_{b}^{p-1}-\norm{\nabla_{\hat{f}_{\theta_{b}}}\mathcal{J}_{t}(\hat{f}_{\theta_{b}})}_{*}
% \end{equation}
% Solving \eqref{eqn:delta_div} for $\delta_{b}$ results in $\delta^{*}_{b}(\lambda)$ as,
\begin{equation*}\label{eqn:delta_star}
\delta_{b}^{*}(\lambda)=\left(\frac{\norm{\nabla_{\hat{f}_{\theta_{b}}}\mathcal{J}_{t}(\hat{f}_{\theta_{b}})}_{*}}{p\lambda}\right)^{\frac{1}{p-1}}.
\end{equation*}

By plugging $\delta_{b}^{*}(\lambda)$ in \eqref{eqn:delta_opt} we obtain a solution to the inner minimization in~\eqref{eqn:dual_problem_recast}:

\begin{equation}\label{eqn:inner_sol}
    \begin{split}
&\inf\limits_{v_{b}}\left[\left<\nabla_{\hat{f}_{\theta_{b}}}\mathcal{J}_{t}(\hat{f}_{\theta_{b}}),v_{b}\right>+\lambda\norm{v_{b}}^{p}\right]\\ 
&= \lambda(\delta_{b}^{*}(\lambda))^{p}-\delta_{b}^{*}(\lambda)\norm{\nabla_{\hat{f}_{\theta_{b}}}\mathcal{J}_{t}(\hat{f}_{\theta_{b}})}_{*}\\
& = \frac{(1-p)\norm{\nabla_{\hat{f}_{\theta_{b}}}\mathcal{J}_{t}(\hat{f}_{\theta_{b}})}_{*}^{\frac{p}{p-1}}}{p^{\frac{p}{p-1}}\lambda^{\frac{1}{p-1}}}. 
    \end{split}
\end{equation}


We can now solve the outer maximization in \eqref{eqn:dual_problem_recast} by using  \eqref{eqn:inner_sol} which leads to
\begin{align}\label{eqn:outer_prob}
    &\sup_{\lambda \geq 0}\biggr\{\frac{1}{B}\sum\limits_{b=1}^B\inf_{v_{b}}\left[\left<\nabla_{\hat{f}_{\theta_{b}}}\mathcal{J}_{t}(\hat{f}_{\theta_{b}}),v_{b}\right> + \lambda(\norm{v_{b}}^{p})\right] - \lambda \epsilon^p\biggl\}\notag\\
    &=\sup_{\lambda \geq 0}\left\{\frac{(1-p)\left(\frac{1}{B}\sum_{b=1}^{B}\norm{\nabla_{\hat{f}_{\theta_{b}}}\mathcal{J}_{t}(\hat{f}_{\theta_{b}})}_{*}^{\frac{p}{p-1}}\right)}{p^{\frac{p}{p-1}}\lambda^{\frac{1}{p-1}}} - \lambda \epsilon^p\right\}.
\end{align}
Note that \eqref{eqn:outer_prob} is a maximization of a concave function in $\lambda$: the solution can be obtained by setting the derivative of~\eqref{eqn:outer_prob} w.r.t. $\lambda$ equal to zero:
\begin{align}\label{eqn:grad_lambda}
        &\frac{\mathrm{d}}{\mathrm{d}\lambda}\left(\frac{(1-p)\left(\frac{1}{B}\sum_{b=1}^{B}\norm{\nabla_{\hat{f}_{\theta_{b}}}\mathcal{J}_{t}(\hat{f}_{\theta_{b}})}_{*}^{\frac{p}{p-1}}\right)}{p^{\frac{p}{p-1}}\lambda^{\frac{1}{p-1}}} - \lambda \epsilon^p\right)\notag\\
        &=\frac{\left(\frac{1}{B}\sum_{b=1}^{B}\norm{\nabla_{\hat{f}_{\theta_{b}}}\mathcal{J}_{t}(\hat{f}_{\theta_{b}})}_{*}^{\frac{p}{p-1}}\right)}{p^{\frac{p}{p-1}}\lambda^{\frac{p}{p-1}}} -  \epsilon^{p} .
    \end{align}
Setting \eqref{eqn:grad_lambda} equal to zero and solving for $\lambda$ gives,
\begin{equation*}\label{eqn:lambda_star}
    \lambda^{*}=\frac{\left(\frac{1}{B}\sum_{b=1}^{B}\norm{\nabla_{\hat{f}_{\theta_{b}}}\mathcal{J}_{t}(\hat{f}_{\theta_{b}})}_{*}^{\frac{p}{p-1}}\right)^{\frac{p-1}{p}}}{p\epsilon^{p-1}}.
\end{equation*}
By plugging $\lambda^{*}$ into \eqref{eqn:outer_prob} we have
\begin{equation}\label{eqn:outer_sol}
    \begin{split}
        % &\sup_{\lambda \geq 0}\left\{\frac{(1-p)\left(\frac{1}{B}\sum_{b=1}^{B}\norm{\nabla_{\hat{f}_{\theta_{b}}}\mathcal{J}_{t}(\hat{f}_{\theta_{b}})}_{*}^{\frac{p}{p-1}}\right)}{p^{\frac{p}{p-1}}\lambda^{\frac{1}{p-1}}} - \lambda \epsilon^p\right\}\\
        &\frac{(1-p)\left(\frac{1}{B}\sum_{b=1}^{B}\norm{\nabla_{\hat{f}_{\theta_{b}}}\mathcal{J}_{t}(\hat{f}_{\theta_{b}})}_{*}^{\frac{p}{p-1}}\right)}{p^{\frac{p}{p-1}}\lambda^{*\frac{1}{p-1}}} - \lambda^{*} \epsilon^p\\
        &=-\epsilon\left(\sum_{b=1}^{B}\norm{\nabla_{\hat{f}_{\theta_{b}}}\mathcal{J}_{t}(\hat{f}_{\theta_{b}})}_{*}^{\frac{p}{p-1}}\right)^{\frac{p-1}{p}}.
    \end{split}
\end{equation}
Finally, combining \eqref{eqn:outer_sol} with \eqref{eqn:dual_problem_recast} we arrive at \eqref{eqn:thm_1_statement} as the closed form solution.\end{proof}
% Using Jimmy's results given in lemma \ref{lem:Jimmy}, with

% \begin{equation*}
%     \begin{cases}
%         p=2\\
%         \mathbf{z} = \tilde{f}\\
%         \hat{\mathbf{z}}_i = \hat{f}_{\theta_i}\\
%         l(\mathbf{z}) = \mathbb{E}_{\tilde{f}}\biggr[\sum \limits_{i=t}^{t+T-1}\gamma^{i-t}r(\hat{\bv{x}}_i,\bv{u}_i)\biggl]\\
%         g_i = \nabla_{\tilde{f}}\mathbb{E}_{\tilde{f}}\bigr(\sum \limits_{i=t}^{t+T-1}\gamma^{i-t}r(\hat{\bv{x}}_i,\bv{u}_i)\bigl)\biggr|_{\tilde{f}=\hat{f_b}}\\
%         n = B
%     \end{cases},
% \end{equation*}

% and relying on Wasserstein duality, we can derive a regularized version of the RHS of \eqref{eqn:DR_start}, namely:

% \begin{equation}\label{eqn:Jimmy}
%     \mathcal{J}_{\Omega_{\epsilon}(\hat{\mu})}(\bv{x}_t, \bv{u}_{t:t+T-1}) \approx \frac{1}{B}\sum\limits_{b=1}^B \mathbb{E}_{\hat{f}_{\theta_b} }\biggr(\sum \limits_{i=t}^{t+T-1}\gamma^{i-t}r(\hat{\bv{x}}_i,\bv{u}_i)\biggl) -\epsilon \biggr(\frac{1}{B}\sum\limits_{b=1}^B \norm{\nabla_{\tilde{f}}\mathbb{E}_{\tilde{f}}\bigr(\sum \limits_{i=t}^{t+T-1}\gamma^{i-t}r(\hat{\bv{x}}_i,\bv{u}_i)\bigl)\biggr|_{\tilde{f}=\hat{f}_{\theta_b}}}^2\biggl)^{\frac{1}{2}}.
% \end{equation}
As in PETS algorithm, we compute the expectation over the estimated state trajectory in \eqref{eqn:Chua_obj} using a particle-based propagation, i.e., propagating $Q \in \mathbb{Z}_{+}$ particles from initial state $\hat{\bv{x}}^\text{q}_{t}=\bv{x}_{t}$ at time $t$, where $\text{q} = 1,2,\dots,Q$. Each particle is propagated using $\hat{\bv{x}}^{\text{q}}_{t+1}\sim\hat{f}_{\theta_{b}}$. Furthermore, as in PETS, we use cross entropy method (CEM) \cite{botev2013cross} to select actions, which iteratively samples $M \in \mathbb{Z}_{+}$ action sequences $\{\bv{u}_{t:t+T-1}^i\}_{i=1,2,\dots,M}$ from a candidate distribution $\mathbf{CEM}(\cdot):=\sN(\varphi_{\bv{u}},\Sigma_{\bv{u}})$. Actions sampled from the candidate distribution are evaluated using  \eqref{eqn:thm_1_statement}, and parameters $\varphi_{\bv{u}}$ and  $\Sigma_{\bv{u}}$ are updated iteratively based on best action sequence among $M$ samples.


 By exploiting the B-NNs ensemble and letting $p=2$ we develop a simplified version of \eqref{eqn:thm_1_statement} that we use in our DR-PETS algorithm that is reported in the following result.
  

\begin{thm}\label{thm:main}
    Given the problem in \eqref{eqn:DR_obj}, the optimization objective can be reformulated as,  
    \begin{equation}\label{eqn:thm_statement}
    \begin{split}
       % &\mathcal{J}_{\Omega_{\epsilon}(\hat{\mu})}(\bv{x}_t, \bv{u}_{t:t+T-1}) \approx 
       &\max_{\bv{u}_{t:t+T-1}}\biggr\{ \frac{1}{B}\sum\limits_{b=1}^B \mathbb{E}_{\hat{f}_{\theta_b} }\biggr[\sum \limits_{i=t}^{t+T-1}\gamma^{i-t}r(\hat{\bv{x}}_i,\bv{u}_i)\biggl]\\ & -\epsilon \biggr(\frac{1}{B}\sum\limits_{b=1}^B \bigg\|
    \sum\limits_{k=1}^{T-1} \gamma^{k} \mathbb{E}_{\hat{\bv{x}}_{t+1:t+k}\sim\hat{f}_{\theta_{b}}}\\
    &\biggr[ \nabla_{\hat{f}_{\theta_{b}}} \ln(\hat{f}_{\theta_{b}}(\hat{\bv{x}}_{t+k}|\hat{\bv{x}}_{t+k-1}, \bv{u}_{t+k-1}))\\
    &\sum\limits_{i=t+k}^{t+T-1} \mathbb{E}_{\hat{\bv{x}}_{i+1}\sim\hat{f}_{\theta_{b}}}\bigr[\gamma^{i-t-k} r(\hat{\bv{x}}_{i}, \bv{u}_{i})\bigl] \biggl]\bigg\|^2\biggl)^{\frac{1}{2}}\biggl\}\\
     & s.t. \quad {\bv{u}_t, \bv{u}_{t+1}, \dots, \bv{u}_{t+T-1}} \in \mathcal{U}.
    \end{split}
    \end{equation}
\end{thm}

\begin{proof}
Considering Theorem \ref{thm:Jimmy}, when $p=2$, the l.h.s. of \eqref{eqn:thm_1_statement} can be rewritten as,
% \begin{equation*}
%     \begin{cases}
%         p=2\\
%         \mathbf{z} = \tilde{f}\\
%         \hat{\mathbf{z}}_i = \hat{f}_{\theta_i}\\
%         l(\mathbf{z}) = \mathbb{E}_{\tilde{f}}\biggr[\sum \limits_{i=t}^{t+T-1}\gamma^{i-t}r(\hat{\bv{x}}_i,\bv{u}_i)\biggl]\\
%         g_i = \nabla_{\tilde{f}}\mathbb{E}_{\tilde{f}}\bigr(\sum \limits_{i=t}^{t+T-1}\gamma^{i-t}r(\hat{\bv{x}}_i,\bv{u}_i)\bigl)\biggr|_{\tilde{f}=\hat{f_b}}\\
%         n = B,
%     \end{cases}
% \end{equation*}
\begin{equation}\label{eqn:Jimmy}
\begin{split}
    % \mathcal{J}_{\Omega_{\epsilon}(\hat{\mu})}(\bv{x}_t, \bv{u}_{t:t+T-1}) \approx
    &\frac{1}{B}\sum\limits_{b=1}^B \mathcal{J}_{t}(\hat{f}_{\theta_{b}})-\epsilon \left(\frac{1}{B}\sum\limits_{b=1}^B \norm{\nabla_{\hat{f}_{\theta_{b}}}\mathcal{J}_{t}(\hat{f}_{\theta_{b}})}^2\right)^{\frac{1}{2}}.
\end{split}
\end{equation}

Now, considering the horizon $T$, during the planning we sample $M$ possibly different sequences of actions, i.e., $\{\bv{u}_{t:t+T-1}^k\}_{k=1,2,\dots,M}$. Thus, the expectation in the first part of \eqref{eqn:Jimmy} is only w.r.t. $\hat{f}_{\theta_{b}}$, and it can be approximated by propagating $Q$ particles for each $k$. Next, we simplify the gradient term $\nabla_{\hat{f}_{\theta_{b}}}\mathcal{J}_{t}(\hat{f}_{\theta_{b}})$ in \eqref{eqn:Jimmy} by unrolling it over time horizon $T$. We use $\hat{\bv{x}}$ to emphasize that the state is predicted using the learned model. For brevity consider $\mathcal{J}_{t}=\mathcal{J}_{t}(\hat{f}_{\theta_{b}})$.
\begin{equation}\label{eqn:gradientVF0}
\begin{split}
    &\nabla_{\hat{f}_{\theta_{b}}}\mathcal{J}_{t}
    = \nabla_{\hat{f}_{\theta_{b}}}\biggr[\sum \limits_{i=t}^{t+T-1}\mathbb{E}_{\hat{f}_{\theta_{b}}}\bigr[\gamma^{i-t}r(\hat{\bv{x}}_i,\bv{u}_i)\bigl]\biggl]\\
    % &=\nabla_{\tilde{f}}\biggr[  r({x}_t,\bv{u}_t) + \gamma\int \tilde{f}({x}_{t+1}|{x}_t,\bv{u}_t) \sum \limits_{i=t+1}^{t+T-1}\mathbb{E}_{\tau\sim (\pi,\tilde{f})}\bigr[\gamma^{i-t-1}r({x}_i,\bv{u}_i)\bigl]d{{x}_{t+1}}\biggl],\\
    % &=\nabla_{\tilde{f}}\biggr[  r({x}_t,\bv{u}_t) + \gamma\int \tilde{f}({x}_{t+1}|{x}_t,\bv{u}_t) \mathcal{J}_{\tilde{f}}({x}_{t+1}, \bv{u}_{t+1:t+T-1})d{{x}_{t+1}}\biggl],\\
    &= \gamma\int\nabla_{\hat{f}_{\theta_{b}}}\biggr( \hat{f}_{\theta_{b}}(\hat{\bv{x}}_{t+1}|\hat{\bv{x}}_t,\bv{u}_t) \mathcal{J}_{t+1}\biggl)d{\hat{\bv{x}}_{t+1}}\\
    &= \phi(\hat{\bv{x}}_t,\bv{u}_t) + \gamma\int \hat{f}_{\theta_{b}}(\hat{\bv{x}}_{t+1}|\hat{\bv{x}}_t,\bv{u}_t) \nabla_{\hat{f}_{\theta_{b}}}\left(\mathcal{J}_{t}\right)d{\hat{\bv{x}}_{t+1}},
\end{split}
\end{equation}
where
$%\begin{equation*}
    \phi(\hat{\bv{x}}_t,\bv{u}_t) = \gamma\int \nabla_{\hat{f}_{\theta_{b}}}\bigr( \hat{f}_{\theta_{b}}(\hat{\bv{x}}_{t+1}|\hat{\bv{x}}_t,\bv{u}_t)\bigl) \mathcal{J}_{t}d{\hat{\bv{x}}_{t+1}}.
$%\end{equation*}

Continuing the unrolling of \eqref{eqn:gradientVF0} with $\hat{f}_{\theta_{b}}(\hat{\bv{x}}_{t+i}) = \hat{f}_{\theta_{b}}(\hat{\bv{x}}_{t+i}|\hat{\bv{x}}_{t+i-1}, \bv{u}_{t+i-1})$ we have:

\begin{align}\label{eqn:gradientVF}
    &\nabla_{\hat{f}_{\theta_{b}}}\mathcal{J}_{t}
    % &= \phi({x}_t,\bv{u}_t) + \gamma\int \tilde{f}({x}_{t+1}|{x}_t,\bv{u}_t) \nabla_{\tilde{f}}\mathcal{J}_{\tilde{f}}({x}_{t+1}, \bv{u}_{t+1:t+T-1})d{{x}_{t+1}},\\
    = \phi(\hat{\bv{x}}_t,\bv{u}_t) + \gamma\int \hat{f}_{\theta_{b}}(\hat{\bv{x}}_{t+1}) \nabla_{\hat{f}_{\theta_{b}}}\mathcal{J}_{t+1}d\hat{\bv{x}}_{t+1}\notag\\
    % &= \phi({x}_t,\bv{u}_t) + \gamma\int f({x}_{t+1})\nabla_{f}\biggr[r(\bv{x}_{t+1}, \bv{u}_{t+1}) + \gamma \int f(\bv{x}_{t+2}) \mathcal{J}_{t+2}d\bv{x}_{t+2}\biggl] d\bv{x}_{t+1},\\
    &= \phi(\hat{\bv{x}}_t,\bv{u}_t) + \gamma\int \hat{f}_{\theta_{b}}(\hat{\bv{x}}_{t+1})\phi(\hat{\bv{x}}_{t+1}, \bv{u}_{t+1})d\hat{\bv{x}}_{t+1}\notag\\
    &+ \gamma^2 \int \hat{f}_{\theta_{b}}(\hat{\bv{x}}_{t+1})\int \hat{f}_{\theta_{b}}(\hat{\bv{x}}_{t+2}) \phi(\hat{\bv{x}}_{t+2},\bv{u}_{t+2})d\hat{\bv{x}}_{t+2}d\hat{\bv{x}}_{t+1} +\notag\\
    & \dots + \gamma^{T-1} \int \hat{f}_{\theta_{b}}(\hat{\bv{x}}_{t+1})\dots\notag\\
    &\int \hat{f}_{\theta_{b}}(\hat{\bv{x}}_{t+T-1}) \nabla_{\hat{f}_{\theta_{b}}}\mathcal{J}_{t+T-1}d\hat{\bv{x}}_{t+T-1:t+1}.
\end{align}

Note that,
$\nabla_{\hat{f}_{\theta_{b}}}\mathcal{J}_{t+T-1}=\nabla_{\hat{f}_{\theta_{b}}}r(\hat{\bv{x}}_{t+T-1},\bv{u}_{t+T-1})=0.$ Then, by substituting $\phi$ in \eqref{eqn:gradientVF} we obtain:
\begin{equation}\label{eqn:gradientVF2}
\begin{split}
&\nabla_{\hat{f}_{\theta_{b}}}\mathcal{J}_{t} = \gamma \int \nabla_{\hat{f}_{\theta_{b}}} \hat{f}_{\theta_{b}}(\hat{\bv{x}}_{t+1})\mathcal{J}_{t+1}d\hat{\bv{x}}_{t+1}\\
    & + \gamma^2 \int \hat{f}_{\theta_{b}}(\hat{\bv{x}}_{t+1})\int \nabla_{\hat{f}_{\theta_{b}}} (\hat{f}_{\theta_{b}}(\hat{\bv{x}}_{t+2}))\mathcal{J}_{t+2}d\hat{\bv{x}}_{t+2:t+1}\\
    & +\dots + \gamma^{T-1} \int \hat{f}_{\theta_{b}}(\hat{\bv{x}}_{t+1})\dots\\
    &\int \nabla_{\hat{f}_{\theta_{b}}}(\hat{f}_{\theta_{b}}(\hat{\bv{x}}_{t+T-1}))\mathcal{J}_{t+T-1}d\hat{\bv{x}}_{t+T-1:t+1}.
\end{split}
\end{equation}

Multiplying each term in \eqref{eqn:gradientVF2} by $\frac{\hat{f}_{\theta_{b}}(\hat{\bv{x}}_{t+i})}{\hat{f}_{\theta_{b}}(\hat{\bv{x}}_{t+i})}$, and observing that $\int \hat{f}_{\theta_{b}}(\hat{\bv{x}}_{t+1})\nabla_{\hat{f}_{\theta_{b}}} \ln(\hat{f}_{\theta_{b}}(\hat{\bv{x}}_{t+1}))\mathcal{J}_{t+1}d\hat{\bv{x}}_{t+1} = \mathbb{E}_{\hat{\bv{x}}_{t+1}}\biggr[ \nabla_{\hat{f}_{\theta_{b}}} \ln(\hat{f}_{\theta_{b}}(\hat{\bv{x}}_{t+1}))\mathcal{J}_{t+1}\biggl]$, we can write:

\begin{align}\label{eqn:gradientVF3}
    &\nabla_{\hat{f}_{\theta_{b}}}\mathcal{J}_{t}= \gamma \mathbb{E}_{\hat{\bv{x}}_{t+1}}\biggr[ \nabla_{\hat{f}_{\theta_{b}}} \ln(\hat{f}_{\theta_{b}})\mathcal{J}_{t+1} \biggl]\notag\\ 
    &+ \gamma^2 \mathbb{E}_{\hat{\bv{x}}_{t+1:t+2}\sim\hat{f}_{\theta_{b}}}\biggr[ \nabla_{\hat{f}_{\theta_{b}}} \ln(\hat{f}_{\theta_{b}}(\hat{\bv{x}}_{t+2}))\mathcal{J}_{t+2}\biggl]+\dots\notag\\
    &+ \gamma^{T-1} \mathbb{E}_{\hat{\bv{x}}_{t+1:t+T-1}\sim\hat{f}_{\theta_{b}}}\biggr[ \nabla_{\hat{f}_{\theta_{b}}} \ln(\hat{f}_{\theta_{b}}(\hat{\bv{x}}_{t+T-1}))\mathcal{J}_{t+T-1} \biggl]\notag\\
    &= \sum\limits_{k=1}^{T-1} \gamma^{k} \mathbb{E}_{\hat{\bv{x}}_{t+1:t+k}\sim\hat{f}_{\theta_{b}}}\biggr[ \nabla_{\hat{f}_{\theta_{b}}} \ln(\hat{f}_{\theta_{b}}(\hat{\bv{x}}_{t+k}))\mathcal{J}_{t+k} \biggl]\notag\\
    % &= \sum\limits_{k=1}^{T-1} \gamma^{k} \mathbb{E}_{\hat{\bv{x}}_{t+1:t+k}\sim\hat{f}_{\theta_{b}}}\biggr[ \nabla_{\hat{f}_{\theta_{b}}} \ln(\hat{f}_{\theta_{b}}(\hat{\bv{x}}_{t+k}|\hat{\bv{x}}_{t+k-1}, \bv{u}_{t+k-1}))\\
    &\sum\limits_{i=t+k}^{t+T-1} \mathbb{E}_{\hat{\bv{x}}_{i+1}\sim\hat{f}_{\theta_{b}}}\bigr[\gamma^{i-t-k} r(\hat{\bv{x}}_{i}, \bv{u}_{i})\bigl] \biggl].
\end{align}


By combining \eqref{eqn:Jimmy} and \eqref{eqn:gradientVF3} we obtain \eqref{eqn:thm_statement}. \end{proof}

% \begin{algorithm}
% \caption{Distributionally robust probabilistic ensembles with trajectory sampling (DRPETS)}
% \begin{algorithmic}[1]
% \State Initialize data $\mathcal{D}$ with a random controller for one rollout.
% \For{Iteration $k = 1$ to $K$}
%     \State Train a probabilistic ensemble dynamics model $\hat{f}_{\theta}$ given data-set $\mathcal{D}$.
%     \For{Time $t = 0$ to $T$}
%         \For{Actions sampled $\bv{u}_{t:t+T} \sim \mathbf{CEM}(\cdot), 1$ to $N$ samples} 
%         \State Propagate state particles $\bv{x}_t^p$ using trajectory sampling and $\hat{f}_{\theta_{b}}\mid\{\mathcal{D}, \bv{u}_{t:t+T}\}$.
%         \State Evaluate actions as per \eqref{eqn:thm_statement}
%         \State Update $\mathbf{CEM}(\cdot)$ distribution.
%         \State Execute first action $\bv{u}_t^*$ (only) from optimal actions $\bv{u}_{t:t+T}^*$.
%         \State Record outcome: $\mathcal{D} \gets \mathcal{D} \cup \{\bv{x}_t, \bv{u}_t^*, \bv{x}_{t+1}\}$.
%         \EndFor
%     \EndFor
% \EndFor
% \end{algorithmic}
% \end{algorithm}

To recap, we aim to solve \eqref{eqn:MPC_prob} despite the unknown transition model \( f \). We address this by formulating a DR version of \eqref{eqn:Chua_obj}, assuming \( f \) follows a distribution within the ambiguity set \( \Omega_{\epsilon}(\hat{\mu}) \). Given the WR-MDP, we define the DR optimization problem in \eqref{eqn:DR_obj}.  
Applying \( p \)-Wasserstein duality, we derive \eqref{eqn:thm_1_statement} as a tractable solution to Problem \ref{prob1}, comprising an empirical estimate over the ensemble performance index and a regularizer. Leveraging the B-NNs ensemble, we further simplify Theorem \ref{thm:Jimmy}, yielding \eqref{eqn:thm_statement}.

% To recap, we want to solve problem \eqref{eqn:MPC_prob} given that the transition model $f$ is unknown. To this end, we consider a DR version of the objective in \eqref{eqn:Chua_obj}, assuming that $f$  is a random variable that belongs to some distribution in the ambiguity set $\Omega_{\epsilon}(\hat{\mu})$. Given the WRMDP, we formulate DR optimization problem presented in \eqref{eqn:DR_obj}.
% Using $p$-Wasserstein duality we derive \eqref{eqn:thm_1_statement} as a tractable solution of Problem \ref{prob1}. This is composed of two terms: an empirical estimate over the ensemble of the performance index, and a regularizer term. Then, we leverage the B-NNs ensemble to further simplify the result of Theorem \ref{thm:Jimmy}, resulting in a regularized objective function as reported in \eqref{eqn:thm_statement}. 


% Considering $B$ empirical models sampled from $\mu$, i.e., $\hat{f}_{\theta_1}, \hat{f}_{\theta_2}, \dots, \hat{f}_{\theta_B} \sim \mu$, and defining the corresponding empirical distribution $\hat{\mu}$, we can derive a lower bound of \eqref{eqn:plan3}, namely \eqref{eqn:DR_start}, that holds with some probability depending on the radius of a Wasserstein ball centered in $\hat{\mu}$. 

% \grtodo{WRMDP is undefined now. I would avoid summaring things in a letter...that's a very short paper. Also, here in this section there is no mention to DR-PETS. If this is the case, then we should not write our story around it. This means changing the title, to start (still, claiming that our novelty is a framework is dangerous and can raise criticisms - maybe just say that we discover a tractable approximation). Btw, since our method relies on an approximation, I think we propose an approximation not  a reformulation.}
% \begin{equation}\label{eqn:gradientVF3}
% \begin{split}
%     \nabla_{\tilde{f}}\mathcal{J}_{\tilde{f}}(\bv{x}_t, \bv{u}_{t:t+T-1})
%     &= \gamma \mathbb{E}_{\bv{x}_{t+1}}\biggr[ \nabla_f \ln(f(\bv{x}_{t+1}))\mathcal{J}_{t+1} \biggl] + \gamma^2 \mathbb{E}_{\{\bv{x}_{t+1:t+2}\}}\biggr[ \nabla_f \ln(f(\bv{x}_{t+2}))\mathcal{J}_{t+2} \biggl] + \dots\\
%     &\qquad \qquad + \gamma^{T-1} \mathbb{E}_{\{\bv{x}_{t+1:t+T-1}\}}\biggr[ \nabla_f \ln(f(\bv{x}_{t+T-1}))\mathcal{J}_{t+T-1} \biggl],\\
%     &= \sum\limits_{k=1}^{T-1} \gamma^{k} \mathbb{E}_{\{\bv{x}_{t+1:t+k}\}}\biggr[ \nabla_f \ln(f(\bv{x}_{t+k}))\mathcal{J}_{t+k} \biggl],\\
%     &= \sum\limits_{k=1}^{T-1} \gamma^{k} \mathbb{E}_{\{\bv{x}_{t+1:t+k}\}}\biggr[ \nabla_{\tilde{f}} \ln(\tilde{f}(\bv{x}_{t+k}|\bv{x}_{t+k-1}, \bv{u}_{t+k-1}))\sum\limits_{i=t+k}^{t+T-1} \mathbb{E}_{x_{i+1}}\bigr[\gamma^{i-t-k} r(x_{i}, \bv{u}_{i})\bigl] \biggl].
% \end{split}
% \end{equation}

% From a practical point of view the implementation of \eqref{eqn:thm_statement} requires some further considerations. Particularly, one should: i) compute the expectation over the empirical models, and ii) define the way to choose the actions. Given an initial state $x_0$ from the true model, the expected value of the next state can be estimated replicating $x_0$ over $P$ particles, and averaging all the different generated trajectories. Moreover, in the above discussion we did not include any specific policy, but we just aimed to maximize a sequence of actions. Thus, a way to choose actions is still pending. From a practical point of view, a way to do this is to assume that the actions are selected according to some probability distribution, e.g., $u \sim \mathcal{N}(m_u, \Sigma_u)$. Note that the distribution of the actions does not necessarily depend on the current state of the system, thus it cannot be referred to a policy $\pi(x_0)$. Instead, the goodness of such actions strongly rely on the MPC method, that plans over an horizon of length $T$ but then applies only the first action. At the next iteration, it updates the parameters $m_u$ and $\Sigma_u$ of the distribution based on its previous experience. 
% Because the action selection does not depend on state and  could be very general (e.g. a uniform distribution), in the previous formulation we omitted the explicit dependence of all expectations on the sampled action.

% The complete DR-PETS algorithm is proposed in Algorithm \ref{alg:STQL}.
% \begin{algorithm}\footnotesize
% % \algsetup{linenosize=\small}
% % \scriptsize
% \caption{DR-PETS algorithm}
% \begin{algorithmic}[1]
% \Require {$N_{train}$, $T$, $N_{init}$, $n\_elite$, $pop\-size$, $P$, $ T_{ep}$, $\epsilon$}
% \Ensure{$\hat{f}_{\theta_b}$, $b=1,\dots,B$}
% \State $\mathcal{D} \varleftarrow \emptyset$
% \For{$t = 0, 1, \dots ,N_{init}$}
% \State Read the state of the system $\bv{x}_t$, execute a random action $\bv{u}_t$, read the new state $\bv{x}_{t+1}$
% \State $\mathcal{D} \varleftarrow \{\mathcal{D} \cup (\bv{x}_t,\bv{u}_t, \bv{x}_{t+1})\}$
% \EndFor
% \State Train $B$ models $\hat{f}_{\theta_b}$ using $\mathcal{D}$
% \State $m_{\pi} \varleftarrow \textit{mean}(\{\bv{u}_i\}_{i=0,\dots, N_{init}})$, $\Sigma_{\pi} \varleftarrow \textit{std}(\{\bv{u}_i\}_{i=0,\dots,N_{init}})$, initialize $\pi \sim \mathcal{N}(m_{\pi}, \Sigma_{\pi})$
% \For{$k = 0, 1, \dots ,N_{train}$}
% \For{$t = 0, 1, \dots , T_{ep}-1$}
% \State Read the state of the system $\bv{x}_t$
% \State $\mathcal{R}, \mathcal{R}_{next} \varleftarrow \{0,0,\dots,0\}$ (initialize two arrays of zeros with shape $(pop\_size \times P, B)$)
% \State Initialize gradients: $\mathcal{L} \varleftarrow \{0,0,\dots,0\}$, (shape: $(pop\_size \times P, B)$)
% \State $\mathcal{X} \varleftarrow \{\bv{x}_t, \bv{x}_t, \dots, \bv{x}_t\}$, (Repeat $\bv{x}_t$ in an array with shape $(pop\_size \times P, B)$)
% \State Sample from $\pi$ two batches of actions $\mathcal{U}$ and $\mathcal{U}_{next}$ of dimension $(pop\_size, T)$
% \State Repeat the sequences of actions $\mathcal{U}$ and $\mathcal{U}_{next}$ in the shapes $(pop\_size \times P, T, B)$
% \For{$i = 0, 1, \dots , T-1$}
% \State $\mathcal{X}_{next},\mathcal{X}_{next,next} \varleftarrow \{0,0,\dots,0\}$ (shape: $(pop\_size \times P,B)$)
% \For{$b = 1, 2, \dots , B$}
% \For{$j = 1, 2, \dots , (pop\_size \times P)$}
% \State Apply the action $\mathcal{U}[j,i,b]$ to $\mathcal{X}[j,b]$
% \State Read the new predicted state $\hat{\bv{x}}_{i+1} \sim \hat{f}_{\theta_b}(\hat{\bv{x}}_i, \mathcal{U}[j,i,b])$ and the reward $r_{i+1}$
% \State $\mathcal{X}_{next}[j,b] \varleftarrow \hat{\bv{x}}_{i+1}$, $\mathcal{R}[j,b] \varleftarrow \mathcal{R}[j,b] + r_{i+1}$
% \If{i == 0}
% \State Compute gradients $\mathcal{L}[j,b] \varleftarrow \nabla_{\hat{f}_{\theta_b}} ln \bigr( \hat{f}_{\theta_b}(\hat{\bv{x}}_{i+1}|\mathcal{X}[j,b],  \mathcal{U}[j,i,b])$
% \EndIf
% \State Apply the action $\mathcal{U}_{next}[j,i,b]$ to $\mathcal{X}_{next}[j,b]$
% \State Read $\hat{\bv{x}}_{i+2} \sim \hat{f}_{\theta_b}(\mathcal{X}_{next}[j,b], \mathcal{U}_{next}[j,i,b])$, and $r_{next_{i+1}}$
% \State $\mathcal{X}_{next,next}[j,b] \varleftarrow \hat{\bv{x}}_{i+2}$, $\mathcal{R}_{next}[j,b] \varleftarrow \mathcal{R}_{next}[j,b] + r_{next_{i+1}}$
% \EndFor
% \EndFor
% \State $\mathcal{X} \varleftarrow \mathcal{X}_{next}$
% \EndFor
% \State Average $\mathcal{R}$, $\mathcal{R}_{next}$, and $\mathcal{L}$ over $P$ (shape: $(pop\_size, B)$)
% \State Average $\mathcal{U}$ over $P$ and $B$ (resulting shape: $(pop\_size, T)$)
% \State $\mathcal{J}_{\Omega_{\epsilon}(\hat{\mu})}(\bv{x}_t, \bv{u}_{t:t+T-1}) \varleftarrow \frac{1}{B}\sum\limits_{b=1}^B \mathcal{R}[:,b] -\epsilon \biggr(\frac{1}{B}\sum\limits_{b=1}^B \norm{\mathcal{L}[:,b]\mathcal{R}_{next}[:,b]}^2\biggl)^{\frac{1}{2}}$
% \State Order $\mathcal{J}_{\Omega_{\epsilon}(\hat{\mu})}(\bv{x}_t, \bv{u}_{t:t+T-1})$ in descending order and accordingly also $\mathcal{U}$
% \State Store the first $n\_elite$ rows of $\mathcal{U}$ in $\mathcal{U}_{elite}$, (new shape $(n\_elite, T)$)
% \State $m_{\pi} \varleftarrow mean(\mathcal{U}_{elite})$, $\Sigma_{\pi} \varleftarrow std(\mathcal{U}_{elite})$
% \State $\bv{u}_t^{\ast} \varleftarrow mean(\mathcal{U}_{elite}[0,0])$
% \State Execute $\bv{u}_t^{\ast}$ on the real system, observe $\bv{x}_{t+1} \sim f(\bv{x}_t, \bv{u}_t^{\ast})$
% \State $\mathcal{D} \varleftarrow \{\mathcal{D} \cup (\bv{x}_t,\bv{u}_t, \bv{x}_{t+1})\}$
% \EndFor
% \State Train $\hat{f}_{{\theta_b}},\;\forall b$ using $\mathcal{D}$
% \EndFor
% \end{algorithmic}\label{alg:STQL}
% \end{algorithm}

% \begin{algorithm}\footnotesize
% % \algsetup{linenosize=\small}
% % \scriptsize
% \caption{PETS algorithm}
% \begin{algorithmic}[1]
% \Require {$N_{train}$, $T$, $N_{init}$, $n\_elite$, $pop\-size$, $P$, $ T_{ep}$}
% \Ensure{$\hat{f}_{\theta_b}$, $b=1,\dots,B$}
% \State $\mathcal{D} \varleftarrow \emptyset$
% \For{$t = 0, 1, \dots ,N_{init}$}
% \State Read the state of the system $\bv{x}_t$, execute a random action $\bv{u}_t$, read the new state $\bv{x}_{t+1}$
% \State $\mathcal{D} \varleftarrow \{\mathcal{D} \cup (\bv{x}_t,\bv{u}_t, \bv{x}_{t+1})\}$
% \EndFor
% \State Train $B$ models $\hat{f}_{\theta_b}$ using $\mathcal{D}$
% \State $m_{\pi} \varleftarrow \textit{mean}(\{\bv{u}_i\}_{i=0,\dots, N_{init}})$, $\Sigma_{\pi} \varleftarrow \textit{std}(\{\bv{u}_i\}_{i=0,\dots,N_{init}})$, initialize $\pi \sim \mathcal{N}(m_{\pi}, \Sigma_{\pi})$
% \For{$k = 0, 1, \dots ,N_{train}$}
% \For{$t = 0, 1, \dots , T_{ep}-1$}
% \State Read the state of the system $\bv{x}_t$
% \State $\mathcal{R} \varleftarrow \{0,0,\dots,0\}$ (initialize array of zeros with shape $(pop\_size \times P, B)$)
% % \State Initialize gradients: $\mathcal{L} \varleftarrow \{0,0,\dots,0\}$, (shape: $(pop\_size \times P, B)$)
% \State $\mathcal{X} \varleftarrow \{\bv{x}_t, \bv{x}_t, \dots, \bv{x}_t\}$, (Repeat $\bv{x}_t$ in an array with shape $(pop\_size \times P, B)$)
% \State Sample from $\pi$ two batches of actions $\mathcal{U}$ of dimension $(pop\_size, T)$
% \State Repeat the sequences of actions $\mathcal{U}$ in the shape $(pop\_size \times P, T, B)$
% \For{$i = 0, 1, \dots , T-1$}
% % \State $\mathcal{X}_{next},\mathcal{X}_{next,next} \varleftarrow \{0,0,\dots,0\}$ (shape: $(pop\_size \times P,B)$)
% % \For{$b = 1, 2, \dots , B$}
% % \For{$j = 1, 2, \dots , (pop\_size \times P)$}
% % \State Apply the action $\mathcal{U}[j,i,b]$ to $\mathcal{X}[j,b]$
% % \State Read the new predicted state $\hat{\bv{x}}_{i+1} \sim \hat{f}_{\theta_b}(\hat{\bv{x}}_i, \mathcal{U}[j,i,b])$ and the reward $r_{i+1}$
% % \State $\mathcal{X}_{next}[j,b] \varleftarrow \hat{\bv{x}}_{i+1}$, $\mathcal{R}[j,b] \varleftarrow \mathcal{R}[j,b] + r_{i+1}$

% % \If{i == 0}
% \State Apply actions $\mathcal{U}$ to $\mathcal{X}$
% \State Propagate $P$ state $\bv{x}_{t}$ particles through ensemble of $B$ models $\hat{f}_{\theta_b}$
% \State Compute $\mathcal{R}$
% % \State Compute gradients $\mathcal{L}$
% % \EndIf
% % \State Apply the action $\mathcal{U}_{next}[j,i,b]$ to $\mathcal{X}_{next}[j,b]$
% % \State Read $\hat{\bv{x}}_{i+2} \sim \hat{f}_{\theta_b}(\mathcal{X}_{next}[j,b], \mathcal{U}_{next}[j,i,b])$, and $r_{next_{i+1}}$
% % \State $\mathcal{X}_{next,next}[j,b] \varleftarrow \hat{\bv{x}}_{i+2}$, $\mathcal{R}_{next}[j,b] \varleftarrow \mathcal{R}_{next}[j,b] + r_{next_{i+1}}$
% % \EndFor
% % \EndFor
% % \State $\mathcal{X} \varleftarrow \mathcal{X}_{next}$
% \EndFor
% \State Average $\mathcal{R}$ and  $\mathcal{U}$ over $P$ (shape: $(pop\_size, B)$)
% \State Average $\mathcal{U}$ over $P$ and $B$ (resulting shape: $(pop\_size, T)$)
% \State $\mathcal{J}(\bv{x}_t, \bv{u}_{t:t+T-1}) \varleftarrow \frac{1}{B}\sum\limits_{b=1}^B \mathcal{R}[:,b]$
% \State Order $\mathcal{J}_{\Omega_{\epsilon}(\hat{\mu})}(\bv{x}_t, \bv{u}_{t:t+T-1})$ in descending order and accordingly also $\mathcal{U}$
% \State Store the first $n\_elite$ rows of $\mathcal{U}$ in $\mathcal{U}_{elite}$, (new shape $(n\_elite, T)$)
% \State $m_{\pi} \varleftarrow mean(\mathcal{U}_{elite})$, $\Sigma_{\pi} \varleftarrow std(\mathcal{U}_{elite})$
% \State $\bv{u}_t^{\ast} \varleftarrow mean(\mathcal{U}_{elite}[0,0])$
% \State Execute $\bv{u}_t^{\ast}$ on the real system, observe $\bv{x}_{t+1} \sim f(\bv{x}_t, \bv{u}_t^{\ast})$
% \State $\mathcal{D} \varleftarrow \{\mathcal{D} \cup (\bv{x}_t,\bv{u}_t, \bv{x}_{t+1})\}$
% \EndFor
% \State Train $\hat{f}_{{\theta_b}},\;\forall b$ using $\mathcal{D}$
% \EndFor
% \end{algorithmic}\label{alg:STQL}
% \end{algorithm}


% \begin{algorithm}\footnotesize
% % \algsetup{linenosize=\small}
% % \scriptsize
% \caption{DR-PETS algorithm}
% \begin{algorithmic}[1]
% \Require {$N_{train}$, $T$, $N_{init}$, $n\_elite$, $pop\-size$, $P$, $ T_{ep}$, $\epsilon$}
% \Ensure{$\hat{f}_{\theta_b}$, $b=1,\dots,B$}
% \State $\mathcal{D} \varleftarrow \emptyset$
% \For{$t = 0, 1, \dots ,N_{init}$}
% \State Read the state of the system $\bv{x}_t$, execute a random action $\bv{u}_t$, read the new state $\bv{x}_{t+1}$
% \State $\mathcal{D} \varleftarrow \{\mathcal{D} \cup (\bv{x}_t,\bv{u}_t, \bv{x}_{t+1})\}$
% \EndFor
% \State Train $B$ models $\hat{f}_{\theta_b}$ using $\mathcal{D}$
% \State $m_{\pi} \varleftarrow \textit{mean}(\{\bv{u}_i\}_{i=0,\dots, N_{init}})$, $\Sigma_{\pi} \varleftarrow \textit{std}(\{\bv{u}_i\}_{i=0,\dots,N_{init}})$, initialize $\pi \sim \mathcal{N}(m_{\pi}, \Sigma_{\pi})$
% \For{$k = 0, 1, \dots ,N_{train}$}
% \For{$t = 0, 1, \dots , T_{ep}-1$}
% \State Read the state of the system $\bv{x}_t$
% \State $\mathcal{R}, \mathcal{R}_{next} \varleftarrow \{0,0,\dots,0\}$ (initialize two arrays of zeros with shape $(pop\_size \times P, B)$)
% \State Initialize gradients: $\mathcal{L} \varleftarrow \{0,0,\dots,0\}$, (shape: $(pop\_size \times P, B)$)
% \State $\mathcal{X} \varleftarrow \{\bv{x}_t, \bv{x}_t, \dots, \bv{x}_t\}$, (Repeat $\bv{x}_t$ in an array with shape $(pop\_size \times P, B)$)
% \State Sample from $\pi$ two batches of actions $\mathcal{U}$ and $\mathcal{U}_{next}$ of dimension $(pop\_size, T)$
% \State Repeat the sequences of actions $\mathcal{U}$ and $\mathcal{U}_{next}$ in the shapes $(pop\_size \times P, T, B)$
% \For{$i = 0, 1, \dots , T-1$}
% \State $\mathcal{X}_{next},\mathcal{X}_{next,next} \varleftarrow \{0,0,\dots,0\}$ (shape: $(pop\_size \times P,B)$)
% \For{$b = 1, 2, \dots , B$}
% \For{$j = 1, 2, \dots , (pop\_size \times P)$}
% \State Apply the action $\mathcal{U}[j,i,b]$ to $\mathcal{X}[j,b]$
% \State Read the new predicted state $\hat{\bv{x}}_{i+1} \sim \hat{f}_{\theta_b}(\hat{\bv{x}}_i, \mathcal{U}[j,i,b])$ and the reward $r_{i+1}$
% \State $\mathcal{X}_{next}[j,b] \varleftarrow \hat{\bv{x}}_{i+1}$, $\mathcal{R}[j,b] \varleftarrow \mathcal{R}[j,b] + r_{i+1}$

% % \If{i == 0}
% \State Apply actions $\mathcal{U}$ to $\mathcal{X}$
% \State Propagate $P$ state $\bv{x}_{t}$ particles through ensemble of $B$ models $\hat{f}_{\theta_b}$
% \State Compute $\mathcal{R}$ and $\mathcal{R}_{next}$
% \State Compute gradients $\mathcal{L}$
% \EndIf
% \State Apply the action $\mathcal{U}_{next}[j,i,b]$ to $\mathcal{X}_{next}[j,b]$
% \State Read $\hat{\bv{x}}_{i+2} \sim \hat{f}_{\theta_b}(\mathcal{X}_{next}[j,b], \mathcal{U}_{next}[j,i,b])$, and $r_{next_{i+1}}$
% \State $\mathcal{X}_{next,next}[j,b] \varleftarrow \hat{\bv{x}}_{i+2}$, $\mathcal{R}_{next}[j,b] \varleftarrow \mathcal{R}_{next}[j,b] + r_{next_{i+1}}$
% \EndFor
% \EndFor
% \State $\mathcal{X} \varleftarrow \mathcal{X}_{next}$
% \EndFor
% % \State Average $\mathcal{R}$, $\mathcal{R}_{next}$, and $\mathcal{L}$ over $P$ (shape: $(pop\_size, B)$)
% % \State Average $\mathcal{U}$ over $P$ and $B$ (resulting shape: $(pop\_size, T)$)
% \State $\mathcal{J}_{\Omega_{\epsilon}(\hat{\mu})}(\bv{x}_t, \bv{u}_{t:t+T-1}) \varleftarrow \frac{1}{B}\sum\limits_{b=1}^B \mathcal{R}[:,b] -\epsilon \biggr(\frac{1}{B}\sum\limits_{b=1}^B \norm{\mathcal{L}[:,b]\mathcal{R}_{next}[:,b]}^2\biggl)^{\frac{1}{2}}$
% \State Order $\mathcal{J}_{\Omega_{\epsilon}(\hat{\mu})}(\bv{x}_t, \bv{u}_{t:t+T-1})$ in descending order and accordingly also $\mathcal{U}$
% \State Store the first $n\_elite$ rows of $\mathcal{U}$ in $\mathcal{U}_{elite}$, (new shape $(n\_elite, T)$)
% \State $m_{\pi} \varleftarrow mean(\mathcal{U}_{elite})$, $\Sigma_{\pi} \varleftarrow std(\mathcal{U}_{elite})$
% \State $\bv{u}_t^{\ast} \varleftarrow mean(\mathcal{U}_{elite}[0,0])$
% \State Execute $\bv{u}_t^{\ast}$ on the real system, observe $\bv{x}_{t+1} \sim f(\bv{x}_t, \bv{u}_t^{\ast})$
% \State $\mathcal{D} \varleftarrow \{\mathcal{D} \cup (\bv{x}_t,\bv{u}_t, \bv{x}_{t+1})\}$
% \EndFor
% \State Train $\hat{f}_{{\theta_b}},\;\forall b$ using $\mathcal{D}$
% \EndFor
% \end{algorithmic}\label{alg:STQL}
% \end{algorithm}


\section{Numerical results}

% \grtodo{Here we go back to DR-PETS...but there is no mention to it in the contributions.}
% \AAC{We evaluate our DR-PETS algorithm under epistemic uncertainty, arising from deploying the learned policy on environments that differ from the nominal system used during training.} To investigate this, we conduct experiments on the Pendulum and Cartpole swing-up tasks, comparing our approach with the PETS algorithm~\cite{chua2018}. We develop custom environments that allow for perturbations in system parameters, thereby explicitly challenging the policy with unseen dynamics. For each task, a 200-step horizon is considered, and an ensemble of \( B = 5 \) neural networks is employed to learn the transition probability functions, with \( Q = 10 \) propagated particles. The B-NN ensemble is trained on the nominal environment for 100 episodes, after which we evaluate the robustness of the obtained DR-PETS policy over a range of perturbed environments. The code for the DR-PETS implementation with the custom environment is available at \url{https://tinyurl.com/595c76t6}.
We evaluate DR-PETS under epistemic uncertainty by testing its learned policy in perturbed environments distinct from the nominal training setup. Experiments on Pendulum and Cartpole swing-up tasks compare DR-PETS with PETS~\cite{chua2018}, using custom environments with parameter perturbations to challenge policy robustness. Each task runs for 200 steps with an ensemble of \( B = 5 \) neural networks modeling transition probabilities and \( Q = 10 \) propagated particles. The B-NN ensemble is trained on the nominal environment for 100 episodes, after which DR-PETS' robustness is assessed across perturbed settings. The implementation and custom environments are available at \url{https://tinyurl.com/595c76t6}.

% \grtodo{Credo dovremmo essere piu' espliciti su come definiamo le incertezze. La cosa interessante e' che facciamo deployment su un sistema differente da quello usato durante il training. La direi proprio cosi'.}
% \subsection{Example: Pendulum task}

% \grtodo{In questo paragrafo dettaglierei quale e' l'incertezza, legandolo al paragrafo precedente (deployment).}
\subsection{Example: Pendulum task}
We start with a nominal pendulum of mass 1 kg and length 0.5 m. Using the algorithm and code from \cite{chua2018}, we train an ensemble of 5 NNs to estimate the transition probability function and solve the MPC planning problem via CEM. DR-PETS differs from PETS by optimizing a distributionally robust MPC objective \eqref{eqn:thm_statement}, whereas PETS uses \eqref{eqn:Chua_obj}. To compare performance, we simulate pendulums with masses varying from 0.5 kg to 1.5 kg i.e. uncertainty set and evaluate total rewards across the set of masses.
% First, we begin with a nominal pendulum having a weight ($w$) of 1 kg and a length ($l$) of 0.5 m. Subsequently, we employ the algorithm and code provided by \cite{chua2018} to train an ensemble of 5 NNs that estimates the transition probability function. We solve the MPC planning problem using this NN ensemble in conjunction with the CEM. What distinguishes DR-PETS from PETS is the distributionally robust MPC objective. PETS optimizes actions ($u$) based on the objective function defined in \eqref{eqn:Chua_obj}, while DR-PETS tackles the objective outlined in \eqref{eqn:thm_statement}, namely  a DR objective function. To assess the performance of DR-PETS relative to PETS, we conduct simulations involving pendulums with varying masses, ranging from $0.5$ kg to $1.5$ kg i.e. the uncertainty set, and then assess the total rewards achieved for each mass value.
% \begin{figure}[h!]
%     \centering
%     \includegraphics[width=1\textwidth]{Figures/Pendulum_reward_DRvsNom_2.png}
%     \caption{Total episodic reward obtained by the PETS (in blue) and DR-PETS (in red) for perturbation of pendulum mass. Shaded region denotes one standard deviation.  }
%     \label{fig:pendulum_plot}
% \end{figure}

\begin{figure}[h!]
    \centering
    \includegraphics[width=0.5\textwidth]{Figures/Pendulum_reward_DRvsNom.png}
    \caption{Total episodic reward obtained by the PETS (in blue) and DR-PETS (in red) for perturbation of pendulum mass. Shaded region denotes half of one standard error.}
    \label{fig:pendulum_plot}
\end{figure}

The plot in figure~\ref{fig:pendulum_plot} is obtained after 50 simulations for each mass perturbation. The performance difference between PETS and DR-PETS is under parameter perturbations is marginal, although the proposed algorithm exhibits more robustness to model parameter uncertainty. Indeed, when the mass value is higher than 1 kg, the DR-PETS reaches slight higher cumulative reward. %This benefit can become very important in more complex tasks.
%\grtodo{Questo paragrafo mi sembra contro producente. Dobbiamo commentare la figura ma non si capisce l'ultima frase...o siamo piu' precisi oppure rimuoviamola...}
\subsection{Example: Cartpole task}
\begin{figure}[h!]
    \centering
    \includegraphics[width=0.5\textwidth]{Figures/cartpole_plot.png}
    \caption{Total episodic reward obtained by the PETS (in blue) and DR-PETS (in red) for perturbation of pole length. Shaded region denotes half of one standard error.}
    \label{fig:cartpole_plot}
\end{figure}
We take a nominal cartpole system with the pole length ($l$) and mass ($m$) being 0.5 m and 1 kg respectively; we employ the algorithm and code provided by \cite{chua2018} to train an ensemble of 5 NNs that estimate the system dynamics. 
% We solve the MPC planning problem over a horizon ($T$) using this NN ensemble in conjunction with the cross-entropy optimizer. What distinguishes DR-PETS from PETS is the distributionally robust MPC objective. PETS optimizes actions ($u$) based on the objective function defined in \eqref{eqn:Chua_obj}, while DR-PETS tackles the objective outlined in \eqref{eqn:thm_statement}, which represents a distributionally robust objective function. 
To assess the performance of DR-PETS relative to PETS, we conduct simulations involving cartpole with varying pole lengths, ranging from 0.2 m to 0.8 m, and then assess the total rewards achieved for each length value.

Figure~\ref{fig:cartpole_plot} presents cumulative reward comparisons of DR-PETS and PETS over 50 simulations against length perturbation. Beyond a 0.6 m pole, both algorithms show a sharp performance decline, but DR-PETS degrades more gradually, maintaining higher cumulative rewards. Additionally, PETS exhibits a wider confidence interval, suggesting greater performance variability.
% The plot in figure~\ref{fig:cartpole_plot} is obtained after 50 simulations for each length perturbation. The plot shows comparison of DR-PETS with PETS for cumulative reward. Beyond the 0.6 m pole length, both algorithms experience a sharp decline in performance as the pole length increases. However, DR-PETS appears to have a slower rate of decline compared to PETS, resulting in DR-PETS having a higher cumulative reward than DR-PETS at pole lengths beyond 0.6 m. Also, the confidence interval for PETS is consistently wider than for DR-PETS, which might indicate that has a higher variability in its performance.
\section*{Acknowledgments}
This research would have been impossible without the foundational contributions and invaluable guidance of Yannis~Paschalidis and Jimmy~Queeney.  
\appendix
\subsection{Special cases: $p=\infty$ and $p=1$}
When $p=\infty$, Wasserstein duality results in 

\begin{equation}
\begin{aligned}
     &\frac{1}{B}\sum\limits_{b=1}^B\inf_{\tilde{f}_b}
    \bigr[\mathcal{J}_{t}(\tilde{f}_{b}) \mid d(\tilde{f}_b, \hat{f}_{\theta_{b}})\leq\epsilon\bigl]
    \\  &=  \frac{1}{B}\sum\limits_{b=1}^B\inf_{\norm{v_{b}}\leq\epsilon}
    \bigr[\mathcal{J}_{t}(\hat{f}_{\theta_{b}}+v_{b})\bigl]
\end{aligned}
\end{equation}

By considering a first order approximation of $\mathcal{J}(\cdot)$ with respect to the inputs $\tilde{f}$, we have

\begin{equation}
\begin{aligned}
    &\frac{1}{B}\sum\limits_{b=1}^B\inf_{\norm{v_{b}}\leq\epsilon}
    \bigr[\mathcal{J}_{t}(\hat{f}_{\theta_{b}}+v_{b})\bigl]
    % &=\frac{1}{B}\sum\limits_{b=1}^B\inf_{\norm{v_{b}}\leq\epsilon}\biggr[\mathcal{J}_{t}(\hat{f}_{\theta_{b}})+\left<\nabla_{\hat{f}_{\theta_{b}}}\mathcal{J}_{t}(\hat{f}_{\theta_{b}}),v_{b}\right>\biggl]\\
    % &=\frac{1}{B}\sum\limits_{b=1}^B\mathcal{J}_{t}(\hat{f}_{\theta_{b}})+\frac{1}{B}\sum\limits_{b=1}^B\inf_{\norm{v_{b}}\leq\epsilon}\biggr[\left<\nabla_{\hat{f}_{\theta_{b}}}\mathcal{J}_{t}(\hat{f}_{\theta_{b}}),v_{b}\right>\biggl]\\
    \\&=\frac{1}{B}\sum\limits_{b=1}^B\mathcal{J}_{t}(\hat{f}_{\theta_{b}})-\frac{\epsilon}{B}\sum\limits_{b=1}^B\norm{\nabla_{\hat{f}_{\theta_{b}}}\mathcal{J}_{t}(\hat{f}_{\theta_{b}})}_{*}.
    \end{aligned}
\end{equation}
As a result, the closed-form solution and adversarial perturbations follow the same form as in
Theorem \ref{thm:Jimmy}. Also note that $p = \infty$ results in the special case, where all inputs are perturbed using the same magnitude $\epsilon$.

%\subsection{Special case: $p=1$}
When $p=1$, all previous derivations hold through \eqref{eqn:delta_opt},

\begin{equation}\label{eqn:delta_opt_spcl}
    \begin{split}
    &\inf\limits_{\delta_{b}\geq0}\left\{\lambda\delta_{b}-\delta_{b}\norm{\nabla_{\hat{f}_{\theta_{b}}}\mathcal{J}_{t}(\hat{f}_{\theta_{b}})}_{*}\right\}\\
    &=\inf\limits_{\delta_{b}\geq0}\left\{\delta_{b}\left(\lambda-\norm{\nabla_{\hat{f}_{\theta_{b}}}\mathcal{J}_{t}(\hat{f}_{\theta_{b}})}_{*}\right)\right\},
    \end{split}
\end{equation}

We can see that in \eqref{eqn:delta_opt_spcl} equals zero when $\lambda\geq\norm{\nabla_{\hat{f}_{\theta_{b}}}\mathcal{J}_{t}(\hat{f}_{\theta_{b}})}_{*}$ and $-\infty$ otherwise. Then while considering the outer maximization problem
\begin{equation*}
    \sup_{\lambda \geq 0}\biggr\{\frac{1}{B}\sum\limits_{b=1}^B\inf_{v_{b}}\left[\left<\nabla_{\hat{f}_{\theta_{b}}}\mathcal{J}_{t}(\hat{f}_{\theta_{b}}),v_{b}\right> + \lambda\norm{v_{b}}\right] - \lambda \epsilon\biggl\},
\end{equation*}
we must have $\lambda^{*}=\max_{b=1,\dots,B}\norm{\nabla_{\hat{f}_{\theta_{b}}}\mathcal{J}_{t}(\hat{f}_{\theta_{b}})}_{*}$, which results in the closed-form solution for the Wasserstein distributionally robust objective to become

\begin{equation}\label{eqn:p=1 objective}
    \frac{1}{B}\sum\limits_{b=1}^B\mathcal{J}_{t}(\hat{f}_{\theta_{b}})-\epsilon\max_{b=1,\dots,B}\norm{\nabla_{\hat{f}_{\theta_{b}}}\mathcal{J}_{t}(\hat{f}_{\theta_{b}})}_{*}
\end{equation}

The form of \eqref{eqn:p=1 objective} resembles Theorem \ref{thm:Jimmy}. However, there is no adversarial interpretation when $p = 1$,
as the magnitude of each adversarial perturbation is chosen to be zero in \eqref{eqn:delta_opt_spcl}.
% \begin{figure}[h!]
%     \centering
%     \includegraphics[width=0.5\textwidth]{Figures/cartpole_plot.png}
%     \caption{Total episodic reward obtained by the PETS (in blue) and DR-PETS (in red) for perturbation of pole length. Shaded region denotes half of one standard error.}
%     \label{fig:cartpole_plot}
% \end{figure}


\bibliographystyle{ieeetr}
\bibliography{root_sub}




\end{document}

\begin{thebibliography}{10}

\bibitem{GARRABE202281}
E.~Garrabe and G.~Russo, ``Probabilistic design of optimal sequential decision-making algorithms in learning and control,'' {\em Annual Reviews in Control}, vol.~54, pp.~81--102, 2022.

\bibitem{Yu2015}
P.~Yu and H.~Xu, ``Distributionally robust counterpart in {M}arkov decision processes,'' {\em IEEE Transactions on Automatic Control}, vol.~61, 01 2015.

\bibitem{pmlr-v168-lellis22a}
F.~D. Lellis, M.~Coraggio, G.~Russo, M.~Musolesi, and M.~di~Bernardo, ``Control-tutored reinforcement learning: Towards the integration of data-driven and model-based control,'' in {\em Proceedings of The 4th Annual Learning for Dynamics and Control Conference} (R.~Firoozi, N.~Mehr, E.~Yel, R.~Antonova, J.~Bohg, M.~Schwager, and M.~Kochenderfer, eds.), vol.~168 of {\em Proceedings of Machine Learning Research}, pp.~1048--1059, PMLR, 23--24 Jun 2022.

\bibitem{Klink2021}
P.~Klink, {\em Model-Based Reinforcement Learning from PILCO to PETS}, p.~165–175.
\newblock Springer International Publishing, 2021.

\bibitem{chua2018}
K.~Chua, R.~Calandra, R.~McAllister, and S.~Levine, ``{D}eep reinforcement learning in a handful of trials using probabilistic dynamics models,'' in {\em Proceedings of the 32nd International Conference on Neural Information Processing Systems}, NIPS'18, (Red Hook, NY, USA), p.~4759–4770, Curran Associates Inc., 2018.

\bibitem{pmlr-v100-okada20a}
M.~Okada and T.~Taniguchi, ``Variational inference mpc for bayesian model-based reinforcement learning,'' in {\em Proceedings of the Conference on Robot Learning} (L.~P. Kaelbling, D.~Kragic, and K.~Sugiura, eds.), vol.~100 of {\em Proceedings of Machine Learning Research}, pp.~258--272, PMLR, 30 Oct--01 Nov 2020.

\bibitem{Iyengar2005}
G.~Iyengar, ``Robust dynamic programming,'' {\em Math. Oper. Res.}, vol.~30, pp.~257--280, 05 2005.

\bibitem{Derman2023}
E.~Derman, Y.~Men, M.~Geist, and S.~Mannor, ``{R}obustness and regularization in reinforcement learning,'' in {\em NeurIPS 2023 Workshop on Generalization in Planning}, 2023.

\bibitem{Gao2023}
R.~Gao, ``Finite-sample guarantees for wasserstein distributionally robust optimization: Breaking the curse of dimensionality,'' {\em Operations Research}, vol.~71, no.~6, pp.~2291--2306, 2023.

\bibitem{derman2020distributionalregularization}
E.~Derman and S.~Mannor, ``Distributional robustness and regularization in reinforcement learning,'' {\em arXiv preprint arXiv:2003.02894}, 2020.

\bibitem{liu2022DRQ}
Z.~Liu, Q.~Bai, J.~Blanchet, P.~Dong, W.~Xu, Z.~Zhou, and Z.~Zhou, ``Distributionally robust {Q}-learning,'' in {\em International Conference on Machine Learning}, pp.~13623--13643, PMLR, 2022.

\bibitem{zhao2023minimax}
F.~Zhao and K.~You, ``Minimax {Q}-learning control for linear systems using the wasserstein metric,'' {\em Automatica}, vol.~149, p.~110850, 2023.

\bibitem{queeney2023optimal}
J.~Queeney, E.~C. Ozcan, I.~Paschalidis, and C.~Cassandras, ``Optimal transport perturbations for safe reinforcement learning with robustness guarantees,'' {\em Transactions on Machine Learning Research}, 2024.

\bibitem{mark2020stochastic}
C.~Mark and S.~Liu, ``Stochastic {MPC} with distributionally robust chance constraints,'' {\em IFAC-PapersOnLine}, vol.~53, no.~2, pp.~7136--7141, 2020.

\bibitem{coppens2021datadriven}
P.~Coppens and P.~Patrinos, ``Data-driven distributionally robust {MPC} for constrained stochastic systems,'' {\em IEEE Control Systems Letters}, vol.~6, pp.~1274--1279, 2021.

\bibitem{wu2022ambiguity}
F.~Wu, M.~E. Villanueva, and B.~Houska, ``Ambiguity tube {MPC},'' {\em Automatica}, vol.~146, p.~110648, 12 2022.

\bibitem{blanchet2016quantifying}
J.~Blanchet and K.~Murthy, ``Quantifying distributional model risk via optimal transport,'' {\em Mathematics of Operations Research}, vol.~44, no.~2, pp.~565--600, 2019.

\bibitem{shi2022distributionally}
L.~Shi and Y.~Chi, ``Distributionally robust model-based offline reinforcement learning with near-optimal sample complexity,'' {\em Journal of Machine Learning Research}, vol.~25, no.~200, pp.~1--91, 2024.

\bibitem{goodfellow2014explaining}
I.~J. Goodfellow, J.~Shlens, and C.~Szegedy, ``{E}xplaining and harnessing adversarial examples,'' in {\em 3rd International Conference on Learning Representations, {ICLR} 2015, San Diego, CA (USA), Conference Track Proceedings}, 5 2015.

\bibitem{botev2013cross}
Z.~I. Botev, D.~P. Kroese, R.~Y. Rubinstein, and P.~L’Ecuyer, ``{T}he cross-entropy method for optimization,'' in {\em Handbook of statistics}, vol.~31, pp.~35--59, Elsevier, 2013.

\end{thebibliography}