Low-Rank Tensor-Network Encodings for Video-to-Action Behavioral Cloning
Authors: Brian Chen
Doruk Aksoy
David J Gorsich
Shravan Veerapaneni
Alex Gorodetsky
Venue: Accepted by TMLR ()

Abstract: We describe a tensor-network latent-space encoding approach for increasing the scalability of behavioral cloning of a video game player’s actions entirely from video streams of the gameplay. Specifically, we address challenges associated with the high computational requirements of traditional deep-learning based encoders such as convolutional variational autoencoders that prohibit their use in widely available hardware or for large scale data. Our approach uses tensor networks instead of deep variational autoencoders for this purpose, and it yields significant speedups with no loss of accuracy. Empirical results on ATARI games demonstrate that our approach leads to a speedup in the time it takes to encode data and train a predictor using the encodings (between 2.6× to 9.6× compared to autoencoders or variational autoencoders). Furthermore, the tensor train encoding can be efficiently trained on CPU as well, which leads to comparable or better training times than the autoencoder and variational autoencoder trained on GPU (0.9× to 5.4× faster). These results suggest significant possibilities in mitigating the need for cost and time-intensive hardware for training deep-learning architectures for behavioral cloning.

Published in Transactions on Machine Learning Research (04/2024)
Low-Rank Tensor-Network Encodings for Video-to-Action Behavioral Cloning
Brian Chen
chenbri@nuicchcom edu
Deportment of Mathematics University of Michigan
Doruk Aksoy"
dorukGeumich. edu
Department of Acrospace Engineering University of Michigan
David J. Gorsich
mil
Ground Vehicle Systems Center
U.S. Army
Shravan Veerapaneni
shraean@umichh.. edu
Department of Mathematics Universily of Michigan
Alex A. Gorodetsky
goroda wrrdda&uuicch.om edu
Department of Acrospace Engineering University of Michigan
Reviewed OpenReview: https: openreeveee net/ pormm? 2d=
Abstract
We describe tensor-network latent-space encoding approach for increasing the scalabil- ity of behavioral cloning of video game player's actions entirely from video streams of the gameplay. Specifically, we address challenges associated with the high computational requirements of traditional deep-learning based encoders such as convolutional variational autoencoders that prohibit their use widely available hardware for large scale data. Our approach uses tensor networks instead of deep variational autoencoders for this purpose, and yields significant. speedups with loss of accuracy. Empirical results ATART games demonstrate that our approach leads to speedup in the time it takes to encode data and train a predictor using the encodings (between 2.6x to 9.6x compared to autocncoders or variational autoencoders). Furthermore, the lensor train encoding can be efficiently trained on CPU to achieve comparable or better training times than the autoencoder and variational autoencoder trained on GPU (0.9x to 5.4x faster). These results suggest significant possi- bilities in mitigating the need for cost and time-intensive hardware for training deep-learning architectures for behavioral cloning.
1 Introduction
Learning by imitation is a paradigm that enables teaching complex lasks through observation of expert behavior (Hussein et al., 2017) including object manipulation (Shen et al., 2022), robotic locomotion (Wang
Equal contribution DISTRIBUTION A. Approved for public release; distribution unlimited. OPSEC#7737

---

Published in Transactions on Machine Learning Research (04/2024)
et al., 2022), and autonomous driving (Bronstein et al., 2022). Several approaches for learning from data have been proposed, including behavioral cloning and inverse reinforcement learning.
In this work, we consider the task of learning directly from expert demonstrations in the form of action and rideo-frame sequences. In this context, learning directly from raw pixel data is computationally challenging due to the high resolution of game frames. This challenge is often addressed via downsampling and usage of large-scale computational resources. An illustrative example is the OpenAl VPT (Baker al., 2022) for the game of Minecraft, which first downsamples the original frame size of 640 x 360 to 128 128. Furthermore, the model is trained using 70, 000 hours of gameplay, has half billion parameters, and is trained for nine days on 720 V100 GPUs.
One approach to lackle the computational challenges of using raw video frames as input is to provide hand- tuned features (such as movement, location, terrain, visibility, etc.) rather than, or in addition to, raw pixel data. With this additional information, important elements of the game are provided and thus do not have to be learned solely from video data. This is the approach taken by some reinforcement learning and imitation learning works, such as Berner et al. (2019) for Dota 2 and Vinyals et al. (2019) for Starcraft 1I. However, such an approach requires human effort and insight in identifying traits that would be useful for a learning algorithm and developing a means for extracting this data. As an alternative to defining hand-crafted features, there have been various approaches proposed for learning a latent space representation via self- supervised learning, where auxiliary tasks are used to train a neural network that maps high nigh-dimensional input data into a lower dimensional space (Chen et al., 2021b). In this setting, autoencoder-based deep learning architectures have been a common method of choice for identifying this latent space (Jaques et al., 2021; Brown et al., 2020; Chen et al., 2021a;b; Amini et al., 2018), although other approaches have also been explored (Sermanet et al., 2017; Pari et al., 2021; Chen et al., 2021b).
In this work, we explore latent space extraction via more computationally efficient tensor network approach, based on the Tensor Train (TT) format (Oseledets, 2011). The TT-format offers a multi-lincar approximation that can be tuned to have a target reconstruction accuracy without the need of extensive hyperparameter tuning. Furthermore, the recent. development of an incremental algorithm, TT-ICE" (Aksoy et al., 2024), allows us to compute TT-decomposition of a large amount of video data faster and with less memory than the standard TT-SVD algorithm TTTdecomposition entirely using basic linear algebra operations such as malrix nultiplication and decompo- sition. Furthermore, it only requires single pass through the data and provides compression guarantees, a significant advantage over backpropagation techniques that can require multiple passes through the training data and do not guarantee any representation accuracies. The lack of guarantees on representation accuracy by deep-learning architectures also leads to the need for validation-data driven hyperparameter tuning to find hyperparameters and architecture settings that lead to high-quality representations. Instead, the com- plexity of the TT representation, represented by the tensor rank, can be found adaptively to guarantee target accuracy on the training data.
critical cnabling innovation for this application. TT-ICE* computes the
While the fo-mrrma may be viewed as having a smaller reprosentational capacity compared to the nonlin- ear architecture of neural networks, there is evidence suggesting that neural networks might exhibit more expressiveness than required for specific tasks. Several works have found that some layers of trained neural networks can be replaced with tensor networks with little to no degradation of accuracy while improving the computational and memory costs of running these networks (Novikov et al., 2015; Garipov et al., 2016; Yu et al., 2017b; Kim al., 2016; Lebedev et al., 2015; Sharma et al., 2021). In our work, we test whether encodings provided by neural networks can similarly be replaced by tensor networks. However, in contrast to the above works, we train the tensor network encodings directly, instead of first training neural network encoding then replacing it (or components of it)) with tensor networks.
In this paper, we demonstrate that this approach does provide a practical and highly effective alternative to deep-learning based encodings by reducing the requirements of hyperparameter tuning, thereby offering shorter training times for behavioral cloning. Moreover, the TT-encoding-based predictors achieve similar or higher behavioral cloning accuracics than predictors built encodings from convolutional autocncoder (AE) or variational autoencoder (VAE) (Kingma & Welling, 2014). Although other neural network ap-
2

---

Published in Transactions on Machine Learning Research (04/2024)
proaches for feature extraction are available (see Chen et al. (2021b) for examples), we chose AEs and VAEs since like the TT, these are trained to minimize a loss that is related to reconstruction error. To summarize, our contributions are as follows:
1. Empirical demonstration that the low dimensional emboddings obtained through TT-ICE" can be used for the downstream task of behavioral cloning.
Empirical demonstration a reduction in training time (2.6 - 9.6x faster when trained on GPU), while achieving comparable behavioral cloning accuracy and gameplay scores compared to butocnecoeerbased approaches for extracting a latent space.
2
Related Work
In machine learning literature, AEs and VAEs are commonly used as feature extraction tools. Polic et al. (2019) use a convolutional autoencoder (CAF) to extract. features from images of a factile sensor. The extracted features are then used to estimate contact object shape, odge position, oricntation and indentation depth using shallow neural networks. Yu & Zhou (2020) extract features from gearbox vibration data using CAE and uses the extracted features to identify fault modes. Pintelas et al. (2021) CAE to compress and denoise high dimensional images to smaller size and then use the encoded images for classification, Zabalza et al. (2016) use a stacked AE to extract features from hyperspectral images and use those features in classification, and Cai et al. (2022) use AFFexxraacted-featuree to determine the level of traffic congestion from video frames.
There is also increasod attention to transformer based methods. Yang et al. (2022) use a transformer encoder to generate pair-wise relations for the objects in the image and a transformer decoder to generate captions. Kim et al. (2021) use a transformer based architecture robot manipulation. However, training complex architectures like transformers increases the computational cost and data requirements even further. In Kim et al. (2021), it is reported that 8 NVIDIA V100 GPUs were used in training.
feature extraction step in dual arm
In addition to neural network-based tools, tensor methods have also been considered as a feature extraction option in machine learning literature. Tensor methods allow the extraction of meaningful features using lincar algebra operations, therefore climinating the need for backpropagation. Taherisadr et al. (2019) use CANDECMMM /PARAFAC (CP) format to reduce the size of EEG signals and increase the computational efficiency of a CNN-based seizure predictor. Zhang et al. (2017) use Tucker format to extract features from PCG signals and a support vector machine based classifier to identify anomalous heart sounds. Fonal & Zdunck (2019) use TT-format to extract features from various multidimensionall datasets and a k-ncarest neighbor classifier to perform multi-clas-elassiffaaaion Zhao et al. (2020) use a TT layer to extract features from high-dimensional video data before passing those features into a recurrent neural network for the task of video summarization.
Although tensor methods and neural networks are both capable of feature extraction, the question of whether the nonlinearity of neural networks can lead to features that are better suited for downstream tasks than the multilinear tensor methods remains. Previous studies indicate that the increased expressiveness of neural networks might be excessive for certain tasks and computationally expensive compared to similar tensor- basod methods. Several works have explored compressing portions of large neural networks after training, including fully connected layers (Novikov et al., 2015) and convolutional layers (Garipov al., 2016; Yu et al., 2017b; Kim et al., 2016; Lebedev et al., 2015; Sharma et al., 2021). The fine-tuning of large pre- trained neural networks can also be simplified by exploiting low rank structure; Hu et al. (2022) explore using low-rank matrices to update the weights of large language models when fine-tuning these models for downstream lasks. Other work has also explored using tensor operations directly within the network architecture itself during training as a means for parameter reduction and model compression; examples include Tensor Train Recurrent Neural Networks (Yang al., 2017), LR-TT-DNN (Qi al., 2023), Tensor Contraction Layers (Kossaifi et al., 2017), Tensor Regression Layers (Kossaifi et al., 2020), and Hidher-OOder Tensor RNNs (Yu et al., 2017a). The ability to replace portions of deep neural networks with
3

---

Published in Transactions on Machine Learning Research (04/2024)
tensor networks either during after training suggests that alternative representations may exist that do not require as much computational power and time to train, while not sacrificing accuracy. However, such works propose embedding tensor networks within neural networks and therefore require optimization and fine-tuning to capture the nonlincarity of the compressed layer(s) accurately. As such, they don't mitigate the issues surrounding the complexity of optimization approaches.
In this work, we explore using a tensor network in TT-format as an alternative to neural network based approaches, where the TT is used to extract features from video frame data for behavioral cloning. This approach is newly made possible by incremental algorithms (such TT-ICE (Aksoy et al., 2024)) with approximation guarantees, which provide the basis for memory efficient training of these models (in contrast to TT-SVD of Oseledets (2011), or other optimization- -baee incremental approaches). In comparison to using neural network based encodings like convolutional AEs and VAEs, we show that using a TT leads to significantly faster training times without sacrificing behavioral cloning accuracy.
3 Methodology
The overarching objective of this work is to extract a game-specific playing strategy from video demonstra- tions of an experienced agent. In this section we first describe our learning framework; then we present the methods to extract a low-dimensional representation and discuss how TT-decomposition fits within this setting. Finally, we briefly present the action predictor that maps latent representations actions.
The imitation learning framework use is shown in Figure 1. It consists of two parts: encoding trans- formation (labeled "Frame encoder" in Figure 1) and an action predictor ("Action predictor" in Figure 1). The encoding transformation serves to reduce the dimensionality of the space so that the action predictor can potentially be learned more efficiently. In this paper we aim to accelerate computation of the encoding transformation.
The encoder works by mapping frames of video sequence into an L dimensional latent-space. We compare three such encoders in this work: two neural network-based encoders determined from an AE and VAE, and a low-rank encoder based on the TTT-ecomposstion. The dimensionality of the latent space is typically determined by the formulation of the encoding transformation. For example, AEs and VAEs would treat this as a hyperparameter over which to optimize, whereas our low-rank approach will adaptively determine this dimension to a user-specified error tolerance. We further note that each of the three methods outlined above also have an ability to reconstruct an image from the latent space, which potentially can be useful for investigating the latent space.
Although there are other neural-network approaches for lcarning latent reprsentation, these mostly change the loss/ / auxiliary objective that is used to train the network. For instance, whereas AEs/VAEs are trained to minimize reconstruction error, contrastive methods instead learn to minimize distance between related samples (such as two distorted views of the same image (Chen et al., 2020)) and maximize a distance between unrelated samples. Because the main difference is the training objective, the training time for AEs AEs/VAEs Es should be similar to contrastive methods, assuming similar neural work architectures. We benchmark our results of using the TT against AEs/VAL' due to the similarity in their training objectives, which is to minimize reconstruction error.
In the rest of this section, we will discuss the methods and tools to compute the latent space representation of the gameplay sequences efficiently in detail.
3.1 Frame Encoder
This section discusses three alternatives for frame encoding. We will first briefly introduce AEs and VAEs and then discuss our TT-based approach in detail.
As mentioned in Section 1, AEs and VAEs have found great use for extracting encodings in the context of imitation learning and behavioral cloning. However, low-rank tensor decompositions have not been widely used for this purpose.
4

---

Published in Transactions on Machine Learning Research (04/2024)
TT AE VAE
Action(s)
Encoding
Stacked encodings
Frame encoder
Action
RGB images
predictor
Figure 1: Imiiation learning framework. Our focus is to enhance latent space learning depicted by the blue bor. We compare the TT, autoencoder (AE) and variational autoencoder (VAE) approaches. Orange cubes represent images from gomeplay videos. The extracted lalent features are then used in the action-prediction network to predict actions. The architecture of the action prediclor remains the same for all three encoding methods.
3.1.1 Autoencoder
AEs are neural network based latent space extraction tool. AEs used for variety of tasks, including dimensionality reduction (Wang et al., 2016), denoising (Majumdar, 2018), and image compression (Sreelak- shmi & Ravi, 2020). AEs consist of two parts: an encoder and a decoder. The encoder takes the input data and maps it to a. latent space representation. The decoder then takes the latent space representation and maps back to the original data. The goal the AE is to learn mapping that minimizes the difference between the original data and the reconstructed data.
AEs powerful tool for unsupervised learning, and they be used to extract meaningful features from data without any prior knowledge. However, ALs be computationally expensive to train, and they be sensitive to the choice of hyperparameters.
3.1.2 Variational Autoencoder
VAEs (Kingma & Welling, 2014) are similar to AEs in utilizing an encoding and decoding network. However, instead of mapping a fixed input to a fixed latent vector, the input is mapped probabilistically to a latent vector 2. A benefit of VAEs that is generative model, allowing one to generate new samples by sampling latent vectors (typically) from a standard normal distribution and passing it through the decoder. An additional potential benefit is that the training of a VAE imposes some structure onto the latent space, as the latent space is trained to be normally distributed.
The AE and VAE are both implemented using an encoder consisting of four convolutional layers, followed by a flattening layer and a final linear layer, and a decoder consisting of a linear layer, a reshaping layer, and four transposed convolutional layers. Additional details on the architecture and training the AEs and VAEs can be found in Appendix A.
3.1.3 Tensor- Train (TT) Format
The TT is compact representation method for multi-dinennsonal arrays. A d-dimensional tensor X E xx"" is said to be in TT-format if it is represented as product of d 3-dimensional ten- sors {g. G, E R called TT-cores. The internal dimensions of TT-cores G; are also referred to as TT-ranks and are determined by error-truncated SVDs through a user-specified relative error tolerance E [0, 1], except for rn = - 1.
In the case of RGB images, a video sequence is a 4-dimensional tensor E 3HxWXCXN with $ [0, 255]. The first three dimensions correspond the height (H), width (W), and color (C) channels; and the last dimension corresponds to the index of an individual image within the video sequence.
5

---

Published in Transactions on Machine Learning Research (04/2024)
In the video application considered in this paper, we first reshape X' into a 6-way tensor of shape 11 X x ng x N, where I HWC. By reshaping X to have smaller dimensions, we aim to limit the maximum TT-rank for each dimension. This also results in a slower growth in overall TT-ranks and yields higher compression. However, increasing the number of dimensions may result in slower computational speed because the TT-format is computed sequentially. In other words, increasing the number of dimensions will increase the number of TT-cores and therefore increase the number of SVDs needed to compute the TT-decomposition.
After reshaping, the tensor of video sequence is represented in TT-format
X ~ 91 3 92 aX1 ... 3x1 96,
(1)
where A 3x1 B denotes the tensor contraction between the third dimension of A and first dimension of B, and 9i x This representation often leads to a reduction in the size of the data. The original (but reshaped) data has storage complexity of O(n"), where n = m R4. The TT-format has storage complexity ((ndr", where r = (n....) Ti-
Note that in Equation 1, the first 5 TT-cores, f1....,9s. correspond to reshapings of the H, W, and C of the images these are common to all images. The final dimension indexes the unique frame and therefore only the final core is unique to a data point. To be more specific, for the i-th image stored in X, we have the reconstruction
(1) = G1 3X1 3X1 Is 3X1 Ge]t,
(2)
where i) denotes the i-th column of Ga- As a result of this, we can use the representations stored in 96 as encodings of the frames and the fifth TT-rank, rs, becomes the size of the latent space L.
One of the methods to compute TT-decomposttion for d-dimensional arrays the TT-SVD algorithm (Os- eledets, 2011, Alg. 1). This algorithm returns an approximation & that satisfies the inequality || X|| F S Sdee || for some predefined relative error upper bound Edea. Moreover, when computed with TT-SVD, the first 5 TT-cores consist of reshapings of orthonormal vectors (i.e., left singular vectors). However, perform- ing the TT-SVD on large lensors requires loading the full tensor into memory and performing a sequence of singular value decompositions. As a result, this approach requires a significant amount of memory. The memory of our existing computational resources limits us to a maximum number of around 30,000 frames for TV-SVD.
One way to reduce the memory and computational requirements of TT-SVD is adopting incremental approach. TT-ICE* (Aksoy et al., 2024, Alg. 3.2) is an incremental algorithu that computes the TT- decomposition for a stream of tensors with a given target accuracy. To train the tensor network, we create batches of B = 2500 frames from gameplay sequences and update the basis vectors stored in TT-cores using TT-ICE". Although TT-ICE" does not require a fixed batch size as shown in Aksoy et al. (2024), we train the tensor network using a fixed batch size to mimic the training of the NN based methods. Once a pass is completed over a batch of images, their latent representations are stacked to Go. Note that while TT-SVD computes the TT decomposition of n, x... X na x N tensor in time O(dna+1N) NN) (Huber et al., 2017), TT-ICE* computes& tensor train using batches over the final dimension. Letting B be the batch size and b be the total number of updates, TT-ICE" calculates the TT decomposition in time (bdnn B). If all batches are used to update the TT, this leads to the same computational complexity as TT-SVD (with bB = N). However, since TT-ICE" learns a generalizable space, not all B frames from a batch are needed to update the TT-cores. TT-ICE" updates the TT-cores by selecting the subset of frames that have an approximation error higher than the threshold. Furthermore, in the edge case where the TT-cores are highly representative of a batch, TT-ICE* skips updating TT-cores and directly computes the latent space embeddings using Algorithm 1. Through these modifications TT-ICE* achieves bB < N. Please refer to Table A.1. in Aksoy et al. (2024) for more details. Even for the worst case of bB = N, TT-ICE" offers a substantial reduction in memory needed to compute the ITT-decomposition in addition to eliminating the necessity of storing all frames at once.
Also note that TT-ICE" explores and increases the number of the basis vectors for cach dimension of the tensor stream adaptively based on the user-specified error tolerance. This removes the need mamually tune the size of the latent space I for the encoded images.
6

---

Published in Transactions on Machine Learning Research (04/2024)
In case problem has special properties to allow a-priori upper bound on TT-rank, the rank and complexity must be adapted in response to received data. The datasets used in this work do not readily come with useful upper-bounds on the TT-ranks. As a result, we set a relative error tolerance of 1% for each training batch and the TT-ranks are automatically determined by TT-ICE" (please refer to Theorem 2 and Theorem of A Assy et al. (2024) for proofs of approximation error)² In the asymptotic limit of the data, TT-ICE" will be able to represent any tensor because of its rank adaptation property.
Next, we discuss how out-of-sample gameplay frames mapped to the latent space by each encoding method.
3.2 Mapping unseen data to latent space
For AE/VAE, the latent space representation of unseen data is acquired directly through a forward pass of the trained encoder with the desired image. On the other hand for TT, the notion of forward pass replacod with sequential projection onto cores.
Note that the first 5 TT-crres contain reshaped orthonormal vectors when TT-cores are trained with TT-ICE". We obtain orthonormal vectors for each dimension by reshaping the first 5 TT-cores as
U reshape (G., (rl11
(3)
These orthonormal vectors are then used to encode any new images from the same game type. We obtain the latent embedding of images by simply projecting the appropriately reshaped images onto U,. The pseudocode of this procedure is provided in Algorithm 1 Note that as the ranks of the TT increase, the computational time of Algorithm 1 grows quadratically with r. In particular, the number of flops for the projection grows at a rate of O(drád-N N max(r, n)), where N is the number of new images to project, n = (1ax.e(1...... R4. and d is the dimensionality of the tensor. For the application considered in this paper, n = 16 and 5.
The AE, VAE, and TT all allow for reconstructing frames from the latent space. As such, we can visualize the reconstructed frames get a sense of what lost in encoding the frames. Figure shows an example of reconstructed frames from the latent. space for each encoding method.
Algorithm Mapping unseen data to latent space by orthogonal projections onto TT-cores
1: Input
2:
TT-ceres trained with IT-ICE* V/ E IR XXXXXXXXXX XN, j-th batch of frames that contain Nj frames to be projected
3:
Output
gj E grax
encodings of the j-th batch of frames
9j reshape (Y,, [R1, n2 m5,N,,
7: for = 1 to 4 do
8: U + reshape (G,, (m)11
Note that for i = 1 we have
9:
- UT g
10: gj f reshape (gjs
11: end for
12: reshape (9o, rr))
13: 9j UT 9j
3.3 Action predictor
To extract the game-specífic playing strategy, we employ a simple MLP based architecture. We stack the latent representation of the last frames together and obtain 4L vector as input the action predictor. Then the action prodictor outputs the predicted action for given timestep. Using the previous four frames
When the occupancy heuristic used, the approximation error guarantees proven TT-ICE apply TT-ICE* well. 2hhi algorithm adapted from Aksoy al. (2024)

---

Published in Transactions on Machine Learning Research (04/2024)
Tensor Train
Autoencoder
VAE
Original
Figure 2: Example reconstruction of Ms. Pac-Man frume using TT (wilh e = 0.011, AE and VAE. TT reconstruction is visually the most similar to the original frame. The size of the latent space for all three models is 11583, which 11.5% of the original size (210 x 160 x 3).
input allows the model to use some amount temporal information and common choice for experiments conducted on Atari games (sce, o-g.; Mnih al. (2015)).
The action predictor consists of 5 hidden layers with 50 neurons each. Every hidden layer is then followed by batch norm and ReLU operation. Finally, the output layer has size equal to the number of possible actions for cach game, which is interpreted as the probability of taking each action. Next, we present the computational experiments we have conducted.
4 Experiments
In this section, we describe our comparative study, the dataset we selected for the experiments, evaluation metrics we used, and results obtained.
4.1 Dataset
The dataset consists of game frames from demonstrations from set of Atari games, in particular: Beam- Rider, Breakout, MsPacman, Pong, Qbert, Seaquest, and Spacelnvadens", Our data is collected using trained Deep Q-Networks (Mnih et al., 2015), which are obtained from the RI-Baselines Zoo package (Raffin, 2018) and further trained using the Stable Baselines platform (Hill al., 2018) for 107 samples (states and ac- tions), using the tuned hyperparameters from RL-Basclines Zoo. The chosen Atari games were seleeted due to the availability of trained models and tuned hyperparameters in the RL-Baselines Zoo package. The reinforcement learning algorithm used to generate the model (Deep Q-Networks) was chosen both due to the availability of trained models and tuned hyperparameters as well as the fact that Deep Q-Networks have been shown perform well on Atari games (Mnih al., 2015).
Rather than using downsized grayscale 84 84 1 images (before stacking) as often done in the literature, instead use the original full-sized RGB images (210 160 3), since work deals directly with image compression and dimensionality reduction. For each game, we generate 200 trajectories, a subset of which is used as a training set and a subset of which is used as a validation set. More details about the split are given in Scction 4.2. Another 50 trajectorics are generated as a test set, which is used to evaluate the action prediction accuracy of the final trained models. We define a trajectory as a sequence of states (frames), rewards, and actions taken by the RL agent, starting the beginning of the game and ending when the RL agent loses life in the game or the end of the game is reached. The datasct can be found at Chen (2023).
We note that have also tested Enduro, but duc computational limitations for the AE/VAB results not presented. In the moderate data case, the AE, VAE took roughly hours per seed compared 4.9 hours for TT-GPU.
8

---

Published in Transactions on Machine Learning Research (04/2024)
4.2 Setup
The procedure for our experiments is follows:
1.
Train an AE, VAE, and TT on the frames from set of gameplay videos. The dimensionality of the latent space L for the TT is determined automatically by TT-ICE" After the TT is trained, an AE and VAE are trained with a latent space size fixed to be equal to the size L of the latent space from the TT. This is done to provide a fair comparison of the methods. The latent size L. for the various games are shown in Table 1. Furthermore, the VAE and AE are trained using a range of L2 regularization coefficients and initial learning rates.
2. For the AE and VAE, an additional portion of the gameplay videos are held out for validation. The validation reconstruction errors are used to select VAE and AE for encoding. Note that for TT-decomposition no such validation procedure is needed.
3. Using the encodings generated by the TT or the AE / VAE input, we train simple foed-forward networks to predict the actions taken by the trained reinforcement learning agent from which data is generated, using a range of hyperparameters.
4. For each of the AE, VAE, and TT encodings, we select the predictor that yields the lowest validation action prediction loss, as measured by cross-entropy loss on a separate validation set.
In our work, we consider training the AE, VAE, and TT separately to first learn an encoding, before training a predictor on the encodings dircctly. This is in contrast to end-to-end approaches, where the encodings would be trained jointly with the predictors. An advantage of training the encoders separately from the action predictors is reduced training time, the expensive encoders are not further finetuned for the particular task of behavioral cloning, A second advantage is that the encodings are smaller than the full frames, leading to lower memory and storage requirements. This approach mimics the approach taken in several reinforcement learning papers, which have found that decoupling training an encoder from training the reinforcement learning agent is an effective approach- examples include Pari et al. (2021); Parisi et al. (2022); Stooke et al. (2021); Yuan al. (2022); Shah & Kumar (2021).
The predictors and AE/ VAE were set up and trained using PyTorch and trained NVIDIA Tesla V100 GPUs. The TT is trained either using GPU (TT-GPU, NVIDIA Tesla V100) or CPU (TT-CPU, Intel Xcon Gold 6154 processors), using the TT-ICE* package timings are reported for both. Note that the only difference between TT-GPU and TT-CPU is the device that the computation took place. This is achieved by replacing NumPy functions (Harris et al., 2020) used for computations with their CuPy counterparts (Okuta al., 2017) for TT-GPU. In parallel to the findings in Aksoy et al. (2024), we have empirically found that this approach most effective when n; = 0(10) for = 1,...,5 and therefore reshaped the image batches to 15 14 x 16 10 X 3 x N, where N is the batch size used for the TT (in our experiments, N = 2500).
In this work, we consider two cases: 1) limited data regime where only 5 gameplay trajectories are used as training data for each game and 2) a moderate data regime where 160 gameplay trajectories are used as training data for each game. A summary of the number of trajectories used in the limited and moderate data settings shown Table In addition to these two scenarios, we also provide a more detailed investigation on the effects of changing the latent space size (Appendix B.2) and the amount of training data (Section 4.4.3) on action prediction accuracy and computational speedup.
Action predictors are trained with the Adam-W W optimizer (Loshchilov & Hutter, 2019) with a cosine learning rate decay (Loshchilov & Hutter, 2017) for fifty epochs in the low data case and five epochs in the moderate data case. These predictors are trained using the same range of initial learning rates and weight decay penalties as the ones used to train the AE and VAE encoders (shown in Appendix A). An additional 1 or 20 trajectorics (in the limited and moderate data regimes, respectively) not used during training are then used as a validation set to evaluate the trained predictors, with the predictors yielding the lowest gross-entropy action prediction loss selected.
9

---

Published in Transactions on Machine Learning Research (04/2024)
Table 1: Latent sizes for the AE, VAE, and TT for the games tested. Note that the latent size is determined automatically for the TT given the prescribed relative error bound of 0.01. That latent size then used for the AE and V.E.
Latent Size L
BeamRider
14645
Breakout
3561
MsPacman
11962
Pong
1287
Qbert
2267
Seaquest
15163
SpaceInvaders 17191
Table 2:
Summary of the number of trajectories used to train the models in the limited and moderate data cases. Note that each row corresponds to separate trajectories.
Limited Data
Moderate Data
Training Set for Training Encoder and Predictor Validation Set for Selecting Best Encoder (ONLY AE/VAE)
5
160
20
Validation Set for Selecting Best Predictor Test Set for Evaluating Final Models
20
50
50
We also include the results of predictor trained using the original, full-sized RGB frames, of dimension 210 x 160 x 3. Because the "no encoding" predictor uses the full-sized images, it is expected that these accuracies will be higher than the predictors using the encodings these are included as a comparison as an upper bound on accuracy we can expect.
Additional details on training and the neural network architectures can be found in Appendix A.
4.3 Metrics
We compare performance of the methods using four metrics: training time, reconstruction quality, action prediction accuracy, and normalized gameplay score. We then report the averaged results for each of those metrics for each regime in Section 4.4.
We compare the time required train each framework against the time needed to train predictor on full frames. The training time for each framework can further be broken down into the time it takes to train the encoder, the time it takes to encode the data, and the time it takes to train the action predictor given the encodings. These results are shown in Appendix B.1 as Figures 9 and 10.
Training Time.
Reconstruction Quality. We measure the reconstruction quality of the TT against AE and VAE using peak-signal-to-noise ratio (PSNR) and mean-squared error (MSE) as metrics.
Action prediction accuracy. We compare the accuracy of the actions suggested by each framework to actions taken by the RL agent. The output of the behavioral cloning agents is a vector of probabilities for each action we calculate the proportion of times the action with the highest probability is equal to the action with the highest probability outputted by the RL agent.
Normalized gameplay score. We normalize gameplay scores achieved by each framework so that a score corresponds to performance of an agent taking random actions and a score of corresponds performance of the RL agent. In particular, given the score s achieved by behavioral cloned agent, the mean score SaL achieved by the RL agent and the mean score 8random achieved by agent taking random
10

---

Published in Transactions on Machine Learning Research (04/2024)
actions, calculate the normalized score to be
8
Syandom
Enormalized
SRL Brandom
4.4 Results
4.4.1
Limited Data Regime
The benefits of bypassing the need for training a neural network encoder can be observed in the total training time of the different methods, as shown in Figure 3. In particular, the training time of the TT-basod predictor trained GPU (TT-GPU) is 8.2x to 9.6x faster compared to the AE and VAE-based predictors and 26. 1 to 29.7x faster compared to the predictor trained with no encoding. When the TT encoding is trained on CPU (but with the predictor still trained on GPC), the TT-predictor is still 3.3x to 5.4x faster than the AE and VAE predictors trained on GPU and 10.4x to 17.5x faster than the predictor trained with no encoding.
AE
VAE
TT-CPU
TT-GPU
Speedup vs Full Frames
30
10
25
20
15
5
10
5
0
BR
B
MP
P
Q
S
S1
BR
B
MP
P
S
S1
(a) 5 Trajectories
(b) 160 Trajectories
Figure 3: Speedup of the end-to-end training procedure of the AE, VAE, and TT-based predictors compared to predictor trained using the full RGB frames. Extracting latent space reduces end-to-end training time in all cases. In the low data regime, using the TT-GPU leads to a 26.1 16.1x to 29.7x faster training time than using the full RGB frames and a 8.2x to 9.6x speedup compared to the AE and VAE. In the moderate data case (160 trajectories), using TT-GPU leads to a 3.4x to 9.8x faster training time compored to using the full RGB frames and 2.6x 4.8x speedup compared the AE and VAF. Note that the difference. in speedups between the trajectory case and 160 trajectory case is due the difference training epochs used each case (50 in the 5 trajeclory case and in the 160 trajectory case); the speedups are similar when the wumaber of training epochs is held fixed while the number of training trajectories is changed, shown in Section 4.4.3. BR: Rider, B:Breakout, MPPMMsPacmman Pacman, P:Pong, Q:Qbert, S:Seaquest, SI:Space Invaders.
This speedup comes with an improvement in reconstruction quality, as shown in Table 3; the TT leads to lower MSE and higher PSNR compared to the AE and VAE six of the seven games. This speedup also comes at no cost in terms of behavioral cloning accuracy (Figure 4a) and achieved gameplay scores (Figure 5a), compared to training using encodings generated by AE VAE.
4.4.2 Moderate Data Regime
In the case where 160 trajectorics are used for training data, we find once again that the TT-based prodictor provides large speedup relative to the AE, VAE, and full frame predictors. In particular, shown in Figure 3, using the TT-encoder trained the no-encoding predictor. Compared to the AE and VAE, the end-to-end training time of the TT-based predictor is 2.6x to 4.8x faster. Training the TT on CPU still usually leads to speedup over the AE and VAE on GPU (0.9x to 2.9x faster), showing the computational efficiency of TT-ICE even on CPU compared to GPU-based deep learning.
GPU led to 3.4x to 9.8x faster training time compared to
These results come with an improvement in reconstruction quality and cost in terms behavioral cloning accuracy gameplay scores. The reconstructions generated by the TT have lower error compared to
11

---

Published in Transactions on Machine Learning Research (04/2024)
Table 3: Reconstruction quality of the TT, AE, and VAE the low data regime. The reported statistics are
MSE (using 255 as marimum pizel value) and PSNR of the reconstructions versus the original frames. For each gome, the bolded values are the smallest MSE and highest PSNR.
MSE
PSNR
TT
AE
VAE
TT
AE
VAE
BeamRider (BR)
6.3
30.5
47.6
46.8
35.3
33.0
Breakout (B)
33.0
67.6
80.9
39.4
32.9
30.3
MsPacman (MP)
27.1
58.0
103.3
38.7
32.5
28.9
Pong (P)
2.0
0.7
56.9
46.7
61.6
36.9
Qbert (Q)
6.8
33.5
68.5
45.6
34.3
30.8
Seaquest (S)
13.9
30.8
74.0
38.6
33.9
29.8
SpaceInvaders (SI)
80.3
144.5
176.4
32.4
28.4
26.2
AE
VAE
No encoding
100
Accuracy (%)
80
60
40
20
0
B
B
Q
(a) 5 Trajectories
(b) 160 Trajectories
Figure 4: Action prediction accuracies of trained predictors using encodings from the TT, AE, and VAE, as well as trained predictor using the full-sized original RGB frames. Training a predictor using TT encodings yields similar and somelimes beller prediction accuracies truining predictor using AE or VAE encodings. Error bounds are colculated as two times the standard error of gameplay scores calculated across ten (in the trajectory case) or three (in the 160 trajectory case) experiments. Please refer to the caption of Figure for the game abbreviations.
the reconstructions generated by the AE and VAE, shown in Table 4. Furthermore, find improved behavioral cloning accuracy of the predictor trained using the TT encodings versus the predictors trained on the AE or VAE encodings for some of the games, as seen in Figure 4b. The gameplay scores achieved by each the encoding methods is comparable, as shown in Figure 5b.
Note that in some games, the gameplay scores achieved by the predictors, including the one trained with encoding, much lower than the scores achieved by the RL agent. This is likely due to issue encountered in imitation learning where the distribution of data used to train the agent differs from the distribution of data encountered while testing the agent, since the agent's act directly affect the future states it encounters. This leads to compounding errors and hurts the performance of the imitation learning agent (Ross & Bagnell, 2010). These low scores are also consistent with Kanervisto et al. (2020), which found that behavioral cloning from human gameplay was unable match hurnan-level gameplay scores for the Atari games considered that paper.
4.4.3 Scalability of the Tensor Train
Differences the speedup between the limited data regime and moderate data regime are due primarily to the number training epochs used in training the AE/VAE. we use 50 training trajectorics in the limited data regime and 5 in the moderate data regime. The reason for this difference is that we found that more training epochs needed in the limited data regime to achieve similar accuracies via the AE/VAF.
12

---

Published in Transactions on Machine Learning Research (04/2024)
TT
AE
VAE
No encoding
Normallzell Score
1.2
0.8
0.4
-0.4
BR
B
MP
S
SI
BR
B
MP
SI
(a) Trajectories
(b) 160 Trajectories
Figure 5: Normalized gameplay scores achieved by each of the behawioral cloning models, trained using either encodings from the TT, AE, VAE, or the original RGB images. score of 0 corresponds to performance of an agent laking random actions and a score of 1 corresponds to performance of the RL agent. Training a predictor using TT encodings yields similor gameplay scores as training a predictor using AE or VAE encodings. Error bounds calculated two times the standard error of gomeplay scores calculated across ten (in the 5 trojectory case) or three (ön the 160 trajectory case) experiments.
Table 4: Reconstruction quality of the TT, AE, and VAE in moderate data regime. The reported statistics are the MSE (using 255 as marimum pizel value) and PSNR of the reconstructions versus the original frames. For each game, the bolded values are the smallest MSE and highest PSNR.
MSE
PSNR AE
TT
AE
VAE
TT
VAE
BeamRider (BR)
0.4
18.0
26.0
58.1
37.6
36.4
Breakout (B)
1.2
7.8
19.8
51.0
40.1
36.8
MsPacman (MP)
1.7
16.2
75.7
48.1
37.6
30.2
Pong (P)
1.8
0.0
6.9
47.0
71.4
40.7
Qbert (Q)
0.8
12.1
36.4
52.7
39.0
33.8
Seaquest (S)
1.0
8.4
52.9
49.5
39.9
31.4
SpaceInvaders (SI)
0.4
18.5
80.2
56.1
36.5
30.1
13

---

Published in Transactions on Machine Learning Research (04/2024)
small To make the speedups of the TT more comparable across different amounts of training data, study to examine the effect of increasing the number of gameplay trajectories used to train the behavioral cloning framework on both training time and prediction accuracy. The AE and VAE are trained using five epochs, matching the training sctup used in the moderate data regime. We ran this for the game of Seaquest, given the relatively low prediction accuracies we achieved compared to some of the other Alari games. As expected, Figure shows that the accuracy increases as more training data used, suggesting that better prediction accuracies can be achieved with more data.
We note that TT's speedup over the AE VAE and encoding remains similar more data is used, long as the number of training epochs is fixed. These results suggest that the timing advantage may remain for the TT even when more data is available for behavioral cloning.
AE
VAE
TT-CPU
TT-GPU
No encoding
100
Action Prediction Accuracy %
Speedup vs. No Encoding
80
4
60
40
2
20
40
80
120
160
40
80
120
Number of Trajectories (a) Speedup
Number of Trajectories (b) Action Prediction Accuracy
Figure 6: Results of the experiments investigating the effect of rumber of trajectories used to train the behav- ioral cloning framework on computational speedup and action prediction accuracy using different encoding methods for the game of Seaquest. As the number of training trajectories increases, action prediction accu- racy also increases. Computational benefits of TT-ICE" over the AE, VAE, no encoding remoin similar across different nwmber of trajectories. Action prediction accuracies are only presented for TT-CPU since TT-CPU and TT-GPU result in the same latent space.
4.4.4 Prediction Accuracy Versus Training Time
A significant difference between training the AE / VAF versus the TT is the number of passes through each data point. Since the amount of training data is limited, multiple passes through the data are used for the AE/VAE to increase the number of update steps. On the other hand, the TT is trained using an algorithm that requires only a single pass through the data. As a result, the AE/VAE process more data points during training than the TT, albeit with repeating data.
To address this difference, we run experiments where each algorithm is trained with one epoch and as such each algorithm processes the same number of data points. The resulting accuracies and timings are plotted in Figure 7. As can be seen, TT-GPU tends to achieve higher accuracies and lower end-to-end training time than the AE/VAE.
4.4.5 Encoding Times and Inference Speed
During inference, new frames are passed through the TT, AE, or VAE, before the resulting latent vector is passed into a predictor. Because the architecture of the prodictor is held fixed, differences inference speed will stem from differences in the time takes encode frames using the various methods. These encoding times compared in Table 5.
14

---

Published in Transactions on Machine Learning Research (04/2024)
AE
VAE
TT-CPU
TT-GPU
100
100
Action Prediction Accuracy %
No Encoding
80
80
60
60
No Encoding
40
40
20
20
0
0
0
2,000
4,000
6,000
0
2,000
4,000
6,000
8,000
Time (s) (a) Pacman
Time (s)
(b) Seaquest
Figure 7: Comparison of behavioral clowing accuracy versus end-lo-end training time for the AE, VAE, and TT, when keeping the number of data points seen by each algorithm the same by using only a single pass through the data. Points in the upper left indicate lower training time and higher accuracy. Note that TT- GPU tends to perform better in terms of lower training time and higher accuracy than the AE and VAE.
In the 5 trajectory case, the encoding time for TT-GPU between 0.4x and 2.8x that of the AE and VAE. In the 160 trajectory case, the encoding time for TT-GPU is between 1.5x and 8.8x that of the AE and VAE. The current Python implementation of TT-ICE* is not optimized for an inference task like the one we have in this work. As shown in Algorithm 1, encoding new data using the TT-encoder is a task that is ial in the dimensions. More elficient parallel approaches may be envisioned that can further accelerate the encoding task. There are multiple works in the literature on scaling such operations (Daas et al., 2022; Solomoník et al., 2014), and perhaps these can be considered in the future. However, even with this inefficiency, our results indicate faster training.
To determine why encoding is slower for the TT, we can investigate both flop counts and actual compute time. Moreover, we see how these vary with dataset size. Flop counts (e.g., the number of multiply-add operations) for the AE and VAE are calculated using the Ptflops package (Sovrasov, 2023). Flop counts for the TT are calculated by counting the number of flops per matrix multiplication in Algorithm 1 (lines 9 and 13).
In Table 6, the ratio in encoding times and flop counts between the TT and AF//VAF are shown. Although the flop counts in the TT are generally higher, this does not lead proportional increase in encoding time (c.g., TT-GPU uses 5.8x many flops the AE in the 5 trajectory case for BeamRider, but the encoding time for TT-GPU is 2.0x as long as the encoding time for the AE). This is partly due to the number of sequential calculations needed for each encoding method; in our experiments, the AE/VAE requires thirteen sequential layers (intermediate layers, including batch norm layers and activation functions), while TT projection requires only five tensor contractions (the number of tensor cores mimus one).
We also explore the scalability of the inference speed of the TT with the amount of data. During training, the tensor ranks of the TT increase as needed to represent new data. The number of flop counts for encoding using the TT grows quadratically with As such, the scalability of the inference speed of the TT depends on the ability of the TT to Iearn generalizable latent structure that can accurately represent new data, as this will lead to fewer and smaller increases in r as new data is introduced.
It difficult to determine this scalability a-priori, but we empirically investigate the growth in the tensor ranks and flop counts in Figure 8. As the amount of data increascs, the maximum rank of the TT grows, but the grotth rate tends level off. The slow growth of the TT ranks suggest that for these Atari games, the TT is able to learn a generalizable latent structure and scales well with additional data.
15

---

Published in Transactions on Machine Learning Research (04/2024)
Table 5: Median time to encode the training and validation data in seconds using the various encodings. the 5 trajectory case, the encoding time for TT-GPU between 0.4x and 2.8x that of the AE and VAE. In the 160 trajectory case, the encoding time for TT-GPU is between 1.5x and 8.8x that of the AE and VAE.
Trajectories
160
Trajectories
AE
VAE
TT-CPU
TT-GPU
AE
VAE
TT-CPU
TT-GPU
BeamRider
3.6
3.6
160.9
7.2
156.9
139.9
1906.5
673.5
Breakout
2.2
0.9
29.7
0.9
27.9
26.1
165.1
54.8
MsPacman
2.9
2.9
111.8
7.9
89.4
90.9
1153.2
484.4
Pong
12.8
12.4
519.2
25.4
433.4
448.1
1994.3
680.5
Qbert
2.0
2.0
64.1
2.9
60.3
60.6
388.5
123.6
Seaquest
3.2
3.1
116.3
7.9
128.3
132.9
1714.4
619.4
SpaceInvaders
2.1
1.9
73.8
71.7
72.3
1387.0
632.6
Table 6: Ratio of encoding times and flop counts between the TT and AE/VAE. Although the flop counts are much higher for the TT than the AE/VAE, E, this does not lead to proportionate increase in encoding time.
Trajectories
160
Trajectories
TT-GPU / AE TT-GPU / VAE TT-GPU
AE
TT-GPU VAE
BeamRider
2.0
2.0
4.3
4.8
Breakout
0.4
1.0
2.0
2.1
Eucoding
MsPacman
2.8
2.8
5.4
5.3
Time
Pong
2.0
2.0
1.6
1.5
Ralio
Qbert
1.5
1.5
2.0
2.0
Seaquest
2.5
2.5
4.8
4.7
Spacelnvaders
1.4
1.6
8.8
8.7
BeamRider
5.8
4.8
12.8
7.8
Breakout
0.8
0.8
3.9
3.0
Flop
MsPacman
7.6
6.5
12.7
8.0
Count
Pong
1.9
1.7
2.3
2.0
Ratio
Qbert
5.3
4.9
7.5
6.2
Seaquest
4.8
3.9
10.4
6.3
SpaceInvaders
5.0
4.5
16.3
9.7
An area of future research will be to improve the encoding time the TT. However, we reinforce that though the encoding time of the TT longer than that of the AE and VAE, the overall time for training the TT on GPU is still much lower than that of the AE and VAE since training the encoding is the most time-intensive task.
4.4.6 Additional Results
Appendix B.1 provides a more detailed investigation into the timing differences between the methods. The speedups achieved by the TT compared the AE and VAE are split into three components:
Time to train the encoder.
Time to encode training and validation data used to train the action predictor.
3. Time to train the action predictor.
To ensure that our results are consistent across different settings, we also examine how changing the latent space size (Appendix B.2) affects the training time and action prediction accuracy of the AE, VAE, and TT. The results are consistent: we find that as the latent space size and number of training trajectorics are changed, TT achieves similar prediction accuracies as the AE and VAE but provides speedup in training time.
16

---

Published in Transactions on Machine Learning Research (04/2024)
Maximum TT-Rank
FLOPs
Beam Rider
Breakout
Pacman
Pong
Qbert
Seaquest.
Space Invaders
10
1040
Maximun TT- Ta-k
104
10
10a
10
10a
1nd
107
10
10
109
103
107
105
100
10
10
104
10a
109
10a
10
10a
10c
10
10a
10a
10
10
201
10a
10
10a
207
109
108
Number of frames
Figure 8: Marimum tensor ronk (red) and flop count (teal) for the TT during training. As training progresses, the marimum lensor rank and the flop count increase, but this growth role decreases. The slow growth the TT ranks suggests that the TT-ICE" algorithm is able to learn generalizable latent structure; as the amount of data seen by the TT-ICE" algorithm increases, smaller updates to the TT ranks needed to represent new data. In cases like this where the tensor ranks of the TT grow slowly, the encoding time for the TT will scale well since the flop count grows quadratically with the TT ranks r.
In Appendix B.3, we provide the gameplay scores achieved by the RL agent used to generate the demon- stration data we used for behavioral cloning. The RL agent achieves gameplay scores that are higher than randomly acting agent and achieves human level performance in two of the seven games.
5 Conclusion
In this work, we have explored using tensor networks as a method for learning a low-dimensional represen- tation of image data for behavioral cloning, In particular, we compare the computational speed of using TT against AEs and VAFs for extracting this low dimensional representation, finding a 2.6x to 9.6x speedup compared to AEs VAEs. Furthermore, the TT-encoding can be efficiently trained on CPU as well, which leads lo comparable This speedup comes at no cost in terms of behavioral cloning accuracy or the gameplay scores achieved by the behavioral cloned agents. As such, the TT-based approach enables learning while having limited access to GPUs, as allows for the expensive dimensionality reduction step to be performed on CPUs at a comparable or faster time than using an AE or VAE on GPUs. Given the large computational requirements of learning directly from video data, future work will aim extend this methodology tackling larger game environments for which the large input space is computational bottleneck.
belter training limes than the AE and VAE trained on GPU (0.9x to 5.4x faster).
Although this paper explores the use of tensor networks for extracting a latent space from video data for behavioral cloning, the approach outlined in this paper may be broadly applicable. Tensor networks have been shown to successfully compress data in a wide variety settings, ranging from EEG signals (Zhang et al., 2017), video data (Zhao et al., 2020; Fonal & Zdunek, 2019), images (Fonal & Zdunek, 2019), and audio (Fonal & Zdunck, 2019). Tensor networks can be generally applied tensorizable features, as many data formats naturally tensors (e.g., videos are 4-way tensors). One and two dimensional data also be reshaped into tensors then compressed via TT. We also note the possibility of using TT to compress tensorizable features and using these compressions along with non-compressod features as input into a learning algorithm.
Abstractly, the benefits this approach are to problems that have separable structure, where certain features of the problem interact limited ways with others. Different tensor network structures define such interac- tions differently, based on network topology. is difficult to a-priori say if a particular application exhibits this structure, and the above literature and the current work continue provide empirical evidence that
17

---

Published in Transactions on Machine Learning Research (04/2024)
this structure exists in quite general problems. A future area of research would be further exploration of data and tasks for which the TT format works well a method of feature extraction.
Another interesting research direction would be to explore a TT-based encoder can be trained using frames from multiple games. If frames from different video games do share some common structure, it is possible that the rank of the TT would grow sublinearly with the number of games included, leading to a more storage-efficiet alternative. It would also be interesting to explore if an encoder trained in this fashion would be able to generalize to similar but out-of-sample games.
Acknowledgments
We acknowledge support from the Automotive Research Center (ARC) in accordance with Cooperative Agreement W56HZV-19-2-0001 with U.S. Army DEVCOM Ground Vehicle Systems Center. This research was supported in part through computational resources and services provided by Advanced Research Com- puting at the University of Michigan, Ann Arbor.
References
Doruk Aksoy, David Gorsich, Shravan Vecrapaneni, and Alex A Gorodctsky. An incremental tensor train decomposition algorithm. SIAM Journal on Scientific Compuling, 46(2):A1047 A1075, 2024.
Alexander Amini, Wilko Schwarting, Guy Rosman, Brandon Araki, Sertac Karaman, and Daniela Rus. Variational autoencoder for end-to-end control of autonomous driving with novelty detection and training de biasing In 2018 LEEE/RSJ International Conference on Intelligent Robots and Systems (IROS), pp. 568-575. IEEE, 2018.
Bowen Baker, Ilge Akkaya, Peter Zhokov, Joost Huizinga, Jie Tang, Adrien Ecoffet, Brandon Houghton, Raul Sampedro, and Jeff Clune. Video pretraining (vpt): Learning to act by watching unlabeled online videos. Advances in Neural Information Processing Systems, 35:24639 24654, 2022.
Christopher Berner, Greg Brockman, Brooke Chan, Vicki Cheung, Przemyslaw Dębiak, Christy Dennison, David Farhi, Quirin Fischer, Shariq Hashme, Chris Hesse, et al. Dota 2 with large scale deep reinforcement lcarning. ar preprint vrxiv:1112 1912006600, 2019.
Eli Bronstein, Mark Palatucci, Dominik Notz, Brandyn White, Alex Kuefler, Yiren Lu, Supratik Paul, Payam Nikdel, Paul Mougin, Hongge Chen, Justin Fu, Austin Abrams, Punit Shah, Evan Racah, Benjamin Frenkel, Shimon Whiteson, and Dragomir Anguelov. Hierarchical mosel-based imitation learning for planning in autonomous driving. In 2022 TEEE/RS /RSJ International Conference on Intelligent Robots and Systems (IROS), pp. 8652-8659. IEEE, 2022.
Daniel Brown, Russell Coleman, Ravi Srinivasan, and Scott Nickum, Safe imitation learning via fast Bayesian reward inference from preferences. In Hal Daumé III and Aarti Singh (eds.), Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pp. 1165-1177. PMLR, 13-18 Jul 2020. URL https://proceediie
mlr. press/ press/vv19//e w119/broom20a. html.
Xiaoyu Cai, Qiongli Jing, Bo Peng, Yuanyuan Zhang, Yuting Wang, Ju Tang, et al. Automatic traffic state recognition based on video features extracted by an autoencoder. Mathematical Problems in Engineering. 2022, 2022.
Brian Chen. Atari Games Dataset. https://deepblue wib.luiicch....... gr87,, 2023.
www.h.
Brian Chen, Siddhant Tandon, David Gorsich, Alex Gorodetsky, and Shravan Veerapaneni. Behavioral cloning in atari games using a combined variational autocncoder and prodictor model. In 2021 IEEE Congress on Evolutionary Computation (CEC), pp. 2077-2084. IEEE, 2021a.
Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. simple framework for contrastive learning of visual representations. In Hal Daumé IIT and Aarti Singh (eds.), Proceedings of the 37th
18

---

Published in Transactions on Machine Learning Research (04/2024)
International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pp. 1597-1607. PMLR, 13-18 Jul 2020. URL https://proceedingg .wl. http://nww......ses/eeeeeeeeeeee ptress/v119/ catps//e211//c .hmll.
Xin Chen, Sam Toyer, Cody Wild, Scott Emmons, Ian Fischer, Kuang-Huei Lee, Neel Alex, Steven H Wang, Ping Luo, Stuart Russell, et al. An empirical investigation of representation learning for imitation. In Thirty-fifth Conference on Neural Information Processing Systems Datasets and Benchmarks Track (Round 2), 2021b.
Hussam Al Daas, Grey Ballard, and Peter Benner. Parallel algorithms for tensor train arithmetic. SIAM Journal on Scientific Computing, 44(1):C25 C53, 2022.
Krzysztof Fonal and Rafal Zdunck. Distributed and randomized tensor train decomposition for feature extraction. In 2019 International Joint Conference on Neural Networks (IJCNN) pp. 1 8. IEEE, 2019.
Timur Garipov, Dmitry Podoprikhin, Alexander Novikov, and Dmitry P. Vetrov. Ultimate tensorization: compressing convolutional and FC layers alike. Coll,, abs/1611.08214
2016. URL http://arxiv.orr
Charles R. Harris, K. Jarrod Millman, Stéfan J. van der Walt, Ralf Gommers, Pauli Virtanen, David Cournapeau, Eric Wieser, Julian Taylor, Sebastian Berg, Nathaniel J. Smith, Robert Kern, Matti Pi- cus, Stephan Hoyer, Marten H. van Kerkwijk, Matthew Brett, Allan Haldane, Jaime Fernández del Río, Mark Wiebe, Pearu Peterson, Pierre Gérard-Marchant, Kevin Sheppard, Tyler Reddy, Warren Weckesser, Hameer Abbasi, Christoph Gohlke, and Travis E. Oliphant. Array programming with NumPy. Nature, 585(7825):357 September 2020. doi: (0.108//441586-002-888888 //doi. //doi.org//0. 00022/8/1156-002222500
URL https:
Ashley Hill, Antonin Raffin, Maximilian Ernestus, Adam Gleave, Anssi Kanervisto, Rene Traore, Prafulla Dhariwal, Christopher Hesse, Oleg Klimov, Alex Nichol, Matthias Plappert, Alec Radford, John Schulman, Szymon Sidor, and Yuhuai Wu. Stable baselines.
com/hill- methiillllsaaaalllaee aatillllsttaleeeee
2018.
Edward J Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, and Weizhu Chen. LoRA: Low-rank adaptation of large language models. In International Conference Learning Representations, 2022. URL sttps://openreeee
Benjamin Huber, Reinhold Schneider, and Sebastian Wolf. A Randomized Tensor Train Singular Value Decomposition, pp. 261-290. Springer International Publishing, Cham, 2017. ISBN 978-3-319-69802- 9. URL
doi: 111-007/778-3-11--69655
Ahmed Hussein, Mohamed Medhat Gaber, Eyad Elyan, and Chrisina Jayne, Imitation learning: A survey of learning methods. ACM Computing Surveys (CSUR), 50(2):1 35, 2017.
Miguel Jaques, Michael Burke, and Timothy M Hospedales. Newtonianvac: Proportional control and goal identification from pixels via physical latent spaces. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 4454-4463, 2021.
Anssi Kanervisto, Joonas Pussinen, and Ville Hautamāki. video games. In 2020 IEEE conference on games (CoG), pp. 558-565, IEEE, 2020.
Benchmarking end-to-end behavioural cloning on
Heecheol Kim, Yoshiyuki Ohmura, and Yasuo Kuniyoshi. Transformer-based deep imitation learning for dual-arm robot manipulation. In 2021 IEEE/RSJ International Conference on Intelligent Robols and Systems (IROS), pp. 8963-899 8972. IEEE, 2021.
Yong-Deok Kim, Eunhyeok Park, Sungioo Yoo, Taelim Choi, Lu Yang, and Dongjun Shin. Compression deep convolutional neural networks for fast and low power mobile applications. In 4th International Conference on Learning Representations, ICLR 2016, San Juan, Puerto Rico, May 2-4, 2016, Conference Track Proceedings, 2016. URL http://arxis: wttps/ii..rg//b/////.eae/ /ttp///bb0/11..053...
Diederik P. Kingma and Max Welling. Auto-cncoding variational bayes In 2nd International Conference on Learning Representations, ICLR 2014, Banff, AB, Canada, April 14-16, 2014, Conference Track Pro- ceedings 2014. URL http://xxxvvvoog//eee/ 131/. 1112..114......
19

---

Published in Transactions on Machine Learning Research (04/2024)
Jean Kossaifi, Aran Khanna, Zachary Lipton, Tommaso Furlanello, and Anima Anandkumar. Tensor con- traction layers for parsimonious deep nets. In Proceedings of the IEEE Conference on Compuler Vision and Pattern Recognition Workshops, pp. 26 32, 2017.
Jean Kossaifi, Zachary C Lipton, Arinbjorn Kolbeinsson, Aran Khanna, Tommaso Furlanello, and Anima Anandkumar. Tensor regression networks. Journal of Machine Learning Research, 21(13)):1-21 2020.
V Lebedev, Y Ganin, M Rakhuba, I Oseledets, and V Lempitsky. Speeding-up convolutional neural networks using fine-tuned cp-decomposition. In 3rd International Conference on Learning Representations, ICLR Conference Track Proceedings, 2015.
Nya Loshchilov and Frank Hutter. SGDR: stochastic gradient descent. with warm restarts. In 5th International Conference on Learning Representations, ICLR 2017, Toulon, France, April 24-26, 2017, Conference Track Proceedings. OpenReview. net,, 2017. URL
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In 7th International Conference on Learning Representations, ICLR 2019, New Orleans, LA, USA, May 6-9, 2019. OpenReview.net, 2019. www.neettoorriee dot.ornnnnd-BeeiCGGG.....
URL
Angshul Majumdar. Blind denoising autoencoder. IEFE transactions on neural nebworks and learning systems, 30(1):312 317, 2018
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A. Rusu, Joel Veness, Marc G. Bellemare, Alex Graves, Martin Riedmiller, Andreas K. Fidjeland, Georg Ostrovski, Stig Petersen, Charles Beattie, Amir Sadik, Ioannis Antonoglou, Helen King, Dharshan Kumaran, Daan Wierstra, Shane Legg, and Demis Hassabis. Human-level control through deep reinforcement learning. Nature, 518:529-55 533, 2015.
Alexander Novikov, Dmitri Podoprikhin, Anton Osokin, and Dmitry P Vetrov. Tensorizing neural networks. Advances in neural information processing systems, 28, 2015.
Ryosuke Okuta, Yuya Unno Daisuke Nishino, Shohei Hido, and Crissman Loomis. Cupy: A numpy- compatible library for nvidia gpu calculations. In Proceedings of Workshop on Machine Learning Systems (LearningSys) in The Thirty-fiss Annual Conference on Neural Information Processing Systems (NIPS), 16.pdf.
2017. URL http: attp://learningsys hwwpeiinssssso.nnnnneoeeeemeeee
Ivan V Oseledets. Tensor-train decomposition. S1AM Journal on Scientific Computing, 33(5)22295 2317, 2011.
Jyothish Pari, Nur Muhammad, Sridhar Pandian Arunachalam, Lerrel Pinto, et al. The surprising effective- ness of representation learning for visual imitation. Robolics: Science and Systems X XVIII, 2021.
Simone Parisi, Aravind Rajeswaran, Senthil Purushwalkam, and Abhinav Gupta, The unsurprising effec- tiveness of pre-trained vision models for control. Tn Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato (eds.), Proceedings of the 39th International Conference on Machine Learning, volume 162 of Proceedings of Machine Learning Research, pp. 17359-17371. PMLR, 17-23 23 Jul 2022. URL https:/ http////oooeeeeeeeee/ee/eeee .mlr mlr press/v162 https///////paaraiii21/ parisii2i2i. html.
Emmanuel Pintelas, Joannis E Livieris, and Panagiotis E Pintelas. A convolutional autoencoder topology for classification in high-dimensional noisy image datasets. Sensors, (222):7771, 2021.
Marsela Polic, Ivona Krajacic, Nathan Lepora, and Matko Orsag. Convolutional autoencoder for feature extraction in tactile sensing. IEEE Robotics and Automation Letters, 4(4):3671-38) 16778,,11777711 2019.
Jun Qi, Chao-Han Huck Yang, Pin- Yu Chen, and Javier Tejedor. Exploiting low-rank tensor-train deep
neural networks based on riemannian gradient descent with illustrations of speech processing. IEEE/ACM Transactions on Audio, Speech, and Language Processing, 31:633-66 -62, 2023. doi: 10.1109/TASLP.2022.. 3231714.
Antonin Raffin. RI baselines zoo. http:://gstau...ooaaaffftt//tttt
2018.
20

---

Published in Transactions on Machine Learning Research (04/2024)
Stéphane Ross and Drew Bagnell. Efficient reductions for imitation learning. In Proceedings of the thir- teenth international conference on artificial intelligence and statislics, pp. 661-668. JMLR Workshop and Conference Proceedings, 2010.
Pierre Sermanet, Corey Lynch, Jasmine Hsu, and Sergey Levine. Time-contrastivy networks: Self-supervised learning from multi-view obscrvation. In 2017 IEEE Conference on Computer Vision and Pattern Recog- nition Workshops (CVPRW), pp. 486 487. IEEE, 2017.
Rutav M Shah and Vikash Kumar. Rrl: Resnet representation for reinforcement learning. In Marina Meila and Tong Zhang (eds.), Proceedings of the 38th International Conference Machine Leurning, volume 139 of Proceedings of Machine Learning Research, pp. 9465-94777 9476. PMLR, 18 24 Jul 2021. URL https://proceeddngg meridiinnnmmr......mmmmmme. /ttp (rres//139/shhhh22 html.
Manish Sharma, Panos P Markopoulos, Flli Saber, M Salman Asif, and Ashley Prater-Bennetie. Con- volutional auto-encoder with tensor-train factorization. In Proceedings of the IEEE/CVP international conference on computer vision, pp. 198-206, 2021.
Hao Shen, Weikang Wan, and He Wang. Learning category-level generalizable object manipulation policy via generative adversarial sell-imilation learning from demonstrations. IEEE Robolics and Aulomation Letters, 7(4):11166-111 11173, 2022.
Edgar Solomoník, Devin Matthews, Jeff R Hammond, John F Stanton, and James Demmel. A massively parallel tensor contraction framework for coupled-cluster computations. Journal of Parallel and Distributed Computing, 74((1))3116-311) 3190, 2014.
Vladislav Sovrasov. Ptflops: a flops counting tool for neural networks in pytorch framework, 2023. URL
https:///ithhub.cmmt wtt.:/cc/////.. htwps//ww...oooooorr////////p/// coutt//ooooooooooeooerorao
K Sreelakshmi and Renjith V Ravi. An encryttioa-teen-compessooo scheme using autoencoder based image compression for color images. In 2020 7th International Conference on Smart Structures and Systems (ICSSS), pp. 1-5. IEEE, 2020.
Adam Stooke, Kimin Lee, Pieter Abbeel, and Michael Laskin. Decoupling representation learning from reinforcement learning. In Marina Meila and Tong Zhang (eds.), Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pp. 9870 9879. PMLR, 18-24 Jul 2021. URL https: httpa:///ccceeeegs..oe. mlr. prwss//13// stooke2oo..2.... html.
Mojtaba Taherisadr, Mohsen Joncidi, and Nazanin Rahnavard. Eog signal dinensionality reduction and classification using tensor decomposition and deep convolutional neural networks. In 2019 IEEE 29th International Workshop on Machine Learning for Signal Processing (MLSP), pp- 1-6. IEFF, 2019.
Oriol Vinyals, Igor Babuschkin, Wojeiech M Czarnecki, Michaël Mathicu, Andrew Dudzik, Junyoung Chung, David H Choi, Richard Powell, Timo Ewalds, Petko Georgiev, et al. Grandmaster level in starcraft ii using multi-agent reinforcement learning. Nature, 575(7782):350-315 2019.
Tianyu Wang, Nikhil Karnwal, and Nikolay Atanasov. Latent policies for adversarial imitation learning.
arXiv preprint Xrc::2200 11299, 2022.
Yasi Wang, Hongxun Yao, and Sicheng Zhao. Auto-encoder based dimensionality reduction. Neurocomputing, 184:232-242, 2016.
Xuewen Yang, Yingru Liu, and Xin Wang. Reformer: The relational transformer for image captioning. Tn Proceedings of the 30th ACM International Conference on Multimedia, pp. 5398 5406, 2022.
Yinchong Yang, Denis Krompass, and Volker Tresp. Tensor-train recurrent neural networks for video classi- fication. In Proceedings of the 34th International Conference on Machine Learning Volume 70, ICML/17, pp.
Jrg, 2017.
Jianbo Yu and Xingkang Zhou. One-dimensional residual convolutional autoencoder based feature learning for gearbox fault. diagnosis. IEFE Transactions on Industrial Informatics, 16(10):6347-6333 6658,, 2020.
21

---

Published in Transactions on Machine Learning Research (04/2024)
Rose Yu, Stephan Zheng, Anima Anandkumar, and Yisong Yue. Long-term forecasting using higher order tensor rnns. Xiv preprint or orxi 71 17711100000 1.00033, 2017a.
Xiyu Yu, Tongliang Liu, Xinchao Wang, and Dacheng Tao. On compressing deep models by low rank and sparse decomposition. In Proceedings of the IEEE conference computer vision and pattern recognition, pp. 7370 7379, 2017b.
Zhecheng Yuan, Zhengrong Xue, Bo Yuan, Xueqian Wang, Yi Wu, Yang Gao, and Huazhe Xu. Pre-trained image encoder for generalizable visual reinforcement learning. Advances in Neural Information Processing Systems, 35:13022 13037, 2022.
Jaime Zabalza, Jinchang Ren, Jiangbin Zheng, Huimin Zhao, Chunmei Qing, Zhijing Yang, Peijun Du, and Stephen Marshall. Novel segmented stacked autocncoder for effective dimensionality reduction and feature extraction in hyperspectral imaging. Neurocompuling, 185:1 10, 2016.
Wenjie Zhang, Jiqing Han, and Shiwen Deng. Heart sound classification based on scaled spectrogram and tensor decomposition. Expert Systems with Applications, 84:220-221 231, 2017.
Bin Zhao, Xuelong Li, and Xiaoqiang Lu. Tth-rmm: trnsor-train hierarchical recurrent neural network for vidco summarization. IEEE Transactions on Industrial Electronics, (8(4):3629-363) 2020.
Appendix
A Additional Details on Training
Architectural details for the encoding part of our ALs and VAEs laid out in Table 7. The decoding portion of these networks mirrors the encoding portion, with convolutional layers replaced by transposed convolutional layers.
Table 7: Architecture for the encoding part of AEs and VAEs. After each intermediate layer also apply batch norm followed by LeakyReLU- these are excluded from the table for simplicity. The asymmetric padding in the first layer is to account for the shape of the input.
Layer
Layer
Inp.
Out.
Kernel Size
Stride
Padding
#
Type
Ch.
Ch.
1
Conv
3
32
3
2
(4, 1)
2
Conv
32
64
3
2
(1, 1)
3
Conv
64
64
3
2
(1, 1)
4
Conv
64
16
1
(0, 0)
5
Flatten
6
Linear
The AEs and VAEs are hyperparameter tuned individually for each game over initial learning rates of 10-2 10-3 and 10 and L2 weight decay penalties of 1, 1,10-2, and 10-4, after which the best model is selected based on validation error and used for encoding. The networks optimizer (Loshchilov & Hutter, 2019) with cosine learning rate decay without restarts (Loshchilov & Hutter, 2017) for fifty epochs in the low data case and five epochs in the moderate data case. The AE is trained to minimize reconstruction error measured via mean squared error while the VAE is trained using the standard ELBO loss (Kingma Welling, 2014). An additional or 20 trajectories (in the limited and moderate data regimes, respectively) are used as validation set in order to hyperparameter tume the networks. Note that the TT does not require this step of validation and does not use this additional data.
trained using the Adam-W W
22

---

Published in Transactions on Machine Learning Research (04/2024)
Training the Encoder
Encoding Data
Training the Predictor
0.8
Rel. time
0.6
0.4
0.50
0.32
01.22
0.50
0.31
0.31
0.32
0.2
0.0g
0.04
0.10
0.04
0.06
0.00
0.00
0
AE
VAB
TT-CPP
TT-GPU
AB
VAR
TT-CPU
TT-GPU
TT-CPU
TT-GPU
VAE
TT-CPU
TT-CPU
AB
VAE
TT-GPU
AB
TT-CPU
TT-
VAB
TI-PPU
TT-GP
Beam Rider
Breakout
Ms.Pacman
Pong
Qbert.
Seaquest
Space Invaders
Figure 9: Time for end-to-ent to-end training procedure in the limited data regime, averaged across all experiments, as a fruction of the time required to train the full-frome predictor. Extracting a lalent space before training a predictor reduces end-to-end training time in all cases, but using the TT-GPU leads to 26.1 to 29.7x reduction in truining time. Compared the AE VAE, using the TT-GPU leads to a 8.2x to 9.6x speedup.
B Additional Results
B.1 Breakdown of Speedups
In Figures 9 and 10, we further break down the timings for training the AE, VAE, and TT-based predictors follows:
Training the encoder: For the AE and VAE, this time the total time required to perform a hyperparameter optimization by training nine AEs VAEs with different network hyperparameters and selecting one based on the validation loss measured on separate validation set. Recall that the TT encoder does not require such procedure. For the TT, this is the time required to perform the TT decomposition (using TT-ICE*) to extract the latent space representation of the input data.
Encoding data: Time to pass all the training and validation data through the AE, VAE, and TT in order to extract latent space representations
Traóning the action predictor: Time to train the action predictor the encodings extracted in the previous step. The architecture of the predictor is same across all three methods. For each model and for each game, nine predictors are trained using different initial learning rates (10-2, 10-3, and 10-4) and L2 regularization penalties (1, 10-2, and 10-4). The predictor yielding the lowest cross-entropy action prediction loss measured on another separate validation set is selected.
Note that the most inte-intensive task is training the encoding. As such, reducing the time it takes to train the encoding by using a TT instcad of an AE or VAE leads to significant reduction in overall training time.
B.2 Effect of Changing Epsilon
We conduct experiments lesting the effect of changing the size of the latent space for the TT, AE, and VAE. In experiments, the latent size of the AE and VAE is set up to match the latent size of the TT. This experiment is conducted using 160 trajectories as training data, with five training epochs for the AE and VAE. In Table 8, the latent size as epsilon is changed is displayed.
In Figure 11, the effect of changing epsilon / the latent size on prediction accuracy is displayed. The general pattern appears to be that decrcasing epsilon increasing the latent size leads to a higher prediction accuracy, as expected. However, for some games, increasing the latent space size past a certain point leads to small negligible improvements in prediction accuracy. The effect of changing the latent space size timings
23

---

Published in Transactions on Machine Learning Research (04/2024)
Training the Encoder
Encoding Data
Training the Predictor
0.79
0.8
Rel. time
0.6
0.42
0.50
0.41
10300
0.41
0.40
0.42
0.4
9,14
p.3s
0.14
0,30
0.2
0
AE
VAB
TT-CPP
TT-GPU
AB
VAR
TL-CPU
TT-GPU
TT-CPU
TT-GPU
VAE
TT-CPU
TT-GPU
AB
VAB
TT-GPU
AB
TT-CPU
TT-GPU
VAB
TT-OPU
TT-GPU
Beam Rider
Breakout
Ms.Pacman
Pong
Qbert.
Seaquest
Space Invaders
Figure 10: Time for the end-to-end training procedure in the moderate data, regime, averaged across all experiments, as a fraction of the time lakes to train a predictor using the full RGB fromes. Extracting a latent space reduces end-to-end training time in all cases, but using TT-GPU leads to a 3.4x to 9.8x reduction in training time compared to no encoding and a 2.6x to 4.8x speedup relative to the AE VAE.
Table 8: Latent size epsilon for the TT is changed. The reported numbers average across three seeds (which use different training data).
Epsilon
0.01
0.02
0.05
0.10
0.15
BeamRider
14645
9867
3324
297
Breakout
3561
957
50
18
MsPacman
11962
6421
867
15
2
Pong
1287
138
4
1
Qbert
2267
1560
308
43
10
Seaquest 15163
8170
1338
37
SpaceInvaders
17191
13795
8035
3994
2097
is shown in Figure 12. Using a larger latent space size tends to increase the end-to-end training time of the models. These results consistent. with findings in the main text. at the different epsilon levels tested, the TT-based predictor yielded similar prediction accuracies as the AE and VAE, while being faster to train.
B.3 Gameplay Performance of RL Agent
This section provides some statistics on the RL agent used to generate the demonstration data.
Table provides the gameplay scores achieved by the RL agent compared to agent that acts randomly by uniformly selecting amongst available actions and a human player. The human player scores are taken from Mnih et al. (2015). The RL agent achieves higher scores than an agent acting randomly and achieves human-level performance in 2 of the games. Table 10 provides the gameplay lengths (number of steps between the beginning of a game to the end) achieved by the RL agent versus an agent that acts randomly. Note that gameplay lengths were not available from Mnih al. (2015). This table indicates that the RI. agent has learned behavior that differs from randomly acting agent.
24

---

Published in Transactions on Machine Learning Research (04/2024)
AE
VAE
TT
No encoding
Action Prediction Accuracy (%)
Beam Rider
Breakout
Pacman
Pong
Qhert
Seaquest
Invaees
100
80
60
40
20
0.15
0.10
0.05
0.00
0.01
0.13
001
0.02
0.01
0.15
0.00
000
0.02
0.01
0.15
0.10
0.00
0.00
0.91
0.15
0.10
0.05
0.02
0.01
0.15
0.10
0.05
0.02
0.01
0.15
0.10
0.05
0.02
0.01
0
Figure 11: Action prediction accuracy of behavioral cloning frameworks using different encoding methods at varying e levels. As the lalent space size is increased for the encoding-based predictors, the action prediction accuracy of the predictors tends to increase. The speedup offered by TT does not cause any performance deterioration. Furthermore, the nonlinearity offered by AE/VAE does not result in improved latent space representation that achieves higher action prediction accuracy with a smaller latent space. Since Pong with = 0.15 and ε = 0.10 returns the same latent size, the accuracy for e = 0.15 is omitted. Please refer to Table for latent space sizes corresponding to each level for each game.
AE
VAE
TT-CPU
TT-GPU
No encoding
Beam Rider
Breakout
Pacman
Pong
Qbert
Seaquest
Spacelavaders
105
Total Training Time (s)
10
10a
102
0.10
0.05
0.02
n.
D010
0.02
0.01
0.35
0.06
n.n2
0.01
0.35
0.10
0.05
0.02
un
0.35
0.10
0.05
0.01
0.15
0.30
400
0.02
0.01
0.16
0.10
0.06
0.02
0.01
Figure 12: Total lime elapsed Irain behavioral cloning frameworks using differrntt encoding methods al
varying E levels. As the latent space size increases, the total training time for the vorious encoding based predictors increases. Across all games and levels, TT-GPU encoding offers the fastest end-to-end training time. Please refer to Table for latent space sizes corresponding to each level for each different game.
Table 9: Comparison of the gameplay scores achieved by the RL agent compared to agent that acts randomly and a human player. The human gameplay scores are taken from Mnih et al. (2015).
Random
RI.
Human (Mnih et al., 2015)
BoamRider
366.1
1012.3
5775.0
Breakout
1.4
64.0
31.8
MsPacman
244.8
1883.4
15693.0
Pong
-20.3
21.0
9.3
Qbert
167.5
5181.0
13455.0
Seaquest
67.2
1962.0
20182.0
25

---

Published in Transactions on Machine Learning Research (04/2024)
Table 10: Comparison of the gameplay lengths (number of time steps from the beginning of the game to the end) achieved by the RL agent compared to agent that moves randomly. Note that gameplay lengths for human player were not available from Mnih et al. (2015).
Random
RL
BeamRider
1289.5
1714.3
Breakout
153.2
982.1
MsPacman
489.7
1025.0
Pong
926.8
1570.7
Qbert
308.8
1390.0
Seaquest
492.2
2414.9
Spacelnvaders
488.7
660.5
26


Reviews:
[{"review_id": "5E3mSMIjGe", "reviewer": "Reviewer_wupi", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "1. The paper is clearly written, presents the scope, work and all components well. It presents earlier work very clearly. It is only in subtle situations where I think the paper could use more details:\n   1. How L is set for AE/VAE (only appears in 4.2)?\n   2. What N is (I assume 2500 given 3.1)?\n   3. How the action prediction accuracy is computed (I assume it is likelihood of RL agent under learned action logits)\n   4. The choice of using a 6-way TT format, and how each individual dimension was selected aren’t clear. \n2. The comparison to AE and VAE baselines was well executed and overall the emphasis being put on speed comparisons given similar accuracies was useful.\n   1. One weakness is the particular choice of VAE selected. Nowadays, SOTA results on discrete domains like ATARI are usually leveraging discrete generative representations (either DAE, VQ-VAE, VQ-GAN, etc, see DreamerV3) or contrastive representations (usually VIT based, see MineClip).\n   2. Would it have led to a situation where the learnt encoder would have a higher accuracy?\n3. The results in Figure 4 and 5 appear quite low.\n   1. Could you comment on the accuracies and if they arise from specifically hard trajectories to predict, or if they are just symptoms of under-fitting BC policies?\n   2. Similarly, Figure 5 has several games at chance or worst, what would be causing this?\n   3. It is also hard to tell what the actual ability level of the base RL agents are. How do the raw scores obtained on each game relate to SOTA agents on Atari? Are they human-level? Are the trajectories representing interesting and complex behaviors?\n   4. Figure 4 and 5 both seem to show that using raw RGB (“no encoding” baseline) are the best options. This usually isn’t the case when there is some transfer gap, could you comment on that?\n4. The setup assumes frame-level encoding, followed by an action predictor which just consumes 4 frames to predict the current action.\n   1. Did you explore a larger number of frame stacking or other architectures for the action predictor?\n   2. How would one integrate the two and leverage videos in the representation?\n5. Most of my worries about the paper are really about the TT format and its implications. Given this isn’t as relevant to this particular paper given its claims, I am not expecting these to be addressed, but would be interested in hearing the authors’ thoughts on them:\n   1. Isn’t a multilinear decomposition too “weak” to represent complex visual scenes in a way that supports generalization? I would have expected non-linearities to help, compared to “just” doing SVD?\n   2. Does the choice of dimensions (i.e. the reshaping done on the 4-dims videos to get to the 6-dim matrix) impact the performance/accuracy of the TT approximation or what it focuses on?\n   3. You mention that TT-SVD is limited to 30’000 frames. Does TT-ICE* have some kind of limitations (e.g. is the accuracy on earlier datapoints reduced as N increases?). Would it scale to O(1M) frames (imagine using the MineCLIP dataset)?\n   4. How would you extend it to handle videos in a way that shares information across the time axis? As far I understand, G6 is independent across frames?", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper explores the use of tensor-network representations (a multi-linear decomposition & approximation to the data) for behavior cloning, in particular as a replacement to AE or VAE representations. The scope and claims are mostly about speed gains from using tensor network representations. They demonstrate this empirically on a few ATARI games.\n\nOverall the paper is very clear about what it intends to do, and what it claims is very sharply and tactically defined. In effect most of my questions are more about what to expect from TT formats and comparing to more recent deep learning methods. But given the paper really only claims computational gains and not much more, I feel like it passes the TMLR two questions test for me, and hence I lean towards acceptance. I feel like people in the TMLR community will benefit from learning about the TT format.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "1. Move requested hyperparameters/choices of (1) earlier in the main text\n2. Discuss low performance in Figure 4 and 5 to ensure they do not affect applicability of these results", "broader_impact_concerns": "No concerns", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "1jk2XGDczB", "reviewer": "Reviewer_k7Jk", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n- This proposed method is both conceptually elegant and empirically effective.\n- The experiments include ablations, such as the latent space size and the amount of training data.\n- The paper is clearly written and provides implementation details.\n\nWeaknesses: I'm mainly concerned about the following aspects:\n- **Training scalability**: It is noted in the paper that the main advantage of the TT-based method is the saving in the training time. However, by comparing the limited data regime with the moderate data regime, we see that the amount of speedup decreases as more data is used during training. This raises a question: how well does the TT-based method work as the data size further scales up? \n- **Memory use and generalizability**: The method learns game-specific representations, i.e. the TT-cores. This means that the number of parameters grows linearly in the number of games, and the TT-based encoder cannot generalize to unseen games.\n- **Inference cost**: the TT-based method's encoding time is high. The paper leaves this as a future direction, but I'm concerned that this will be a main bottleneck in adapting the proposed method in more general settings.\n- Related to the above two points, it might be better if the authors could add some discussions on the applicability of their method. For example, for what kind of data do we expect an advantage of the proposed method?\n- Several design choices could be better explained:\n    - Choices of encoder baselines? For example, contrastive methods and masked prediction have been shown to produce better representations than (V)AE on images.\n    - Choices of tasks: why these choices of Atari games? Disclaimer that I'm not an expert in RL so I apologize if this is a silly question.\n    - Choices of the data generation process: currently DQN is used; (how) would the results be different as the data sources change?", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper aims to improve the computational efficiency of behavioral cloning. For extracting low-dimensional latent representations, it proposes an efficient dimension reduction method that is based on tensor decomposition.\n- The behavioral cloning task is instantiated as action prediction in gameplays.\n- The model used for action prediction consists of 2 parts: an encoding algorithm, and an action predictor.\n  - The action predictor is a standard MLP, which is kept the same for different encoders.\n- The main contribution of the paper lies in the choice of encoder: The paper proposed to use a tensor decomposition method to obtain a low-dimensional latent. The method of choice is TT-ICE, which is an incremental version of TT-SVD. The proposed TT-based method is better than baseline methods using AE/VAE, in terms of both computation time and performance.\n\nHere are some more detailed results:\n- Training time:\n    - when using limited data (i.e. 5 trajectories per game): the TT-based method is around 9 times faster than AE/VAE, and around 28 times faster than predicting directly on pixels.\n    - when using moderate data (i.e. 160 trajectories per game): the TT-based method is around 3 times faster than AE/VAE.\n    - for both data regime, when the TT-based method is ran on CPUs, it is still faster than the other methods trained on GPUs.\n- Performance:\n    - pSNR and MSE: TT-based method is significantly better.\n    - action prediction accuracy: TT-based method is mostly better, except for the game S and SI under limited data.\n    - normalized gameplay scores: TT-based method is comparable with AE; both are better than VAE.\n- Encoding time (related to inference time): TT-based method is no better.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "- Please consider including more discussions on training scalability, generalizability, inference cost, and other design choices, as mentioned in weakness.\n- Please consider comparing with more baselines, such as contrastive learning or masked prediction.\n- There's a minor typo: Sec 4.1: \"Rather than use\" -> \"Rather than using\"", "broader_impact_concerns": "There are no direct ethical concerns.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "kejsRi66aQ", "reviewer": "Reviewer_S6tk", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n1. The use of TTN as an encoder for low-dimensional feature extraction is interesting in the scenario of image data for behavior cloning. \n\n2. The TTN exhibits some performance gain in the experiments compared with AEs and VAEs. \n\n3. The discussion of the advantages of using TTN on GPUs is an important new contribution. \n\nWeaknesses: \n1. The technical novelty of this work is limited. The TTN model was proposed by Oseledets in 2011, and it has already been applied to speech, image, and signal processing. In particular, Yang et al., (2017) put forth using the TTN for video data processing, and Qi et al., (2023) even used low-rank Tensor-Train Network for speech processing. \n\n[1] I. V. Oseledets, “Tensor-train decomposition,” SIAM J. Sci. Comput., vol. 33, no. 5, pp. 2295–2317, 2011.\n\n[2] Y. Yang, D. Krompass, and V. Tresp, “Tensor-train recurrent neural networks for video classification,” in Proc. Int. Conf. Mach. Learn., 2017, pp. 3891–3900.\n\n[3] Qi, J., Yang, C.H.H., Chen, P.Y. and Tejedor, Exploiting Low-Rank Tensor-Train Deep Neural Networks Based on Riemannian Gradient Descent With Illustrations of Speech Processing. IEEE/ACM Transactions on Audio, Speech, and Language Processing, 31, pp.633-642, 2023.\n\n2. The authors do not show theoretical support to explain the low-rank TTN for the successful application in this work. The authors can refer to the recent theoretical work of Qi et al. on TTN as an encoder exhibiting low-dimensional representations of input features, where many theoretical bounds are provided in the contexts of end-to-end TTN and pre-trained TTN models. \n\n[4] Qi, J., Yang, C.H.H., Chen, P.Y. and Hsieh, M.H., Theoretical error performance analysis for variational quantum circuit based functional regression. npj Quantum Information, 9(1), p.4, 2023. \n\n[5] Qi, J., Yang, C.H.H., Chen, P.Y. and Hsieh, M.H., Pre-training Tensor-Train Networks Facilitates Machine Learning with Variational Quantum Circuits. arXiv preprint arXiv:2306.03741, 2023. \n\n3. The title is \"low-rank tensor-network encoding,\" but the authors only discuss the TTN for feature extraction. If the authors use the tensor network, the authors need to compare other tensor decomposition methods like CP and Tucker decompositions, e.g., the CP decomposition for CNN below:\n\n[6] Lebedev, V., Ganin, Y., Rakhuba, M., Oseledets, I. and Lempitsky, V., 2014. Speeding up convolutional neural networks using fine-tuned cp-decomposition. arXiv preprint arXiv:1412.6553.\n\n4. The speed-up achieved by TTN on GPUs is not fully explained. Different setups of TT ranks can greatly affect the running speed but with a compromise of lowering the empirical performance. The compromise of the efficiency of TTN and its empirical performance should be further discussed in both theory and experiments.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "In this work, the authors leverage a tensor-train network as the feature encoder dealing with raw video data in the context of video-to-action behavior cloning. The use of TTN for feature extraction is interesting in the new applied scenarios, and the authors' experiments could verify the effectiveness of TTN.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "(1) To improve the quality of this paper, more theoretical explanations are expected for something uncommon using TTN for the new application. It is unclear why the TTN can be effectively applied in video-to-action behavioral cloning. \n\n(2) More low-rank tensor methods are expected to be attempted. \n\n(3) More discussions on the efficiency of low-rank TTN on GPUs should be expected.", "broader_impact_concerns": "There is no ethical issue in this work.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper proposed a tensor-network latent-space encoding approach to improve video-to-action game playing based on behavioral cloning. 

In the initial round, the reviewers' comments are concentrated on the clarity of the presentations and the lack of broad comparisons to existing methods. The authors' rebuttal has addressed most of the reviewers' concerns. In the final recommendations, the remaining concerns are
- [Not an issue] lacking technical novelty in Tensor-Decomposition: Given the review criteria of TMLR and the application of video-to-action game playing with tensor network is new, I don't consider it as a weakness.
- [Major] inference cost: two reviewers demanded more analysis on the inference cost, in addition to the training cost.
- [Minor] comparison to other representation learning methods: to demonstrate the effectiveness of the proposed embedding method, the authors should add comparisons to other representation learning methods, such as contrastive learning. Though the reviewers also acknowledged that the current results already show performance advantages. 

Based on my own reading and the reviewers' suggestions, I recommend accept with minor revision. The authors should at least address the major issue (second bullet point) in the revised version.
Decision: Accept (poster)