# Efficient Algorithm Level Error Detection for Number-Theoretic Transform
Assessed on FPGAs

Kasra Ahmadi, Saeed Aghapour, Mehran Mozaffari Kermani, Senior Member, IEEE,
and Reza Azarderakhsh, Member, IEEE

###### Abstract

Polynomial multiplication stands out as a highly demanding arithmetic process
in the development of post-quantum cryptosystems. The importance of number-
theoretic transform (NTT) extends beyond post-quantum cryptosystems, proving
valuable in enhancing existing security protocols such as digital signature
schemes and hash functions. Due to the potential for errors to significantly
disrupt the operation of secure, cryptographically-protected systems,
compromising data integrity, and safeguarding against side-channel attacks
initiated through faults it is essential to incorporate mitigating error
detection schemes. This paper introduces algorithm level fault detection
schemes in NTT multiplication, representing a significant enhancement compared
to previous research. We evaluate this through the simulation of a fault
model, ensuring that the conducted assessments accurately mirror the obtained
results. Consequently, we attain a notably comprehensive coverage of errors.
Finally, we assess the performance of our efficient error detection scheme on
FPGAs to showcase its implementation and resource requirements. Through
implementation of our error detection approach on Xilinx/AMD Zynq Ultrascale+
and Artix-7, we achieve a comparable throughput with just a 9% increase in
area and 13% increase in latency compared to the original hardware
implementations.

###### Index Terms:

Fault detection, fast Fourier transform (FFT), FPGA, number-theoretic
transform (NTT).

##  I Introduction

Fast Fourier Transform (FFT) [1] algorithms which are used to compute the
Discrete Fourier Transform (DFT) have various applications, spanning from
digital signal processing to the efficient multiplication of large integers.
When the coefficients of the polynomial are specifically chosen from a finite
field, the resulting transform is known as the Number Theoretic Transform
(NTT) [2], and it can be computed using FFT algorithms designed for operations
within this particular finite field. An efficient approach for polynomial
multiplication, NTT holds significant importance in post-quantum
cryptosystems, e.g., lattice-based cryptosystems which is regarded as a
leading contender for quantum-secure public key cryptography, primarily
because of its broad applicability and security proofs grounded in the worst-
case hardness of established lattice problems. The NTT has proven to be a
potent tool that facilitates the computation of this operation with quasi-
polynomial complexity $O(n.lgn)$. Several recent studies have been conducted
on optimizing NTT [3], [4], [5], [6], [7], [8], and [9]. In addition to
polynomial multiplication, the use of NTT has the potential to significantly
enhance existing schemes by improving their security parameters. NTT is widely
used in signature schemes, hash functions, and identification schemes. As a
result, incorporating efficient error detection mechanisms in NTT for
polynomial multiplication will not only improve the security and reliability
but also mitigate the risk of fault attacks in the respective algorithms in
post-quantum cryptography. Moreover, such schemes could at least alleviate the
attack surface or be add-ons to other approaches.

###  I-A Related works

We catogorize relevant studies into three subsets. The initial category
centers on fault attacks and countermeasures applied to post-quantum
cryptographic schemes, particularly those extensively utilizing NTT. Ravi _et
al._ [10] presented the first fault injection analysis of the NTT. The authors
found a significant flaw in the way NTT is implemented in the pqm4 library
[11]. The identified vulnerability, referred to as twiddle-pointer, is
exploited to demonstrate practical and effective attacks on Kyber and
Dilithium [12]. Several other authors have reported employing fault injection
on structured lattice-based schemes as a foundation for attacks [13], [14],
[15], and [16] threatening the security of implementations.

The focus of the second category lies in offering fault detection
methodologies pertaining to the algorithmic level of classical cryptographic
schemes. In [17], the authors suggested an effective algorithmic-level error
detection for the ECSM window method, aiming to identify both permanent and
transient errors. In [18], the authors presented error detection schemes
designed for the double-and-add-always and Montgomery ladder ECSM algorithm
utilized in non-supersingular elliptic curves.

The third category focuses on developing fault detection schemes specifically
for NTT. In [19], the authors integrate the advantages of bitslicing, a
software implementation technique where a datapath of an $n$-bit processor is
treated as $n$ parallel single-bit datapaths, to devise a fault countermeasure
for the NTT used in Dilithium [20]. Sarker _et al._ [21] and [22], presented
error detection architectures of NTT based on recomputation with encoded
operands. The authors achieved high error coverage with low area overhead.
Despite the recomputation mechanism, the latency is elevated. We note that
despite high error coverage, for fast implementations, these works might
increase the total time to levels not acceptable.

###  I-B Our Major Contributions

  * •

We have proposed an algorithm level fault detection scheme of NTT
multiplication which is widely used in lattice based cryptography. We achieved
high error coverage with less latency compared to previous works. This was
achieved through the implementation of algorithm-level error detection for NTT
section and partial recomputation for pre-computation sub-block. We performed
simulations for single and burst fault injection using our proposed scheme.
The simulation outcomes demonstrated that our approach is capable of detecting
diverse types of faults with high error coverage and offers protection against
the NTT fault attack presented in [10] and those with respective fault models.

  * •

Our error detection method is deployed on Xilinx/AMD Zynq Ultrascale+ and
Artix-7. The results of our implementations indicate that we can attain a
significantly high level of error coverage with only a 13% increase in latency
and a 9% area bloat.

##  II Preliminaries

###  II-A Number Theoretic Transform (NTT)

The NTT is a specific form of the discrete Fourier transform, where the
coefficient ring is selected to be a finite field that includes the
appropriate roots of unity. Let $\mathbb{R}=\mathbb{Z}[X]/(X^{n}+1)$ be the
ideal lattice. Also, $n$ is power of 2 and $q\,\equiv\,1\,$(mod $2n$) where
$q$ is a prime number. Let $f(x)=f_{0}+f_{1}x+f_{2}x^{2}+...+f_{n-1}x^{n-1}$
be an irreducible polynomial of degree $n$ and let $\omega$ be a primitive
$n$-th root of unity in $\mathbb{Z}_{q}$ which $\omega^{n}\,\equiv\,1$ (mod
$q$) and $a=(a[0],a[1],...a[n-1])\in\mathbb{Z}_{q}$ and
$b=(b[0],b[1],...b[n-1])\in\mathbb{Z}_{q}$. The transformations
$\widetilde{a}=\text{NTT}(a)$ and $\widetilde{b}=\text{NTT}(b)$ are defined as
$\widetilde{a}[i]\equiv\sum_{j=0}^{n-1}a[j]\omega^{ij}\,$(mod $q$) and
$\widetilde{b}[i]\equiv\sum_{j=0}^{n-1}b[j]\omega^{ij}\,($mod $q$) for
$i=0,1,...,n-1$, and we have $\text{INTT}(\text{NTT}(a))=a.$

The NTT of a sequence $a$ is derived as:

$\widetilde{a}=NTT(a)=\left[\begin{array}[]{cccc}\omega^{0}&\omega^{0}&...&\omega^{0}\\\
\omega^{0}&\omega^{1}&...&\omega^{n-1}\\\
\omega^{0}&\omega^{2}&...&\omega^{2(n-1)}\\\ .&.&...&.\\\
\omega^{0}&\omega^{n-1}&...&\omega^{(n-1)^{2}}\end{array}\right]*\left[\begin{array}[]{c}a(0)\\\
a(1)\\\ a(2)\\\ .\\\ a(n-1)\end{array}\right]$ (1)

In Algorithm 1, the iterative NTT implementation derives the NTT of a
specified polynomial $a$. The $\text{bit\\_reverse}(k)$ function (presented in
Line 5) rearranges the input $k$ where the new placement of elements is
determined by reversing the binary representation of the operand. Lines 9 and
10 perform the butterfly operation. Every butterfly module, taking inputs $a$
and $b$ and producing outputs $c$ and $d$, executes the following butterfly
computation: $c=a+b\,\omega^{k}$ and $d=a-b\,\omega^{k}$.

Algorithm 1 Iterative NTT [24]

Input: $a\in\mathbb{Z}_{q}$ of length $n=2^{k}$, $\omega$ is primitive $n$-th
root of unity

Output: $\widetilde{a}=\text{NTT}(a)$ in bit-reversed order

1: $\widetilde{a}=a$

2: for $s=k$ to $1$:

3: $m=2^{s}$

4: for $k=0$ to $2^{k-s}-1$:

5: $\omega=\omega^{\text{bit-reverse}(k)\,.\,m/2}$

6: for $j=0$ to $m/2-1$:

7: $u=\widetilde{a}[k.m+j]$

8: $t\equiv\omega\,.\,\widetilde{a}[k.m+j+m/2]$ mod $q$

9: $\widetilde{a}[k.m+j]\equiv(u+t)$ mod $q$

10: $\widetilde{a}[k.m+j+m/2]\equiv(u-t)$ mod $q$

11: return $\widetilde{a}$

###  II-B Negative Wrapped Convolution

NTT can be used to perform polynomial multiplication. However, performing the
NTT provides a cyclic convolution, thus computing $c\equiv a\,.\,b$ mod
$(X^{n}+1)$ would require applying the NTT of length $2n$. This problem can be
resolved by utilizing negative wrapped convolution, which prevents the
doubling of input length [23]. Let $\psi$ be a primitive 2$n$-th root of unity
in $\mathbb{Z}_{q}$ where $\psi^{2}=\omega$. Moreover, there exists a pre-
process module which is described in Algorithm 2. We can achieve post-process
function by changing $\psi$ to $\psi^{-1}$ in Algorithm 2. The negative
wrapped convolution of $a$ and $b$ is defined as $c=\text{post-
process}(\text{INTT}(\text{NTT}(\hat{a})\,\Circle\text{\,NTT}(\hat{b})))$
where $\hat{a}=\text{pre-process}(a)$, $\hat{b}=\text{pre-process}(b)$, and
the symbol $\Circle$ represents component-wise multiplication.

Algorithm 2 Pre-processing step for NTT calculations

Input: $a\in\mathbb{Z}_{q}$ of length $n$, $\psi$ primitive 2$n$-th root of
unity

Output: $\hat{a}=\text{pre-process}(a)$

1: $\hat{a}=a$

2: for $i=0$ to $n-1$:

3: $\hat{a}[i]=a[i]\,.\,$$\psi^{i}$

4: return $\hat{a}$

##  III Proposed Error Detection Scheme

In this section, we present our error detection schemes on polynomial
multiplication, $\text{NTT}(\hat{a})\,\Circle\,\text{NTT}(\hat{b})$ where
$\hat{a}=\text{pre-process}(a)$ and $\hat{b}=\text{pre-process}(b)$. We have
devised an error detection scheme at the algorithm level on the component-wise
NTT multiplication sub-block,
$\text{NTT}(\hat{a})\,\Circle\,\text{NTT}(\hat{b})$, as well as an error
detection scheme for the pre-process function, $\hat{a}=\text{pre-
process}(a)$, involving careful utilization of recomputation using shifted
operands.

Figure 1: Concurrent error detection scheme for NTT operation which is used in
[25].

####  III-1 Component-wise NTT Multiplication

Our focus is on $\text{NTT}(\hat{a})\,\Circle\,\text{NTT}(\hat{b})$ operation
in this section. Jou _et al._ [25], proposed an algorithm level error
detection scheme on FFT network. Their proposed error detection scheme is
depicted in Fig. 1. In order to achieve an efficient scheme, we have adapted
the aforementioned approach to component-wise NTT multiplication. The primary
focus is to use an error detection scheme with small overhead and high error
detection ratio with respect to efficient realizations of the NTT. From NTT
definition, we can use the relationships presented below:

$\widetilde{a}(0)=\sum_{j=0}^{n-1}a(j)\omega^{0}=\sum_{j=0}^{n-1}a(j),\quad
j=0,1,...,n-1$ (2)
$\widetilde{b}(0)=\sum_{j=0}^{n-1}b(j)\omega^{0}=\sum_{j=0}^{n-1}b(j),\quad
j=0,1,...,n-1$ (3)

The result of the NTT multiplication for index 0 can be written as:

$\widetilde{a}(0)\widetilde{b}(0)=\sum_{j=0}^{n-1}a(j)\sum_{j=0}^{n-1}b(j)$
(4)

We note that (4) is used in our proposed error detection scheme.

If we rotate the input sequence of (1) by one, then we get:

$\displaystyle\left[\begin{array}[]{c}\widetilde{a}(0)\\\ \widetilde{a}(1)\\\
\widetilde{a}(2)\\\ .\\\ \widetilde{a}(N-1)\end{array}\right]$
$\displaystyle=\rho*\left[\begin{array}[]{c}\hat{a}(1)\\\ \hat{a}(2)\\\ .\\\
\hat{a}(N-1)\\\ \hat{a}(0)\end{array}\right]$ (15)

or

$\displaystyle\left[\begin{array}[]{c}\widetilde{a}(0)/\omega^{0}\\\
\widetilde{a}(1)/\omega^{1}\\\ \widetilde{a}(2)/\omega^{2}\\\ .\\\
\widetilde{a}(n-1)/\omega^{(n-1)}\end{array}\right]$
$\displaystyle=\theta*\left[\begin{array}[]{c}\hat{a}(1)\\\ \hat{a}(2)\\\ .\\\
\hat{a}(n-1)\\\ \hat{a}(0)\end{array}\right]$ (26)

where the symbol * denotes matrix multiplication,
$\rho=\left[\begin{array}[]{cccc}\omega^{0}&\omega^{0}&...&\omega^{0}\\\
\omega^{1}&\omega^{2}&...&\omega^{0}\\\
\omega^{2}&\omega^{4}&...&\omega^{0}\\\ .&.&...&.\\\
\omega^{n-1}&\omega^{2(n-1)}&...&\omega^{0}\end{array}\right]$ in (5), and
$\theta=$$\left[\begin{array}[]{cccc}\omega^{0}&\omega^{0}&...&\omega^{0}\\\
\omega^{0}&\omega^{1}&...&\omega^{n-1}\\\
\omega^{0}&\omega^{2}&...&\omega^{2(n-1)}\\\ .&.&...&.\\\
\omega^{0}&\omega^{n-1}&...&\omega^{(n-1)^{2}}\end{array}\right]$ in (6). From
$\alpha$ (in (1)) + $\beta$ (in (6)), we get the following:

$\displaystyle\left[\begin{array}[]{c}\widetilde{a}(0)(\alpha+\beta\omega^{-0})\\\
\widetilde{a}(1)(\alpha+\beta\omega^{-1})\\\
\widetilde{a}(2)(\alpha+\beta\omega^{-2})\\\ .\\\
\widetilde{a}(n-1)(\alpha+\beta\omega^{-(n-1)})\end{array}\right]$
$\displaystyle=\theta*\left[\begin{array}[]{c}\alpha\hat{a}(0)+\beta\hat{a}(1)\\\
\alpha\hat{a}(1)+\beta\hat{a}(2)\\\ .\\\
\alpha\hat{a}(n-2)+\beta\hat{a}(n-1)\\\
\alpha\hat{a}(n-1)+\beta\hat{a}(0)\end{array}\right]$ (37)

where $\alpha$ and $\beta$ in (7) are scalars which can be selected
arbitrarily. Without losing generality and for the sake of simplicity, we
select $\alpha=2$ and $\beta=1$. Let us denote $\hat{a}^{m}$ and
$\hat{b}^{m}$, respectively, as the inputs which are shifted by $m$. In other
words, we can achieve:

$\widetilde{a}=\frac{1}{(\alpha+\beta\omega^{-k})}NTT(\alpha\hat{a}+\beta\hat{a}^{1})$
(38)

$\widetilde{b}=\frac{1}{(\alpha+\beta\omega^{-k})}NTT(\alpha\hat{b}+\beta\hat{b}^{1})$
(39)

Let us denote $c$ as the component-wise NTT multiplication of $\hat{a}$ and
$\hat{b}$. We can get:

$\displaystyle c$ $\displaystyle=NTT(\hat{a})\,\Circle\,NTT(\hat{b})=$
$\displaystyle\frac{1}{(\alpha+\beta\omega^{-k})^{2}}NTT(\alpha\hat{a}+\beta\hat{a}^{1})\,\Circle\,NTT(\alpha\hat{b}+\beta\hat{b}^{1})$
(40)

The proposed error detection scheme is depicted in Fig. 2. The Decoder_2 sub-
block multiplies its input with $\frac{1}{(\alpha+\beta\omega^{-k})^{2}}$,
which is the inverse of $(\alpha+\beta\omega^{-k})^{2}$ in
$\mathbb{R}=\mathbb{Z}[X]/(X^{n}+1)$. The scheme for error detection involves
comparing $c(0)$, obtained from (10), with
$\sum_{j=0}^{n-1}a(j)\sum_{j=0}^{n-1}b(j)$.

Figure 2: Proposed algorithm level error detection scheme for NTT
multiplication module.

####  III-2 Pre-processor

The pre-processor module performs element-by-element multiplication on two
input arrays. In proposing the error detection scheme, we have applied an
additional recomputing with shifted operands. As shown in Fig. 3, the encoding
and decoding modules constitute shifting which is free in hardware. The pre-
processing module considers a slight overhead linked to NTT multiplication. As
a result, the error detection scheme employing recomputation does not notably
impact the overall performance.

Figure 3: Proposed error detection scheme for pre-process module through
recompuation with shifted operands.

##  IV Simulations and Fault Model

In this section, we present the utilized fault model and the simulation
results acquired by implementing Python code to show error detection coverage,
employing the presented fault model, and utilizing our efficient error
detection schemes in this paper.

###  IV-A Fault Model

We assume that faults occur in either the complex multiplier or adder, or a
specific group of input or output lines in the butterfly module which is used
in NTT. As depicted in Fig. 4, faults can happen at modules $\\{1,2,3\\}$
marked in the figure, which may result in erroneous outputs of a butterfly
sub-block used in NTT module. Fault occurrence follow a normal distribution
within butterfly modules in Fig. 2. Furthermore, as we provide fault detection
schemes over NTT multiplication, faults can happen during component wise
multiplication module.

Figure 4: The utilized fault model in this work for the butterfly sub-block.

###  IV-B Error Coverage Simulation Results

To assess the error coverage, we employed Python3 to perform simulations,
applying the described fault models and error detection techniques to both the
NTT and pre-process modules. The simulations encompassed a million samples,
using parameters identical to the standard Kyber $(n=256)$ and BLISS signature
[26] $(n=512)$. As shown in Table I, with increase in the number of faults, we
can attain an error detection ratio close to 100% for both modules. Higher
error coverage is achievable for identifying burst errors. We have done
extensive experiments injecting burst adjacent faults, the results of which
are not presented for the sake of brevity.

Table I: Achieved error detection ratios of the proposed schemes with 1, 2, 4,
8, and 16 faults occurrence for one million samples

Parameters |  |  ${}^{1}n=256,$${}^{2}\omega=17$, |  $n=512,\omega=2,401$,
---|---|---|---
| ${}^{3}q=3,329$ | $q=12,289$
Component |  | Pre-process | NTT | Pre-process | NTT
| Mult | Mult
| 1 | 99.6% | 51% | 99.7% | 53%
Number of | 2 | 99.9% | 67% | 99.9% | 70.6%
faults | 4 | 100% | 90.3% | 100% | 90.9%
| 8 | 100% | 99% | 100% | 99%
| 16 | 100% | 99.9% | 100% | 99.9%

1$n=$ Polynomial size, 2$\omega=$ Twiddle factor, 3$q=$ Prime number

##  V Xilinx/AMD FPGA Implementation results

To confirm the efficiency of our presented approach, we chose to assess its
performance by applying it to the NTT multiplication involving a 256-degree
polynomial. We conducted a benchmark for implementation on different FPGAs,
i.e., Xilinx/AMD Zynq Ultrascale+ and Artix-7. The results clearly demonstrate
that our proposed efficient error detection schemes maintain a modest overhead
while effectively achieving a high level of fault detection. We utilized the
High-Level Synthesis (HLS) Vitis development environment to transform our
proposed schemes into Register Transfer Level (RTL) hardware descriptions. The
generated IP imported to Vivado to report power, utilization, and latency.
Tables II and III demonstrate the results of our implementations and the
derivations for area, timing, power, and energy. In the Vivado synthesis tool
context, area is determined by combining Slices and DSPs, with a conversion
ratio where one DSP is deemed equivalent to 100 Slices. Our proposed error
detection designs achieve a maximum operational frequency of around 140 MHz
across all the FPGAs. Our proposed error detection scheme imposes a maximum
overhead of 9% in additional area and introduces a latency increase of up to
13% in clock cycles, at most considering different designs. Our simulation and
implementation code is publicly available in our
github111https://github.com/KasraAhmadi/NTT_Error_Detection.

###  V-A Implementation Optimization

Although HLS has shifted the design entry level of abstraction from RTL to
C/C++, practical implementation often requires significant source code
rewriting to make it HLD-ready. We have carefully taken this into effort as
discussed here. This includes the incorporation of pragmas to attain
satisfactory performance. Our intended area and timing efforts implementation
was realized by strategically inserting the following pragmas into our
program.

Table II: AMD/Xilinx Zynq Ultrascale+, xczu4ev-sfvc784-2-i Implementation
Results

Strategy |  | Area | Timing
---|---|---|---
|  | Effort | Effort
Scheme |  | Our | Baseline | Our | Baseline
scheme | work | scheme | work
| LUTs | 1,469 | 1,435 | 3,207 | 3,165
Area | FFs | 1,188 | 1,122 | 2,227 | 2,161
CLBs | 343 | 340 | 733 | 683
| DSPs | 33 | 30 | 54 | 51
Power (W) |  | 0.46 | 0.45 | 0.52 | 0.51
@ 140 MHz |
| Latency | 3,703 | 3,438 | 2,241 | 1,979
Timing | [CCs]
Total time | 26,439 | 24,547 | 16,000 | 14,130
| [ns]
Energy (nJ) |  | 12,161 | 11,046 | 8,320 | 7,206

Our Scheme: The design which includes error detection scheme.

Baseline work: The design which does not include any error detection scheme.

Table III: AMD/Xilinx Artix-7, xc7a75ti-ftg256-1L Implementation Results Strategy |  | Area | Timing
---|---|---|---
|  | Effort | Effort
Scheme |  | Our | Baseline | Our | Baseline
scheme | work | scheme | work
| LUTs | 1,569 | 1,530 | 2,956 | 2,939
Area | FFs | 1,693 | 1,597 | 2,979 | 2,883
SLICEs | 673 | 656 | 1,240 | 1,253
| DSPs | 33 | 30 | 54 | 51
Power (W) |  | 0.2 | 0.19 | 0.28 | 0.26
@ 140 MHz |
| Latency | 3,749 | 3,482 | 2,267 | 2,003
Timing | [CCs]
Total time | 26,767 | 24,861 | 16,186 | 14,301
| [ns]
Energy (nJ) |  | 5,353 | 4,723 | 4,532 | 3,718

####  V-A1 Pre-calculate the decoder module values

To enhance the efficiency of the proposed error detection scheme, we have the
option to pre-calculate the decoder module in (10), specifically computing
$\frac{1}{(\alpha+\beta\omega^{-k})^{2}}$ for various values of the parameter
$k$ and save them in memory.

####  V-A2 Utilizing task-level pipelining

We perform loop unrolling on the outer loop in Algorithm 1 (Line 2) and
transform the remaining code into a function named "stage," illustrating each
stage in the NTT architecture. We undertook this approach to leverage task-
level pipelining, allowing functions and loops to operate simultaneously. This
enhances the concurrency of the RTL implementation, resulting in an overall
increase in design throughput.

####  V-A3 Pipelining the stage function

This pragma was inserted to facilitate instruction-level pipelining, aiming to
enhance throughput and clock frequency within each NTT stage function. It is
important to acknowledge that this optimization entails the trade-off of
utilizing extra digital resources.

##  VI Conclusion

In this paper, we proposed an efficient algorithm level error detection scheme
over polynomial multiplication using NTT. We have aimed for low hardware
overhead and low latency, suitable for deeply-embedded systems. The presented
scheme effectively safeguards cryptographic algorithms that employ NTT
multiplication. By performing simulations, we demonstrated that the proposed
error detection scheme achieved extensive coverage of faults. Moreover, we
implemented our proposed design on two Xilinx/AMD FPGAs. Regarding overhead
and latency, the implementation led to minimal additional expenses in hardware
(at most 9% area, 13% latency, 7% power, and 13% energy). With high error
coverage and acceptable overhead, the proposed schemes in this works are
suitable for resource-constrained and sensitive usage models.

## Acknowledgments

This work was supported by the US National Science Foundation (NSF) through
the award SaTC-1801488.

## References

  * [1] W. T. Cochran, J. W. Cooley, D. L. Favin, H. D. Helms, R. A. Kaenel, W. W. Lang, G. C. Maling, D. R. Nelson, C. M. Rader, and P. Welch, “What is the fast Fourier transform?,” _IEEE Trans. Audio Electroacoust_ , 1967.
  * [2] J. W. Cooley and J. W. Tukey, “An algorithm for the machine calculation of complex Fourier series,” _Math. Comp._ , vol. 19, pp. 167-172, 1975.
  * [3] J. Hermelink, S. Streit, E. Strieder, and K. Thieme, “Adapting belief propagation to counter shuffling of NTTs,” _IACR Transactions on Cryptographic Hardware and Embedded Systems_ , vol. 2023, Nov. 2022.
  * [4] S. Kurniawan, P. Duong-Ngoc, and H. Lee, “Configurable memory based NTT architecture for homomorphic encryption,” _IEEE Trans. Circuits Syst. II_ , Exp. Briefs, vol. 70, no. 10, pp. 3942-3946, 2023.
  * [5] B. Li, Y. Yan, Y. Wei, and H. Han, “Scalable and parallel optimization of the number theoretic transform based on FPGA,” _IEEE Trans. Very Large Scale Integr. (VLSI) Syst._ , vol. 32, no. 2, pp. 291-304, Feb. 2024.
  * [6] Y. Geng, X. Hu, M. Li, and Z. Wang, “Rethinking parallel memory access pattern in number theoretic transform design,” _IEEE Trans. Circuits Syst. II,_ Exp. Briefs, vol. 70, no. 5, pp. 1689-1693, May 2023.
  * [7] J. Mu, H. Tan, J. Wu, H. Lu, C. Chang, S. Chen, S. Liang, J. Ye, H. Li, and X. Li, “Energy-efficient NTT design with one-bank SRAM and 2-D PE array,” _Design, Automation Test in Europe Conference_(DATE), vol. 19, pp. 3353-3365, 2023.
  * [8] T. Tosun, E. Savas, “Zero-value filtering for accelerating non-profiled side-channel attack on incomplete NTT-based implementations of lattice-based cryptography,” _IEEE Trans. Inf. Forensics Security_ , vol. 19, pp. 3353-3365, 2024.
  * [9] S. Aghapour, K. Ahmadi, M. Anastasova, M. M. Kermani, and R. Azarderakhsh, “PUF-Kyber: Design of a PUF-based Kyber architecture benchmarked on diverse ARM processor,” _TechRxiv_ , 2023\.
  * [10] P. Ravi, B. Yang, S. Bhasin, F. Zhang, and A. Chattopadhyay, “Fiddling the twiddle constants-fault injection analysis of the number theoretic transform,” _IACR Transactions on Cryptographic Hardware and Embedded Systems_ , pp 447-481,__ 2023\.
  * [11] M. J. Kannwischer, J. Rijneveld, P. Schwabe, and K. Stoffelen, “pqm4: Testing and benchmarking NIST PQC on ARM Cortex-M4,” IACR, USA, Tech. Rep. 2019/844, 2019\.
  * [12] R. Avanzi, J. Bos, L. Ducas, E. Kiltz, T. Lepoint, V. Lyubashevsky, J. M. Schanck, P. Schwabe, G. Seiler, D. Stehle. Crystals-Kyber–Algorithm Specifications and Supporting Documentation, 2024\.
  * [13] N. Bindel, J. Buchmann, and J. Kramer, “Lattice-based signature schemes and their sensitivity to fault attacks,” in _Proc. Workshop Fault Diagnosis Tolerance Cryptogr. (FDTC)_ , pp. 63-77, 2016.
  * [14] L. G. Bruinderink and P. Pessl, “Differential fault attacks on deterministic lattice signatures,” _IACR Transactions on Cryptographic Hardware and Embedded Systems,_ vol. 2018, no. 3, pp. 21-43, 2018.
  * [15] K. Mus, S. Islam, and B. Sunar, “QuantumHammer: A practical hybrid attack on the LUOV signature scheme,” in _Proc. ACM SIGSAC Conference on Computer and Communications Security_ , 2020.
  * [16] P. Ravi, M. Jhanwar, J. Howe, A. Chattopadhyay, and S. Bhasin, “Exploiting determinism in lattice-based signatures: Practical fault attacks on pqm4 implementations of NIST candidates,” in _Proc. AsiaCCS_ , pp. 427-440, 2019\.
  * [17] K. Ahmadi, S. Aghapour, M. M. Kermani, and R. Azarderakhsh, “Efficient error detection schemes for ECSM window method benchmarked on FPGAs,” IEEE Trans. Very Large Scale Integr. (VLSI) Syst,_to appear 2024_.
  * [18] A. Dominguez-Oviedo and M. A. Hasan, “Algorithm-level error detection for ECSM,” _Centre Appl. Crypto. Res._ , Univ. Waterloo, ON, Canada, Tech. Rep., TR-2009-05, 2009.
  * [19] R. Singh, S. Islam, B. Sunar, and P. Schaumont, “Analysis of EM fault injection on bit-sliced number theoretic transform software in Dilithium,” _ACM Transactions on Embedded Computing Systems,_ to appear 2024 _._
  * [20] L. Ducas, E. Kiltz, T. Lepoint, V. Lyubashevsky, P. Schwabe, G. Seiler, and D. Stehle, “CRYSTALS-Dilithium: A lattice-based digital signature scheme,” _IACR Transactions on Cryptographic Hardware and Embedded Systems,_ vol. 2018, pp. 238-268, 2018.
  * [21] A. Sarker, M. M. Kermani, and R. Azarderakhsh, “Hardware constructions for error detection of number-theoretic transform utilized in secure cryptographic architectures,” _IEEE Trans. Very Large Scale Integr. (VLSI) Syst._ , vol. 27, no. 3, pp. 738-741, 2018.
  * [22] A. Sarker, A. C. Canto, M. M. Kermani, and R. Azarderakhsh, “Error detection architectures for hardware/software co-design approaches of number-theoretic transform,” _IEEE Trans. Comput.-Aided Des. Integr. Circuits Syst._ , vol. 42, no. 7, pp. 2418-2422, 2022.
  * [23] V. Lyubashevsky, D. Micciancio, C. Peikert, and A. Rosen, “SWIFFT: A modest proposal for FFT hashing,” in _Proc. Fast Software Encryption._ Berlin, Germany: Springer, pp. 54-72, 2008.
  * [24] N. Zhang, Q. Qin, C. Zhou, S. Yin, S. J. Wei, and L. Liu, “NTTU: An area-efficient low-power NTTuncoupled architecture for NTT-based multiplication,” _IEEE Trans. Comput.,_ vol. 69, no. 4, pp. 520-533, 2020.
  * [25] J. Y. Jou and J. A. Abraham, “Fault-tolerant FFT networks,” _IEEE Trans. Comput.,_ vol. 37, no. 5, pp. 548-561, May 1988.
  * [26] L. Ducas, A. Durmus, T. Lepoint, and V. Lyubashevsky, “Lattice signatures and bimodal Gaussians,” in _Proc. 33rd Annu. Cryptology Conf._ , pp. 40-56, 2013.


# Learning from Imperfect Demonstrations
via Adversarial Confidence Transfer

Zhangjie Cao∗, Zihan Wang∗, Dorsa Sadigh Emails<EMAIL_ADDRESS><EMAIL_ADDRESS><EMAIL_ADDRESS>Science, Stanford
University, CA, USA* means Equal Contribution. Author ordering determined by
coin flip over a Google Hangout.

###### Abstract

Existing learning from demonstration algorithms usually assume access to
expert demonstrations. However, this assumption is limiting in many real-world
applications since the collected demonstrations may be suboptimal or even
consist of failure cases. We therefore study the problem of learning from
imperfect demonstrations by learning a confidence predictor. Specifically, we
rely on demonstrations along with their confidence values from a different
_correspondent_ environment (source environment) to learn a confidence
predictor for the environment we aim to learn a policy in (target environment—
_where we only have unlabeled demonstrations_). We learn a common latent space
through adversarial distribution matching of multi-length partial trajectories
to enable the transfer of confidence across source and target environments.
The learned confidence reweights the demonstrations to enable learning more
from informative demonstrations and discarding the irrelevant ones. Our
experiments in three simulated environments and a real robot reaching task
demonstrate that our approach learns a policy with the highest expected
return. We show the videos of our experiments on our website.

##  I Introduction

Standard imitation learning algorithms assume that collected demonstrations
are provided by experts, who can always achieve successful outcomes [1, 2, 3].
However, this assumption is usually violated in real-world applications.
Expert demonstrations require domain expertise and thus are expensive to
obtain. On the other hand, we often have easy access to a plethora of
imperfect demonstrations ranging from expert behavior to noise. This is
evident when collecting human demonstrations on a robot arm since controlling
manipulators with high degrees of freedom can be challenging for human users
[4], or a bounded rational demonstrator may have difficulty providing optimal
demonstrations [5, 6].

Prior work on learning from imperfect demonstrations can be roughly
categorized into confidence-based and ranking-based methods: Confidence-based
methods learn a confidence over demonstrations to re-weight them [7, 8], while
ranking-based methods recover the reward function [9, 10, 11]. However, all
these works require some level of annotation, e.g., confidence or rankings of
the demonstrations. For complex environments, large-scale annotations are
required to cover the whole state-action space, which is labor-intensive and
expensive to obtain. Furthermore, the annotations can often contain a
significant amount of noise when collected through crowdsourcing platforms.

In this paper, we focus on confidence-based methods. As shown in Fig. 1, our
insight is that instead of using manually annotated confidence over
demonstrations, we leverage confidence annotations in a different but related
environment—the _source_ environment. The source environment may have a
different state-action space from the _target_ environment—the environment we
are operating in. The source environment and the target environment should
have a correspondence mapping between state-action pairs [12]. We assume that
collecting offline demonstrations along with their confidence annotations is
easy in the source environment. As an example, the source and target
environments can be a simulated robot and a real robot respectively or two
different robot arms with different degrees of freedom. Our key idea is to
leverage the correspondence between state-action pairs in the source and
target environments to learn a confidence measure without the restrictive
assumption of access to confidence annotations of target demonstrations.

Figure 1: Overview of the problem setting.

This problem setting introduces a new challenge: The source and target
environments have different state-action spaces. That means the dimensions of
the states and actions and even the semantic meaning of each dimension can be
different between the two environments, so we cannot directly apply the source
confidence predictor to the target state-action pairs. We instead need to
transfer the source confidence predictor to the target environment.

We propose an approach that leverages the assumption of the correspondence
between state-action pairs in the source and target environments leading to
the existence of a common latent space, where one can learn a shared
confidence predictor for both environments. Our approach aligns both the
feature-level distributions and the confidence-level distributions of partial
trajectories with different lengths across domains by a domain adversarial
loss. We design a multi-length partial trajectory matching, which preserves
temporal relationships across consecutive states and actions, and enables an
accurate matching of trajectories.

Our contribution is a new confidence-based imitation learning algorithm that
learns from imperfect demonstrations. Instead of relying on manual annotations
in our environment, we leverage confidence-annotated demonstrations in a
related environment—where collecting this data is easier—to learn a confidence
predictor for our environment of interest. We experiment with several
simulation environments and a real robot reaching task. Our results suggest
that our approach achieves both more accurate confidence prediction and higher
performance of the imitation policy than other methods.

##  II Related Work

Imitation Learning. Imitation learning aims to learn a policy by imitating the
behavior of expert demonstrations, where mainstream methods can be categorized
into behavior cloning (BC), inverse reinforcement learning (IRL), and
generative adversarial imitation learning (GAIL). BC directly learns the
policy mapping state to action by supervised learning over state-action pair
data in the demonstrations [13, 14, 15]. IRL first recovers the reward
function from demonstrations and learns a policy using reinforcement learning
based on the learned reward function [16, 1, 2]. GAIL aligns the occupancy
measure between the policy and the demonstrations with adversarial learning
[3, 17, 18]. These standard imitation learning methods usually assume that all
of the demonstrations are from experts who successfully finish the task. This
work goes beyond this assumption by allowing for learning from suboptimal or
even malicious demonstrations.

Imitation Learning from Imperfect Demonstrations. Imitation learning from
imperfect demonstrations aims to recover a well-performing policy from
demonstrations that can range from random noise to optimal. Ranking-based and
confidence-based methods are the two core approaches to address this problem.
Ranking-based methods recover the reward function from provided ground-truth
[9, 19], self-generated [10, 11], or interactively collected rankings [20,
21]. In our setting, we can only generate rankings and learn the reward in the
source environment, which cannot be applied to the target environment with a
different state space. Target rankings can be self-generated by injecting
noise in the target demonstrations as in [10], but that would require the
self-generated demonstrations to be ranked worse than the original ones. This
may not hold in our setting with the potential of nearly random demonstrations
in our dataset.

Confidence-based methods typically learn the confidence predictor from
confidence annotations [7, 8] but the annotations are assumed to be on target
demonstrations. We go beyond this assumption in our setting, since the
confidence predictors are from a different environment. Cao et al. learn the
target confidence predictor from the reward of demonstrations in the source
environment but still assume the two environments share the same state space
[22]. We target a more general setting, where the source and target
environments have different state-action spaces and dynamics.

Other than learning from suboptimal demonstrations, prior works also consider
learning from demonstrations with different dynamics and develop feasibility
metrics that learn from this data [22, 23]. However, unlike this paper, these
works require the same state space across the two dynamics.

Correspondence of Dynamics. To learn the correspondence [24] between different
dynamics, there are mainly two categories of methods: learning a translation
mapping [25], or a common latent space [26, 27] for states and actions. To
learn a translation mapping, Tzeng et al. [28] weakly align pairs of images
with paired image observations. Ammar el al. [29] utilize unsupervised
manifold alignment to find correspondence between states of demonstrations.
Kim et al. [30] and Zhang et al. [12] recently translated states and actions
across environments by learning a unidirectional mapping of states and a cycle
of forward mapping and inverse mapping between actions. Learning a translation
mapping is often much more challenging than simply finding a common latent
space. A common latent space is sufficient for the transfer of a confidence
predictor across different dynamics.

To learn a common latent space, domain randomization methods align the
different dynamics by randomizing the dynamics with hand-crafted augmentations
[31, 32, 33], but they require the differences between the two dynamics to be
covered by the augmentations. Gupta et al. [34] require pairs of states from
different dynamics to learn the latent space. The assumptions of paired data
or covering augmentations are often restrictive in practice. We remove these
assumptions and only assume the existence of correspondence and learn the
invariant representation space using multi-length partial trajectory matching.

##  III Problem Statement

Let the source and target environments be represented by two deterministic
Markov decision processes (MDP), $M^{\text{src}}$ and $M^{\text{tar}}$, where
$\mathcal{M}^{\bullet}=\langle\mathcal{S}^{\bullet},\mathcal{A}^{\bullet},p^{\bullet},\mathcal{R}^{\bullet},\rho_{0}^{\bullet},\gamma^{\bullet}\rangle,\bullet\in\\{\text{src},\text{tar}\\}$.
$\bullet=\text{src}$ indicates the source domain, while $\bullet=\text{tar}$
indicates the target domain. $\mathcal{S}^{\bullet}$ is the state space,
$\mathcal{A}^{\bullet}$ is the action space,
$p^{\bullet}:\mathcal{S}^{\bullet}\times\mathcal{A}^{\bullet}\rightarrow\mathcal{S}^{\bullet}$
is the transition function, $\rho_{0}^{\bullet}$ is the initial state
distribution,
$\mathcal{R}^{\bullet}:\mathcal{S}^{\bullet}\times\mathcal{A}^{\bullet}\rightarrow\mathbb{R}$
is the reward function, and $\gamma^{\bullet}$ is the discount factor. A
policy for $M^{\text{tar}}$ is
$\pi^{\text{tar}}:\mathcal{S}^{\text{tar}}\times\mathcal{A}^{\text{tar}}\rightarrow[0,1]$
defining a probability distribution over the target action space in a given
state. To evaluate a policy $\pi^{\text{tar}}$, we use the expected return,
which is defined as
$\eta_{\pi^{\text{tar}}}=\mathbb{E}_{s_{0}\sim\rho_{0}^{\text{tar}},\pi^{\text{tar}}}\left[\sum_{t=0}^{\infty}(\gamma^{\text{tar}})^{t}\mathcal{R}^{\text{tar}}(s_{t},a_{t},s_{t+1})\right]$,
where $t$ indicates the time step. A trajectory $\xi^{\bullet}$ is drawn from
either the source or target environments, and is a sequence of state-action
pairs:
$\xi^{\bullet}=\\{s^{\bullet}_{0},a^{\bullet}_{0},s^{\bullet}_{1},\dots,a^{\bullet}_{T-1},s^{\bullet}_{T}\\}$.

Problem Definition. In our problem setting, we are given a set of imperfect
target demonstrations
$\Xi^{\text{tar}}=\\{\xi^{\text{tar}}_{1},\xi^{\text{tar}}_{2},\dots\\}$ in
the target environment $M^{\text{tar}}$, and a set of imperfect source
demonstrations
$\Xi^{\text{src}}=\\{\xi^{\text{src}}_{1},\xi^{\text{src}}_{2},\dots\\}$ along
with a measure of confidence $c^{\text{src}}$ in the source environment
$\mathcal{M}^{\text{src}}$, where the confidence function
$c^{\text{src}}:\mathcal{S}^{\text{src}}\times\mathcal{A}^{\text{src}}\rightarrow\mathbb{R}$
evaluates how useful the input state-action pairs are and if they are
available or easy to access. Our goal is to find a well-performing policy
$\pi^{\text{tar}}$ for the target environment. Note that we go beyond standard
imitation learning by allowing these demonstrations in both environments to
range from useless or even harmful demonstrations to nearly optimal or fully
optimal ones. At the high level, our approach is to learn a confidence
function
$c^{\text{tar}}:\mathcal{S}^{\text{tar}}\times\mathcal{A}^{\text{tar}}\rightarrow\mathbb{R}$,
which assigns a value to how confident we are in the optimality of the action
$a^{\text{tar}}_{t}$ in a particular state $s^{\text{tar}}_{t}$. We can then
reweight the target state-action pairs and conduct standard imitation learning
on the reweighted state-action pairs to learn a policy. To learn
$c^{\text{tar}}$, we would like to learn a confidence predictor
$f^{\text{tar}}$ for the target environment by leveraging knowledge of the
confidence function $c^{\text{src}}$ from the source environment.

Assumptions. To transfer the confidence predictor from the source to the
target environment, we assume that the source and target environments are
_correspondent_ , which means that there exists a relational mapping between
the state-action pairs of the source agent and the target agent [12].
Specifically, there exists a relation on states:
$\Phi:\mathcal{S}^{\text{src}}\times\mathcal{S}^{\text{tar}}$. There also
exists a mapping from source state-action pairs to the target action space,
$H_{1}:\mathcal{S}^{\text{src}}\times\mathcal{A}^{\text{src}}\rightarrow\mathcal{A}^{\text{tar}}$,
and another mapping from target state-action pairs to the source action space,
$H_{2}:\mathcal{S}^{\text{tar}}\times\mathcal{A}^{\text{tar}}\rightarrow\mathcal{A}^{\text{src}}$.
The mappings $\Phi$, $H_{1}$, and $H_{2}$ satisfy the following: if
$(s^{\text{src}},s^{\text{tar}})\in\Phi$, then $\forall
a^{\text{src}}\in\mathcal{A}^{\text{src}},(s_{\text{next}}^{\text{src}},s_{H_{1}}^{\text{tar}})\in\Phi$
and $\forall
a^{\text{tar}}\in\mathcal{A}^{\text{tar}},(s_{\text{next}}^{\text{tar}},s_{H_{2}}^{\text{src}})\in\Phi$.
Here, $s_{\text{next}}^{\text{src}}$ is the successor state from
$s^{\text{src}}$ when taking the action $a^{\text{src}}$, and
$s_{H_{1}}^{\text{tar}}$ is the successor state from state $s^{\text{tar}}$
when taking the action $H_{1}(s^{\text{src}},a^{\text{src}})$, and
$s_{\text{next}}^{\text{tar}}$ and $s_{H_{2}}^{\text{src}}$ are defined
similarly. The assumption builds a connection between the source and target
agents. Without this assumption, the two environments may have no relations to
each other, and there will not be any reason for knowledge transfer across
them. This assumption is often satisfied in many real-world applications.
Specifically, we are interested in two common settings: (1) The source
environment being a simulation of a real robot that would define the target
environment. For example, the source environment can be a simulated robotic
arm picking an apple from a bowl in simulation, while the target environment
can be a real robotic arm picking an apple from a bowl. (2) The source and
target environments are different robots performing the same task in a similar
context. For example, the source environment can be a simpler 3-DoF robot arm
reaching the center of a bowl without colliding with obstacles, while the
target robot can be a 7-DoF robot performing the same task. We note that the
correspondence assumption is less restrictive compared to assumptions in
related prior work such as assuming the same dynamics across the environments
[7, 10].

Figure 2: (left) The overview of the learning process consisting of three
stages: In the first stage, we train the source encoder $E^{\text{src}}$ and
the shared decoder with source confidence-labeled demonstrations
$s^{\text{src}},a^{\text{src}}$. In the second stage, we align the latent
distributions of source and target state-action pairs with a multi-length
partial trajectory matching, where we explain the second stage in detail in
the right figure. In the third stage, we use the learned $E^{\text{tar}}$
chained with $F$ as the target confidence predictor. (right) Multi-length
partial trajectory matching: we align the feature-level distributions (output
of encoders $E^{\text{src}}$ and $E^{\text{tar}}$) of different lengths’
partial trajectories with different discriminators $D_{1},D_{2},\dots,D_{n}$
and align the confidence-level distributions (output of the shared decoder
$F$) of different lengths’ partial trajectories with different discriminators
$D^{\prime}_{1},D^{\prime}_{2},\dots,D^{\prime}_{n}$.

##  IV Adversarial Confidence Transfer

We aim to leverage the confidence annotation in the source environment and
transfer the confidence knowledge to the target environment. Our key insight
is to map the source and target state-action pairs to a common latent space
$\mathcal{Z}$, and enforce similarity between the latent features
($z^{\text{src}}$ and $z^{\text{tar}}$) for correspondent state-action pairs.
Here, $(s^{\text{src}},a^{\text{src}})$ and $(s^{\text{tar}},a^{\text{tar}})$
satisfy $(s^{\text{src}},s^{\text{tar}})\in\Phi$ and
$(s_{\text{next}}^{\text{src}},s_{\text{next}}^{\text{tar}})\in\Phi$. We then
build a shared decoder to predict the confidence from the latent features.

Since the source and target environments have different state-action spaces,
we learn two different encoders
$E^{\text{src}}:\mathcal{S}^{\text{src}}\times\mathcal{A}^{\text{src}}\rightarrow\mathcal{Z}$
and
$E^{\text{tar}}:\mathcal{S}^{\text{tar}}\times\mathcal{A}^{\text{tar}}\rightarrow\mathcal{Z}$
to map the source and target state-action pairs into a common latent space
$\mathcal{Z}$. We then learn a shared decoder
$F:\mathcal{Z}\rightarrow\mathbb{R}$ to predict the confidence from latent
features with confidence-labeled demonstrations in the source environment. The
target confidence predictor can be computed by $F\circ E^{\text{tar}}$. Now,
the challenge is how to learn the encoders to ensure that the common latent
space satisfies the above requirements.

###  IV-A Multi-length Partial Trajectory Matching

To map correspondent state-action pairs to a similar latent feature, we
develop a distribution matching objective to align the latent distribution of
source and target state-action pairs. We can directly align the distribution
of state-action latent features, i.e., matching
$E^{\text{src}}(s^{\text{src}}_{t},a^{\text{src}}_{t})$ and
$E^{\text{tar}}(s^{\text{tar}}_{t},a^{\text{tar}}_{t})$. However, this does
not capture the temporal dependency between state-action pairs over a
trajectory. To address this issue, we incorporate a prior that captures this
temporal dependency when learning the encoders. Imagine a setting where a
simulated robot and a real robot are both making a burger. The state-action
pairs corresponding to placing the ingredients (e.g., the patty, lettuce, and
tomatoes) are always after placing the bottom bun. If we only match the state-
action pairs one at a time between the simulated and the real robot, placing
the bottom bun may match placing any of the ingredients or even the top bun,
but if we also match the consecutive steps of the state-action pairs, placing
the bottom bun in both environments will more likely capture the information
that the bottom bun comes before the rest of the ingredients. So, as shown in
Fig. 2 (right), we propose a multi-length partial trajectory alignment that
emphasizes learning the temporal relationship between state-action pairs.
Specifically, we match the latent feature distribution of length $k$,
$(k=1,2,\dots,K)$ partial trajectories, which preserves the temporal
relationship between consecutive states and actions, and make the latent
features of correspondent state-action pairs more likely to be aligned.

###  IV-B Feature-Level and Confidence-Level Matching

We can now develop our learning algorithm that aligns the latent distribution
of state-action pairs. We first train the source encoder $E^{\text{src}}$ and
the shared decoder $F$ with source confidence-labeled state-action pairs using
a regression loss:

$\displaystyle\mathcal{L}_{\text{reg}}=\mathbb{E}_{(s^{\text{src}},a^{\text{src}})\sim\xi^{\text{src}},\xi^{\text{src}}\sim\Xi^{\text{src}}}\left[\lVert
F(E^{\text{src}}(s^{\text{src}},a^{\text{src}}))-c^{\text{src}}(s^{\text{src}},a^{\text{src}})\rVert\right].$
(1)

After learning $E^{\text{src}}$ and $F$, we fix the parameters of
$E^{\text{src}}$, and $F$ and use the latent feature space of $E^{\text{src}}$
as the common feature space. We then only need to make the target encoder
$E^{\text{tar}}$ encode the target state-action pairs into this shared feature
space and match the source latent feature distribution.

We learn the latent feature distribution alignment using a generative
adversarial network. Since we only want the partial trajectories of the same
length to be matched, for each length of partial trajectories $k$, we adopt a
discriminator $D_{k}$ and a loss $\mathcal{L}^{k}_{\text{fea}}$ to match the
latent distribution of length-$k$ partial trajectories. Specifically, the
discriminator $D_{k}$ aims to discriminate the latent features of source and
target length-$k$ partial trajectories, while the target encoder
$E^{\text{tar}}$ is trained to confuse the discriminator. The loss for the
discriminator $D_{k}$ can be derived as a binary classification loss:

$\displaystyle\mathcal{L}^{k}_{\text{fea}}=$
$\displaystyle-\mathbb{E}_{(s_{1},a_{1},\dots,s_{k},a_{k})\sim\xi,(\xi,\dot{)}\sim\Xi^{\text{src}}}$
(2)
$\displaystyle\left[\log\left(D_{k}(\text{Cat}\left(E^{\text{src}}(s_{1},a_{1}),\dots,E^{\text{src}}(s_{k},a_{k})\right)\right)\right]$
$\displaystyle-\mathbb{E}_{(s_{1},a_{1},\cdots,s_{k},a_{k})\sim\xi,\xi\sim\Xi^{\text{tar}}}$
$\displaystyle\left[\log\left(1-D_{k}(\text{Cat}\left(E^{\text{tar}}(s_{1},a_{1}),\dots,E^{\text{tar}}(s_{k},a_{k})\right)\right)\right].$

We concatenate the encoded features of multiple consecutive state-action pairs
as the latent feature for the length-$k$ partial trajectories. The
discriminator $D_{k}$ is trained to minimize the discriminator loss
$\mathcal{L}^{k}_{\text{fea}}$ and the target encoder $E^{\text{tar}}$ is
trained to maximize the discriminator loss.

However, the above losses only match the latent feature distributions of
state-action pairs, which is confidence-agnostic. So there may still exist
mismatches of state-action pairs with different confidence values. As shown in
Fig. 2 (right), to address the issue, we introduce a confidence-level
distribution matching to match the predicted confidence of source and target
state-action pairs, which further guides the latent features to be matched in
a confidence-aware way. Specifically, we use another discriminator
$D^{\prime}_{k}$ for matching the length-$k$ confidence sequence of length-$k$
partial trajectories using a similar binary classification loss:

$\displaystyle\mathcal{L}^{k}_{\text{con}}=$
$\displaystyle-\mathbb{E}_{(s_{1},a_{1},\cdots,s_{k},a_{k})\sim\xi,(\xi,\dot{)}\sim\Xi^{\text{src}}}$
$\displaystyle\left[\log\left(D^{\prime}_{k}(\text{Cat}\left(F(E^{\text{src}}(s_{1},a_{1})),\dots,F(E^{\text{src}}(s_{k},a_{k}))\right)\right)\right]$
$\displaystyle-\mathbb{E}_{(s_{1},a_{1},\dots,s_{k},a_{k})\sim\xi,\xi\sim\Xi^{\text{tar}}}$
$\displaystyle\left[\log\left(1-D^{\prime}_{k}(\text{Cat}\left(F(E^{\text{tar}}(s_{1},a_{1})),\dots,F(E^{\text{tar}}(s_{k},a_{k}))\right)\right)\right].$

Similarly, $D^{\prime}_{k}$ is trained to minimize the discriminator loss
$\mathcal{L}^{k}_{\text{con}}$ and $E^{\text{tar}}$ is trained to maximize the
discriminator loss.

$\small\min_{E^{\text{src}},F}\mathcal{L}_{\text{reg}}$ (3)

$\small\min_{D_{k},D^{\prime}_{k}}\mathcal{L}^{k}_{\text{fea}}+\mathcal{L}^{k}_{\text{con}}(k\in[1,K])$
(4)

$\small\max_{E^{\text{tar}}}\sum_{k=1}^{K}\left(\mathcal{L}^{k}_{\text{fea}}+\lambda\mathcal{L}^{k}_{\text{con}}\right)$
(5)

To optimize the whole network, we first train the source encoder and the
shared decoder with the source confidence-labeled data as shown in Eqn. (3)
(the first stage in Fig. 2 (left)). We then iteratively train the feature-
level and confidence-level discriminators and the target encoder as shown in
Eqn. (4) and Eqn. (5), where $\lambda$ is the trade-off between
$\mathcal{L}_{\text{fea}}$ and $\mathcal{L}_{\text{con}}$ (the second stage in
Fig. 2 (left)). After convergence, we chain the target encoder
$E^{\text{tar}}$ and the confidence decoder $F$ to form the target confidence
predictor as shown under Conf. Score Inference in Fig. 2.

We can then predict the confidence of each state-action pair $(s,a)$ in the
target environment as $F(E^{\text{tar}}(s,a))$. With the confidence of each
state-action pair, when conducting imitation learning on the target
demonstrations, we use the confidence to re-weight each state-action pair in
the imitation loss as in [7]. Then the imitation learning policy can learn
more from useful demonstrations with higher confidence.

Relation to Existing Works. Our method relaxes assumptions in prior works [7,
22] and targets a more realistic problem. However, we argue that our work is
compatible with prior works. The common latent feature space in our work can
be regarded as a new space of state-action pairs for the source and target
environments, where the correspondence between state-action pairs is
preserved, especially w.r.t. confidence. With the common latent feature space,
we can easily extend prior works to the setting of different state spaces by
conducting the algorithm, that is originally performed on states, on the
latent feature. Our contribution is not only providing a new approach to
learning from imperfect demonstrations but also expanding the usage of
algorithms from prior work by introducing a new space to operate in.

##  V Experimental Results

We conduct experiments on two MuJoCo environments, a simulated and a real
Franka Panda arm, and have released the code here. We compare our method
(Ours) with baselines GAIL [3], Dynamics Cycle-Consistency (DCC) [12], and
variants of our method by gradually adding modules: starting from GAIL without
confidence weighting, we first add feature-level matching as Ours-Feature. We
then add the confidence-level matching as Ours-Confidence. Adding multi-length
partial trajectory matching to Ours-Confidence derives Ours. We show the
results of using the ground-truth confidence to reweight the target
demonstrations (Oracle). The details about our baseline implementations are at
our website.

For the Mujoco environments, we evaluate the average return of $100$ rollouts.
For both the simulated robot and sim-to-real environments, we evaluate the
average return of $500$ rollouts. In all the experiments, we compute the
results over $10$ runs and report the mean and the standard deviation. We
compute the p-values using the student’s t-test. The details on the process of
demonstration generation for all environments are shown on our website.

(a) 1-joint

(b) 2-joint

(c) 4-leg

(d) 5-leg

Figure 3: Illustration of source and target MuJoCo environments. The left two
figures are Reacher and the right two are Ant.

MuJoCo Environment. We create 4 different MuJoCo environments (see Fig. 3):
1-joint and 2-joint reacher, 4-leg and 5-leg ant. 1-joint reacher and 4-leg
ant are the source environments, and the 2-joint reacher and 5-leg ant are the
target environments. The task for the reacher is to reach the red point from
its initial configuration. The task for the ant is to move towards the right
horizontally as fast as possible. We train the optimal policy using the RL
algorithm TRPO [35]. For the Reacher environment, we select the random policy,
a partially-trained policy, and the optimal policy to collect demonstrations
with mixed confidence. For the Ant environment, we select the random policy,
three partially-trained policies, and the optimal policy to collect
demonstrations with mixed confidence.

The expected return for these experiments is shown in Table I. Ours
outperforms GAIL and DCC, and is comparable with the Oracle in both
environments. Also, Ours outperforms Ours-Confidence, which demonstrates the
efficacy of multi-step partial trajectory matching. Ours-Confidence
outperforms Ours-Feature and Ours-Feature outperforms GAIL, which demonstrate
the efficacy of confidence-level and feature-level matching respectively. For
2-joint reacher and 5-leg ant, the highest p-values comparing with baselines
(between Ours and DCC) are $0.147$ and $0.031$ (statistically significant)
respectively.

| Reacher | Ant
---|---|---
GAIL |  -47.13$\pm$5.67 |  507.85$\pm$338.80
DCC |  -41.11$\pm$5.35 |  1218.59$\pm$231.74
Ours-Feature |  -43.97$\pm$4.68 |  1059.02$\pm$280.95
Ours-Confidence |  -41.82$\pm$6.80 |  1263.04$\pm$343.02
Ours |  -39.53$\pm$4.32 |  1556.43$\pm$159.12
Oracle |  -28.73$\pm$3.97 |  1622.55$\pm$179.43

TABLE I: The expected return for MuJoCo environments.

Figure 4: The top row and the bottom row are the expected return and sample
trajectories for 2-joint reacher and 5-legged ant respectively. The light grey
line in the Ant environment is the positive x-axis, which is the direction the
ant is supposed to move toward.

Figure 4 and 4 show the learning curves of the algorithms. We observe that all
methods have similar convergence rates. Figure 4 and 4 show a rollout
generated by each of the policies. The rollout from the oracle policy reaches
the goal in the most effective manner following the shortest trajectory
length. Ours is the next closest to the Oracle, while the GAIL baseline fails
at reaching the goal position at times.

Figure 5: Illustration of different trajectories and the 5 DoF robot arm.

Method | OT | OSC
---|---|---
Return | Success | Return | Success
GAIL |  -902.2$\pm$152.2 |  42.8$\pm$11.9 |  -1065.8$\pm$105.8 |  21.2$\pm$8.6
DCC |  -495.6$\pm$46.5 |  76.4$\pm$3.4 |  -723.7$\pm$121.4 |  49.5$\pm$15.2
Ours-Feature |  -1198.4$\pm$80.4 |  31.2$\pm$7.5 |  -868.1$\pm$265.7 |  37.6$\pm$18.3
Ours-Confidence |  -781.9$\pm$12.7 |  59.6$\pm$0.8 |  -821.2$\pm$121.4 |  42.0$\pm$11.9
Ours |  -154.8$\pm$12.8 |  92.8$\pm$1.0 |  -702.2$\pm$104.4 |  52.4$\pm$7.9
Oracle |  -55.6$\pm$2.3 |  100.0$\pm$0.0 |  -613.9$\pm$159.4 |  59.6$\pm$13.2

TABLE II: The expected return and the success rate among $500$ trials (%) of
GAIL, DCC, Ours-Single, Ours w/o $\mathcal{L}_{\text{con}}$, Ours, and Oracle
for the two setups OR and OSC in the simulated robot experiments.

Simulated Robot. As shown in Fig. 5, we create a task to move the end-effector
of a 5 DoF Panda Franka Robot arm toward the center of a plate without
colliding with the boundaries. The reward function consists of the negative L2
distance to the center; collisions with the wall or failing to reach the table
within the time limit; and positive reward when the end-effector reaches the
plate.

The source robot arm has 7 DoF and the target has 5 DoF, where the joints
marked by red in Fig. 5 are disabled. We show different kinds of trajectories:
(green–optimal trajectory reaching the center of the plate), (blue–suboptimal
trajectories reaching the plate but not the center), (red–failure cases
colliding with the wall), (violet–another type of failure case failing to
reach the plate within the time limit). We create two settings: (1) (OT)
consists of the optimal (green) and out-of-time (violet) trajectories; and (2)
(OSC) consists of optimal (green), suboptimal (blue), and collision (red)
trajectories. In the OT setting, we use the joint and end-effector position as
the state, which serves as an easy setup since the end-effector position is
shared between the 7-DoF and the 5-DoF. In the OSC setting, we use the joint
position, velocity and torque as the state, which is difficult to align. For
both settings, we use the 3D end-effector position difference as the action
space.

Expected return and the success rate are shown in Table II. Ours outperforms
GAIL and DCC with a large margin, which indicates that the learned target
confidence predictor can assign higher confidence to useful demonstrations
than DCC, and GAIL. The order of the performance from high to low is Ours,
Ours-Confidence, Ours-Feature, and GAIL, which demonstrates the efficacy of
our multi-length partial trajectory matching, confidence-level matching and
feature-level matching respectively. The highest p-values when comparing with
baselines (between Ours and DCC) are $5.64\times 10^{-12}$ for the success
rate and $1.25\times 10^{-13}$ for the expected return in the OR setting, and
$0.006$ for the success rate and $0.031$ for the expected return in the OSC
setting, demonstrating statistical significance.

(a) Simulated Robot

(b) Real Robot

(c) Expected Return

(d) Success Rate

Figure 6: (a-b) Illustration of the simulated robot and the real robot arm
environments. Different colors indicate different areas to place the objects.
(c-d) Expected return and success rate.

Sim-to-Real Environment. In the sim-to-real environment, we use a simulated
and a real Franka Panda Arm as the source and the target respectively (both
with 7-DoF). The task is to move the cube to the upper layer of the shelf (see
Fig. 6). On the shelf, there is a large stack of books on the right and a
small stack on the left. The middle area is empty. We assign positive rewards
to success of placing the cube on the shelf, with rewards $200$, $300$, and
$100$ for placing on the left, the middle and the right respectively. We
penalize time by giving a $-1$ reward for each time step. The average number
of steps for all demonstrations is about $300$. If the arm fails to put the
cube within $1000$ steps, it receives reward $0$. We use the joint position,
velocity and the end effector position as the state space and joint forces as
the action space.

Fig. 6(c) and 6(d) show Ours outperforms baselines DCC and GAIL. The
performance from high to low is: Ours, Ours-Confidence, Ours-Feature, and
GAIL, which demonstrates that multi-length partial trajectory matching,
feature-level and confidence-level matching are all necessary to learn a
common latent space and an accurate target confidence predictor for the real
robot. The p-value between Ours and the highest baseline, DCC, is $0.0004$ for
the expected return and $0.0052$ for the success rate, demonstrating
statistical significance.

Varying Composition of Demonstrations. We conduct experiments by varying the
composition of source demonstrations but fixing the target demonstrations in
the two Mujoco environments to test the robustness of confidence predictor
under different demonstration mixtures. We do not change the target
demonstrations since that will also influence the imitation learning
performance and we cannot test the efficacy of the confidence predictor. In
the _Demonstration_ column in Table III, we show the number of demonstrations
that are collected from different policies from random to optimal in Reacher
(3 policies) and Ant (5 policies). From top to bottom, the optimality of
source and target demonstrations deviates more and more. We observe that the
performance drops with larger deviation but does not drop too much even on the
last row. This demonstrates that the proposed approach can work stably even
with confidence distribution shift between demonstrations.

Reacher | Ant
---|---
Demonstration | Expected Return | Demonstration | Expected Return
1-5-94 |  -36.69$\pm$5.65 | 48-49-97-5-1 |  1525.03$\pm$176.91
1-49-50 |  -37.38$\pm$10.11 | 48-73-73-5-1 |  1489.85$\pm$268.22
1-94-5 |  -37.96$\pm$6.61 | 48-97-49-5-1 |  1436.78$\pm$176.91
47-48-5 |  -37.61$\pm$5.50 | 72-73-49-5-1 |  1351.44$\pm$115.59
94-1-5 |  -39.49$\pm$2.56 | 97-48-49-5-1 |  1345.96$\pm$299.08

TABLE III: Expected return with respect to varying compositions of the source
demonstrations for Reacher and Ant. The numbers in the demonstration column
indicate the number of the demonstrations collected from random to optimal
policies (3 policies for Reacher and 5 policies for Ant).

##  VI Conclusion

Summary. We propose an algorithm to learn from imperfect demonstrations, where
the demonstrations can be suboptimal or even fail. We learn a confidence
predictor by leveraging confidence labels and demonstrations in a different
but correspondent environment. We show the policy learned by our method
outperforms other baselines in various environments.

Limitations and Future Work. We require the source and target environments to
be correspondent. In the future, we plan to relax this assumption by defining
weaker correspondence only between confidence instead of the strict definition
of correspondence on dynamics.

##  VII Acknowledgements

We would like to thank FLI grant RFP2-000, NSF Award Number 1849952 and
1941722, and ONR for their support.

## References

  * [1] P. Abbeel and A. Y. Ng, “Apprenticeship learning via inverse reinforcement learning,” in _ICML_ , 2004, pp. 1–8.
  * [2] B. D. Ziebart, A. L. Maas, J. A. Bagnell, and A. K. Dey, “Maximum entropy inverse reinforcement learning.” in _AAAI_ , 2008.
  * [3] J. Ho and S. Ermon, “Generative adversarial imitation learning,” in _NeurIPS_ , vol. 29, 2016.
  * [4] D. P. Losey, K. Srinivasan, A. Mandlekar, A. Garg, and D. Sadigh, “Controlling assistive robots with learned latent actions,” in _ICRA_ , 2020.
  * [5] H. A. Simon, _Models of bounded rationality: Empirically grounded economic reason_. MIT press, 1997.
  * [6] M. Kwon, E. Biyik, A. Talati, K. Bhasin, D. P. Losey, and D. Sadigh, “When humans aren’t optimal: Robots that collaborate with risk-aware humans,” in _HRI_ , 2020.
  * [7] Y.-H. Wu, N. Charoenphakdee, H. Bao, V. Tangkaratt, and M. Sugiyama, “Imitation learning from imperfect demonstration,” in _ICML_ , 2019.
  * [8] S. Zhang, Z. Cao, D. Sadigh, and Y. Sui, “Confidence-aware imitation learning from demonstrations with varying optimality,” in _Conference on Neural Information Processing Systems (NeurIPS)_ , 2021.
  * [9] D. Brown, W. Goo, P. Nagarajan, and S. Niekum, “Extrapolating beyond suboptimal demonstrations via inverse reinforcement learning from observations,” in _International Conference on Machine Learning_. PMLR, 2019, pp. 783–792.
  * [10] D. S. Brown, W. Goo, and S. Niekum, “Better-than-demonstrator imitation learning via automatically-ranked demonstrations,” in _Conference on Robot Learning_. PMLR, 2020, pp. 330–359.
  * [11] L. Chen, R. Paleja, and M. Gombolay, “Learning from suboptimal demonstration via self-supervised reward regression,” _arXiv preprint arXiv:2010.11723_ , 2020.
  * [12] Q. Zhang, T. Xiao, A. A. Efros, L. Pinto, and X. Wang, “Learning cross-domain correspondence for control with dynamics cycle-consistency,” in _International Conference on Learning Representations_ , 2021. [Online]. Available: https://openreview.net/forum?id=QIRlze3I6hX
  * [13] M. Bain and C. Sammut, “A framework for behavioural cloning.” in _Machine Intelligence 15_ , 1995.
  * [14] S. Ross and D. Bagnell, “Efficient reductions for imitation learning,” in _AISTATS_ , 2010.
  * [15] S. Ross, G. Gordon, and D. Bagnell, “A reduction of imitation learning and structured prediction to no-regret online learning,” in _AISTATS_ , 2011\.
  * [16] A. Y. Ng, S. J. Russell _et al._ , “Algorithms for inverse reinforcement learning.” in _ICML_ , 2000, pp. 663–670.
  * [17] J. Fu, K. Luo, and S. Levine, “Learning robust rewards with adverserial inverse reinforcement learning,” in _ICLR_ , 2018.
  * [18] H. Xiao, M. Herman, J. Wagner, S. Ziesche, J. Etesami, and T. H. Linh, “Wasserstein adversarial imitation learning,” _arXiv preprint arXiv:1906.08113_ , 2019.
  * [19] D. S. Brown and S. Niekum, “Deep bayesian reward learning from preferences,” _arXiv preprint arXiv:1912.04472_ , 2019.
  * [20] V. Myers, E. Biyik, N. Anari, and D. Sadigh, “Learning multimodal rewards from rankings,” in _Proceedings of the 5th Conference on Robot Learning (CoRL)_ , 2021.
  * [21] E. Novoseller, Y. Wei, Y. Sui, Y. Yue, and J. Burdick, “Dueling posterior sampling for preference-based reinforcement learning,” in _Conference on Uncertainty in Artificial Intelligence_. PMLR, 2020, pp. 1029–1038.
  * [22] Z. Cao and D. Sadigh, “Learning from imperfect demonstrations from agents with varying dynamics,” _IEEE Robotics and Automation Letters (RA-L)_ , 2021.
  * [23] Z. Cao, Y. Hao, M. Li, and D. Sadigh, “Learning feasibility to imitate demonstrators with different dynamics,” in _Proceedings of the 5th Conference on Robot Learning (CoRL)_ , 2021.
  * [24] C. L. Nehaniv, K. Dautenhahn _et al._ , “The correspondence problem,” _Imitation in animals and artifacts_ , vol. 41, 2002.
  * [25] M. E. Taylor, P. Stone, and Y. Liu, “Transfer learning via inter-task mappings for temporal difference learning.” _Journal of Machine Learning Research_ , vol. 8, no. 9, 2007.
  * [26] F. Sadeghi and S. Levine, “Cad2rl: Real single-image flight without a single real image,” _arXiv preprint arXiv:1611.04201_ , 2016.
  * [27] H. Bharadhwaj, Z. Wang, Y. Bengio, and L. Paull, “A data-efficient framework for training and sim-to-real transfer of navigation policies,” in _2019 International Conference on Robotics and Automation (ICRA)_ , 2019, pp. 782–788.
  * [28] E. Tzeng, C. Devin, J. Hoffman, C. Finn, P. Abbeel, S. Levine, K. Saenko, and T. Darrell, “Adapting deep visuomotor representations with weak pairwise constraints,” in _Algorithmic Foundations of Robotics XII_. Springer, 2020, pp. 688–703.
  * [29] H. B. Ammar, E. Eaton, P. Ruvolo, and M. Taylor, “Unsupervised cross-domain transfer in policy gradient reinforcement learning via manifold alignment,” in _Proceedings of the AAAI Conference on Artificial Intelligence_ , vol. 29, no. 1, 2015.
  * [30] K. H. Kim, Y. Gu, J. Song, S. Zhao, and S. Ermon, “Cross domain imitation learning,” 2019.
  * [31] J. Tobin, R. Fong, A. Ray, J. Schneider, W. Zaremba, and P. Abbeel, “Domain randomization for transferring deep neural networks from simulation to the real world,” in _2017 IEEE/RSJ international conference on intelligent robots and systems (IROS)_. IEEE, 2017, pp. 23–30.
  * [32] X. B. Peng, M. Andrychowicz, W. Zaremba, and P. Abbeel, “Sim-to-real transfer of robotic control with dynamics randomization,” in _2018 IEEE international conference on robotics and automation (ICRA)_. IEEE, 2018, pp. 3803–3810.
  * [33] S. Zakharov, W. Kehl, and S. Ilic, “Deceptionnet: Network-driven domain randomization,” in _Proceedings of the IEEE/CVF International Conference on Computer Vision_ , 2019, pp. 532–541.
  * [34] A. Gupta, C. Devin, Y. Liu, P. Abbeel, and S. Levine, “Learning invariant feature spaces to transfer skills with reinforcement learning,” _arXiv preprint arXiv:1703.02949_ , 2017.
  * [35] J. Schulman, S. Levine, P. Abbeel, M. Jordan, and P. Moritz, “Trust region policy optimization,” in _ICML_ , 2015.