S2P: State-conditioned Image Synthesis for Data Augmentation in Offline Reinforcement Learning
Abstract
Offline reinforcement learning (Offline RL) suffers from the innate distributional shift as it cannot interact with the physical environment during training.
To alleviate such limitation, state-based offline RL leverages a learned dynamics model from the logged experience and augments the predicted state transition to extend the data distribution.
For exploiting such benefit also on the image-based RL, we firstly propose a generative model, S2P (State2Pixel), which synthesizes the raw pixel of the agent from its corresponding state.
It enables bridging the gap between the state and the image domain in RL algorithms, and virtually exploring unseen image distribution via model-based transition in the state space.
Through experiments, we confirm that our S2P-based image synthesis not only improves the imagebased offline RL performance but also shows powerful generalization capability on unseen tasks.
Introduction
[Figure]
Out [Figure]
ExtendedDistribution
Deep learning algorithms have shown significant development thanks to the large pre-collected dataset, such as
SQuAD [49] in natural language processing (NLP) and
ImageNet [4] in computer vision.
On the contrary, reinforcement learning (RL) requires an online trial-and-error in training process to collect the data by interacting with the environment which hinders its utilization in many real-world applications.
Due to this intrinsic property of current online RL algorithms, there exist some approaches that try to deploy large and diverse pre-recorded datasets without online interaction with the environment, which is called offline RL.
However, recent studies have observed that the current online RL algorithms [10, 36] perform poorly in an offline setting.
It is primarily attributed to the large extrapolation error when the Q-function is evaluated on out-of- distribution actions, which is called the distributional shift [26, 22, 8].
That is, due to the offline setting that limits online data collection, the offline RL has struggled to generalize beyond the given offline dataset.
Even though some offline RL methods [27, 25, 56] achieve reasonable performances in some settings, their training is still limited to behaviors within the given offline dataset distribution, and detours the evaluation on out-of-distribution data rather than directly addressing such an empty space of the offline dataset.
Thus, there exists a growing need for the development of algorithms specialized to directly address such out-of-distribution data by extending the offline dataset distribution’s support.
To alleviate the fixed dataset distribution problem, recent studies propose some data augmentation strategies.
In state-based RL, model-based algorithms [2, 1, 28, 16] which learn a dynamics model from the pre-recorded dataset and augment the dataset with generated state transitions have emerged as a promising paradigm.
As the model-based approach trains dynamics models in a supervised manner, it allows a stable training process and generates reliable state transition data for augmentation.
Thus, it can be a plausible choice that it enables the generalization into the unseen state-action by performing dynamics-consistent planning on unseen state distribution.
When it comes to the image domain, however, there is still no augmentation strategy to mitigate the aforementioned distribution shift.
Even though some model-based image RL methods [12, 11] that propose to learn latent dynamics using reconstruction error from ELBO objective [18, 53, 24] can be exploited for generating image transition data in a similar manner to the state-based methods, the quality of the output images from these approaches is not satisfactory because 1) their focus is on learning latent representation suitable for the RL network’s inputs rather than generating high-quality and accurate images, and 2) ELBO-based objective cannot generate photo-realistic images compared to other generative models such as GAN [9] or diffusion [14, 43] and it usually produces blurry outputs.
Above all, 3) model-based image RL only exploits image input and it makes the generative model fail to capture the accurate dynamics and the details of the agent’s posture or objects in the image [42, 44, 3].
These undesired properties discourage offline RL algorithms from adding the reconstruction output of model-based image RL to their training data as an augmentation strategy.
Therefore, we propose S2P (State2Pixel) which utilizes multi-modal input (the state, and the previous image observation of the agent) to synthesize the realistic image from its corresponding state.
The key element of S2P is a multi-modal affine transformation (MAT) which effectively exploits both state and image cross-modality information.
Unlike previous learned affine transformation [19–21, 46, 35], which leverages a single domain input, MAT fuses the cross-modal representation from the state and the image to produce the scale and the bias modulation parameters.
This multi-modality of S2P makes it possible to generate the dynamics-consistent images from the reliable state transition while preserving high-quality image generation capability.
To sum up, our work makes the following key contributions
• We propose a state-to-pixel generative model (S2P) which generates dynamics-consistent image and multimodal affine transformation (MAT) module for aggregating cross-modal inputs.
• To the best of the author’s knowledge, this work is the first to propose image augmentation for offline image RL and overcome innate fixed distribution problem by implicitly leveraging reliable state transition.
• We evaluate our S2P on the DMControl [54] benchmark environments with the offline {+setting, and it results in 1 . 
+}
− 3 .
0 x higher offline RL performance by augmenting the generated synthetic image transition data. • Even with the state distributions of the unseen tasks, S2P can generalize to unseen image distribution, and we show that the agent can be trained by offline RL with these generated images only.
Related Work 2.1 Image Synthesis
Generative Adversarial Networks (GAN) [9] based deep generative models enjoy huge success in synthesizing high-resolution photo-realistic images via style mapping function and the learned affine transformation [19–21, 46, 35].
StyleGAN [19] firstly utilizes a style vector w and Adaptive Instance
Normalization (AdaIN) [5, 15] in the generative networks to disentangle the latent space and control the scale-specific synthesis.
Following studies [20] pinpoint that the AdaIN operation which leads to information loss in the feature magnitude makes undesired droplet-like artifacts in the synthesized images and proposes weight demodulation by assuming the variance of input features.
SPADE [46] proposes an architecture to synthesize the image using its corresponding semantic masks and spatially learned affine transformation.
ManiGAN [35] suggests Text-Image Affine Combination Module (ACM) which enables the network to manipulate the images using text descriptions given by users.
The difference between our proposed MAT and the previous studies is that we leverage cross-modal data, state and image, to estimate the modulation parameters for the learned affine transformation whereas other studies use a single data type such as text or image.
2.2 Offline Reinforcement Learning & Data Augmentation
Offline RL [6, 30, 34] is the task of learning policies from a given static dataset, which is different from online RL that learns useful behaviors through trial-and-error in the environment.
Prior offline RL algorithms are designed to constrain the policy to the behavior policy used for offline data collection via direct state or action constraints [8, 38], maximum mean discrepancy [26], KL divergence [56, 61, 17], or learning conservative critics [27, 25].
However, most of these methods are limited to exploiting the state-action distribution of the given static dataset, rather than exploring and extending the distribution.
As the offline setting prohibits online interaction with the environment, we suggest the synthetic data generation method to enable the offline RL agent to virtually explore and extend the distribution by leaving the support of the dataset.
Recent works in model-based state RL that involves learning a policy with a dynamics model [22, 2, 1, 29, 59, 60] suggest the need to augment the data with generated transitions from the model for extending the data distribution.
Also, augmentation strategies in image domain [31, 51, 32, 57, 58] emphasize the importance of image augmentation for sample efficiency and robust representation learning.
But, these works focus on purely image manipulation techniques on the given image such as cropping rather than generating image transition.
Some image-based methods [47, 11, 12] that use the variational model to train the latent dynamics are studied in a similar concept to the state-based ones.
However, the generated image from these methods is the byproduct of learning the effective image representation rather than the main purpose of these works, and as these methods only utilize the image input, it leads to missing objects or inaccurate dynamics of the agent in the reconstructed image.
Therefore, to bridge the gap between the model-based state transition data augmentation techniques and image augmentation, we propose the method that generates dynamics-consistent image transition data along timesteps with multi-modal inputs.
Method 3.1 S2P Generator
Architecture.
The goal of S2P (State2Pixel) generator is to synthesize the image ˆ I t which perfectly represents all the information of its corresponding state s t .
Unfortunately, a state-sole condition cannot formulate a single deterministic rendered image because, in most cases, state does not provide the agent’s position from the global coordinate, but rather from an egocentric coordinate.
Also, image-based RL algorithms utilize sequential images to capture the agent’s velocity using the change of the background, e.g. ground checkerboard, between input images.
It means that the image of the current step I t is dependent not only on the current state s t , but also on the image of the previous step I t − 1 .
We, therefore, build the generator G to synthesize the image ˆ I t from both the state s t and the previous image I t − 1 so that the generated image ˆ I t can preserve the dynamics-consistency in the physical environment.
[Equation]
At the first layer of S2P, the input image I t − 1 and the state s t are projected to the feature space via convolution and MLP encoders respectively.
Both features are then fed to the hierarchical generator block which consists of several residual connections [13] and the upsampling layer.
After passing through each generator block, the spatial size of the feature map is doubled while the channel dimension is halved.
The image features are converted to the RGB images at the last layer of the generator with a single convolution layer.
[Figure] We observe that the input signals, i.e. state s t and image I t − 1 , become attenuated as they pass through deeper generating layers and the network produces images with poor quality.
So, similar to recent style-based image synthesis algorithms [19–21, 46, 35], we adopt a learned affine transformation architecture to inject auxiliary signals to the generator.
The difference between the previous stylebased generative models and S2P is that we propose a multimodal affine transformation (MAT) to produce the learnable modulation parameters, γ and β , with the cross-modality representation via a multimodal feature extractor and state-to-latent mapping function.
The overall architecture of our proposed S2P is depicted in Figure 2.
A non-linear latent mapping function g : S → W which is implemented as an 8-layer MLP produces a latent code w ∈ W from the given state s in the state space S .
The latent code w is spatially expanded as the same size of the input feature of MAT module x , and the conditioned image I t − 1 is also linearly interpolated to make its size same as x .
The spatially expanded w and the resized I t −are channel-wisely concatenated and fed to the multimodal feature extractor to fuse the state and image cross-modality representation.
Each estimator then produces the learnable scale γ and bias β for effective cross-modal affine transformation.
Finally, our proposed MAT operation is defined as:
[Equation]
where µ c ( · ) and σ c ( · ) are the channel-wise mean and the standard deviation of the input feature of
MAT x i from the i th block of the generator, and ⊙ denotes Hadamard element-wise product.
The design of MAT is illustrated in Figure 3.
In addition, it is shown that the neural network is biased toward learning a low frequency mapping and has difficulty in representing a high frequency information [48, 40].
To mitigate such undesired tendency, we do not use naïve state vector s as input, but employ a positional encoding with the high frequency function ψ : R → R 2 L which is defined as:
[Equation]
where x indicates each component of state vector s .
We utilize multi-scale discriminators in [55] to increase the receptive field without deeper layers or larger convolution kernels for alleviating the overfitting.
Two discriminators with identical architecture are adopted during training, and the synthesized and real images which are resized to several spatial sizes are fed to each discriminator.
Loss Function.
Our S2P generator is trained by linearly combined multiple objectives.
First, we leverage a pixel-wise L 1 loss between the output of the generator ˆ I t and the real image I t ,
[Equation]
[Equation] [Figure] In addition to the pixel-wise loss, we also utilize the ImageNet [50] pre-trained VGG19 [52] network to calculate the perceptual similarity loss L per between two images,
[Equation]
where ϕ ( i ) denotes the i-th layer of VGG19.
We implement the adversarial objective L adv for both S2P generator and multi-scale discriminators same as pix2pixHD [55].
The difference is that we replace the least square loss with the hinge-based loss [37] and we condition state information to the discriminator D so that the generator is induced to produce the dynamics-consistent outputs.
[Equation]
The total loss function to optimize the S2P generator can be defined as:
[Equation]
where λ 1 , λ 2 , and λ 3 are the hyperparameters to balance among the objectives.
3.2 Offline reinforcement learning with synthetic data
We consider the Markov decision process (MDP) M
= ( I , S , A , T, r, ρ 0 , γ ) , where I denotes the image space, S the state space corresponding to I , A the action space, T ( s ′ | s, a ) the transition dynamics, r ( s, a ) the reward function, ρ 0 the initial distribution, and γ the discount factor.
We denote the discounted image visitation distribution of a policy π using d π M
( I ) :
= (1 − γ ) (cid:80) ∞
t =0 γ t P ( I t =
I | π ) , where P ( I t
= I | π ) is the probability of reaching image observation I at time t by using π in
M .
Similarly, we denote the image-action visitation distribution with d π M ( I, a ) :
= d π M ( I ) π ( a | I ) .
The objective of RL is to optimize a policy π ( a | I ) that maximizes the expected discounted return
J ( π ) =
− γ E ( I,a ) ∼ d π M ( I,a )
[ r ( s, a )] .
In the offline RL setting, the algorithm has access to a static dataset D = { ( I t , s t , a t , r t , I t +1 , s t +1 ) }
Nt =0 collected by unknown behavior policy π β .
In other words, the dataset D is obtained from d π β M ( I, a ) :
= d π β M ( I ) π β ( a | I ) and the goal is to find the best possible policy using the static dataset without online interaction with the environment.
To utilize the dynamics-consistent transition data for augmentation in offline RL, we take the modelbased approach that trains an ensemble of dynamics and reward model
ˆ T θ ( s ′ , r | s, a ) , which outputs the predicted next state, and reward ˆ r
( s, a ) .
Once a model has been learned, we can construct the learned MDP (cid:99) M
= ( I , S , A , ˆ T, ˆ r, ρ 0 , γ ) , which has the same spaces, but uses the learned dynamics and reward function.
Naively optimizing the RL objective with the (cid:99) M is known to fail in the offline RL setting, both in theory and practice [22, 59], due to the distribution shift and model-bias.
To overcome these, we take an uncertainty estimation algorithm like bootstrap ensembles [45, 39], and obtain u ( s, a ) , an estimate of uncertainty in dynamics.
Then we could utilize the uncertainty penalized reward ˜ r ( s, a ) =
ˆ r
( s, a ) − λu ( s, a ) , where λ is a hyperparameter.
We consider the following uncertainty quantification method that uses the maximum learned variance over the ensemble, u ( s, a ) = max i =1 ,...
N || Σ iθ ( s, a ) || F [59, 22].
As a final process, we train offline RL with the following hybrid objective : J ( π ) =1 − γ E
( I,a ) ∼ d f ( I,a ) [ r ( s, a )] , where d f ( I, a ) =
fd π β M ( I, a ) + (1 − f )
d η (cid:102)
M ( I, a ) , f ∈ [0 , 1] is the ratio of the datapoints drawn from the offline dataset D , and η ( ·| s ) is the state rollout distribution used with the trained dynamics ensemble model ˆ T θ , and (cid:102)
M is same as (cid:99) M except that the reward is ˜ r ( s, a ) instead of ˆ r ( s, a ) .
Samples from d η (cid:102) M ( I, a ) can be obtained by rollout η in (cid:102) M and convert the obtained state transitions { ( s t , a t , r t , s t +1 ) }
Nt =0 into { ( I t , a t , r t , I t +1 ) }
Nt =0 by using the trained S2P [Table] generator in Section 3.1.
For implementation [Table] we collect synthetic image transition data in separate replay buffer D model and train the agent by any offline RL algorithms with the sampled mini-batches from D and D model by the ratio of f and 1 − f .
The overall algorithm and more training details are summarized in Appendix B
[Table]
Experiments 4.1 Environments & Data collection
We evaluate our method on a large subset of the dataset from the DeepMind Control (DMControl) suite [54].
It includes 6 environments, which were typically used for online image-based RL benchmarks.
However, to the best of our knowledge, all of these environments have never been properly evaluated in an offline image-based RL setting.
The datasets in these benchmarks are generated as follows:
random : rollout by a random policy that samples actions from a uniform distribution.
mixed : train a policy using state-based SAC [10] until 500k steps for finger, cheetah and 100k steps for the others, then randomly samples trajectories from the replay buffer.
500k, 100k steps are minimum steps required to reach the expert level performance for each task.
expert : train a policy using state-based SAC.
After convergence we collect trajectories from the converged policy.
4.2 Offline Reinforcement Learning
To validate whether the S2P can help improve the offline RL performance, we evaluate our algorithm with recent offline RL algorithms like CQL [27] which utilizes conservative training of the critic, and
IQL [25] which trains critic by implicitly querying actions near the distribution of the dataset, and
SLAC-off [33] which is state-of-the-art online image-based RL algorithm.
We use this SLAC-off with the offline setting.
We also compare the policy constraint-based offline RL algorithms like
BEAR [26], and behavior cloning, BC .
The results for these two algorithms are in Appendix B. To extend the offline RL into the image-based setting, we follow the image encoder architecture from [33] and train a variational model using the offline data.
Then, we train in the latent space of this model.
We include the offline RL results on different types of environments and data in Table 1.
The results on the originally given offline dataset (50k) are shown in the left column of each algorithm, and the results on the S2P-based augmented dataset are shown in the right column of each algorithm.
The
S2P augments the same amount of the original offline dataset (50k).
As it physically has more data (100k) than the original dataset (50k), for a reference, we also include the results on the 100k offline dataset in the parenthesis in Table 1.
Overall, the S2P-based method achieves better performance than the 50k dataset, even exceeding the 100k dataset’s score in some tasks.
As the S2P answers how to generate the image transition data, we also have to consider where to generate the image transition data through S2P. Specifically, we use a random policy as η ( ·| s ) in the mixed, expert dataset, and a policy trained by the state-based offline RL as η ( ·| s ) in the random dataset.
These strategies are considered due to the following two assumptions.
Firstly, as the random dataset only has random behavior, the dataset may not have any meaningfully rewarded states even with the augmentation by the random policy, especially in the locomotion environment (e.g. cheetah, walker cannot leave the initial states by the random policy).
As the S2P’s objective is to extend the distribution of the given dataset, the policy trained in an offline manner can help leave the support of the dataset compared to the naive random policy.
Secondly, as the non-random datasets may have relatively biased state-action distributions that receive meaningful rewards compared to the random dataset, it is difficult to get out of the support of these datasets by the trained policy as most of the state-action induced by the trained policy are included in the similar distribution of these datasets.
However, the random policy can be effective as it can bring some exploration effects like noise injection or increasing entropy in [36, 10].
[Table]
To verify such assumptions, we analyze the effect of different types of rollout distribution η ( ·| s ) in offline RL performance (Table 2).
We denote 50k dataset as the results of the given original offline dataset, and +S2P(random η ) as the results of the data augmented by rollout with the random actions, and +S2P(offRL η ) as the results of the data augmented by rollout with the state-based offline RL policy.
As expected, the random policy is more effective in the expert dataset.
Also, we could find the opposite phenomenon in the random dataset, and trade-offs between these two strategies in the mixed dataset.
4.3 Comparison with model-based image RL
To show why the multi-modal inputs are necessary for augmenting the image transition data in offline image RL, we compare our S2P with the previous model-based image RL algorithm, Dreamer [12], as it can also reconstruct the images of the agent by training the reconstruction error from ELBO objective only using uni-modal inputs (previous images of the agent).
For comparison, we trained
Dreamer in an offline manner with the same dataset used for training S2P, and predicted future images with the episode context obtained from 5 consecutive ground truth images.
Even though the
Dreamer saw more previous steps’ images compared to S2P (only a single image of the previous step), it still has difficulty in generating accurate posture, which supports the S2P’s advantages (Fig 4).
It is because the priority of the model-based image RL algorithms is learning the effective latent representation for RL tasks, and they do not utilize a broader source of supervision from the state inputs.
To prove that state-inconsistent images from the model-based image RL method cannot improve offline RL performance at the S2P level, we perform the same experiment in Table 1, but replace the augmented images from S2P with images from Dreamer.
We observe that the image augmentation from Dreamer even degrades the original RL performance in several tasks, and the performance improvements with S2P excels the augmentation from Dreamer by a large margin in all baselines (Table 3).
Thus, we could say that the inaccurate posture and quality of the images generated by the model-based method trained with the uni-modal inputs (images) are not sufficient for augmenting image transition data in the offline setting.
More experiments and details of the augmentation are in the Appendix.
t t = 3 t = 1 (a)
GT
Baseline + PE + SAT + MAT (b)
Cheetahjump Walkerrun GT NoRef (c) t =t =
t = 3 t =
[Figure]
dataset method 50k
[Table]
4.4 Ablation
To observe how each component of the S2P contributes to the quality of the synthesized images, we perform an ablation study on the model architecture of the S2P and show its qualitative results on
Figure 5(a).
A baseline architecture without any contribution of our proposed method, i.e. positional encoding (PE) and multimodal affine transformation (MAT), shows the worst image quality.
We report that a simple application of the high frequency function ψ to the input state s t (PE) results in a noticeable increase in image quality.
We also address the necessity of the input image I t − 1 to estimate the modulation parameters, γ , β , for the learned affine transformation.
We ablate the input I t − 1 in
MAT and utilize only the spatially expanded input state which is called State Affine Transformation (SAT).
The generator which replaces the MAT module with SAT has difficulty in exploiting dynamics information from the previous image I t − 1 and the translation error is accumulated as the generator recurrently synthesizes the long-horizon trajectory.
We can observe such dynamics inconsistency especially in the locomotion tasks, e.g. cheetah, and walker, which expresses the velocity in images by the change of the checkerboard in the ground.
Compared to SAT, our proposed MAT which leverages both the state s t and the previous image I t − 1 shows better performance in reconstructing not only the posture of the agent but also its dynamics-consistent background.
4.5 Zero-shot Task Adaptation
To validate whether the S2P can help the offline RL process in settings that require generalization to tasks that are different from the given dataset, we construct two environments cheetah-jump and walker-run .
In cheetah-jump , which is referred from [59], the agent should solve a task that is different from the original purpose of the behavior policy.
Specifically, we relabel the rewards in the cheetah-run-mixed dataset to reward the cheetah to jump as high as possible.
Then, we generate the image transition data from the states whose z positions are bigger than a threshold to validate the S2P-based augmentation’s advantage in tasks that require generalization.
By training with the relabeled reward, the agent with the S2P achieves a higher return and learns to bounce back and forth to take a leap higher (Figure 5(b)), even though the batch data contain little jumping motion.
To investigate whether the S2P can generate unseen image distributions from unseen state distribution, we collect walker-run state dataset by the same way of other mixed datasets in Section 4.1.
Then, we generate images from these unseen states by recurrently using the S2P generator and we apply offline
RL on these generated image transition data.
Even though the state distribution is totally different from walker-walk dataset as the agent should run instead of walk, not only the S2P successfully generalizes to the unseen image distributions (Figure 5(b)), but also the agent can be trained to run by offline RL only with these synthesized images beyond the expertise of the dataset (Table 4).
More details are in Appendix B.
[Figure]
returns of the walker-run and cheetah-jump tasks.
We include mean undiscounted return of the episodes in the batch data for comparison.
We attribute such satisfactory task generalization to the model architecture of the S2P which leverages both I t − 1 and s t for synthesizing the I t .
The posture of the agent is deterministic with the sole-state condition and the background of the image such as the ground checkerboard is dependent both on the state and the previous image.
Therefore, our S2P exploits the state input to generate the posture of the agent and exploits the image input to generate the background.
It is well shown in Figure 5(c) where we intentionally replace all the image input with the zero matrix O H × W × 3 during the inference phase.
We observe that S2P still perfectly reconstructs the posture of the agent with its corresponding state only while it fails to recover the background and the ground checkerboard as we expected.
Conclusion
We firstly present the state-to-pixel (S2P) algorithm that synthesizes the raw pixel from its corresponding state and the previous image.
As the augmentation paradigm of the S2P is generating dynamics-consistent image transition data, we demonstrate that S2P not only improves the imagebased offline RL performance but also shows powerful generalization capability on unseen tasks.
Even though S2P shows promising results in offline RL by data augmentation techniques, the assumption that datasets consist of pairs of images and states is still a strong assumption.
Thus, for future work, we plan to further extend the idea with more relaxed assumptions such as unpaired datasets or extend other state-based applications to the image-based RL algorithms.
Acknowledgement
This research was supported by Future Challenge Program through the Agency for Defense Development funded by the Defense Acquisition Program Administration and by Institute of Information & communications Technology Planning & Evaluation (IITP) grant funded by the Korea government(MSIT)
[NO.2021-0-01343, Artificial Intelligence Graduate School Program (Seoul National
University)]
Checklist
For all authors... (a) Do the main claims made in the abstract and introduction accurately reflect the paper’s contributions and scope?
[Yes] See Section(b)
Did you describe the limitations of your work?
[Yes] See Section(c)
Did you discuss any potential negative societal impacts of your work?
[Yes] This work inherits the potential negative societal impacts of reinforcement learning.
We do not anticipate any additional negative impacts that are unique to this work.
(d) Have you read the ethics review guidelines and ensured that your paper conforms to them?
[Yes] 2.
If you are including theoretical results... (a) Did you state the full set of assumptions of all theoretical results?
[N/A] (b) Did you include complete proofs of all theoretical results?
[N/A] 3.
If you ran experiments... (a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)?
[Yes] We included code in the supplemental materials.
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)?
[Yes] See Section 3 and Appendix (c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?
[Yes] See Section(d)
Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)?
[Yes] See Appendix 4.
If you are using existing assets (e.g., code, data, models) or curating/releasing new assets... (a) If your work uses existing assets, did you cite the creators?
[Yes] See Section(b)
Did you mention the license of the assets?
[Yes] See Appendix.
(c) Did you include any new assets either in the supplemental material or as a URL?
[Yes] We include the code & model in the supplemental materials.
(d) Did you discuss whether and how consent was obtained from people whose data you’re using/curating?
[N/A] (e) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content?
[N/A] 5.
If you used crowdsourcing or conducted research with human subjects... (a) Did you include the full text of instructions given to participants and screenshots, if applicable?
[N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable? [N/A] (c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation?
[N/A]
A Image generation A.1 Architecture specification
The generator architecture is implemented with several MAT Residual Blocks followed by bilinear upsampling as shown in Figure 6(b).
The architecture of the residual block is largely implemented from Park et al.
[46] where we replace the SPADE module with MAT in SPADE Residual block. We also leverage the Spectral Normalization [41] to all the convolution layers in the generator.
The latent mapping network g for generating the latent code w is implemented with an 8-layer of MLP with channel-wise normalization at the first layer, which is the same as the style mapping function in StyleGAN The first MLP layer of the latent mapping function is different from the physical environment as the number of the input state, ns , is different. The dimension of the input state is increased 2 L times by a high frequency positional encoding function ψ .
We set the hyperparameter
L as 10 and concatenate the input and output of ψ , which leads to 21 times increase in the channel dimension.
The specification of the entire architecture is shown in Figure 7.
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
A.2 Training details
We implement our proposed generator architecture with the public deep learning platform PyTorch and train it on a single NVIDIA RTX A6000 GPU.
We train 30 epochs for each task and the generated image size is set to 128 × 128 .
An Adam optimizer [23] with the learning rate of 0.0002 is utilized and the batch size is set to 16.
3x3 Conv (3, 512)
Linear (512, 512)
MAT Residual Block (1024, 1024)
MAT Residual Block (1024, 1024)
MAT Residual Block (1024, 1024)
MAT Residual Block (1024, 512)
MAT Residual Block (256, 128)
MAT Residual Block (512, 256)
MAT Residual Block (128, 64) 3x3 Conv (64, 3), Tanh 4x4x
Upsample (2x)
Upsample (2x)
Upsample (2x)
Upsample (2x)
Upsample (2x) 128x128x
Output shape 4x4x 8x8x 8x8x 16x16x 16x16x 32x32x 32x32x 64x64x 64x64x 128x128x 128x128x 128x128x 8x8x t =t =
t =t =
t = 5 t =
[Figure]
A.3 Additional results of image synthesis
We report additional qualitative results on the multiple environments of DMControl in Figure 8.
We show that our proposed S2P generates high-quality images regardless of the environment.
We recurrently generate a single trajectory exploiting the current state and the previous images which are also the output of the generator from the previous state.
Also, we evaluate the quality of generated images quantitatively with metrics (FID score, LPIPS, PSNR, SSIM) which are frequently adopted
[Table]
for evaluating image quality in Table 6.
Our S2P outperforms Dreamer in all the quantitative results which indicates the generated image from S2P has better quality than from Dreamer.
B Offline RL Experiments Details B.1 Algorithm
Input: offline dataset D , state rollout distribution η ( ·| s ) .
Train the probabilistic dynamics models ˆ T θ ( s ′ , r | s, a ) =
N ( µ θ ( s, a ) , Σ θ ( s, a )) on D .
Train the image generator G ( s t , I t − 1 ) on D .
for i =
to K do Randomly sample I 0 , s 0 ∼ D .
Get τ s ∼ ( s 0 , a 0 , r 0 , s 1 , ..., s M ) by using the η ( ·| s ) and
ˆ T θ ( s ′ , r | s, a ) .
Generate τ I ∼ ( I 0 , a 0 , r 0 , I 1 , ...
I M ) from τ s by G .
Save τ I in D model .
for
Apply any offline RL algorithm with the dataset sampled from D, D model with the ratio of f , 1 − f .
B.2 Ablation studies B.2.1 Uncertainty types
We conduct experiments on how the choice of the uncertainty affects the performance.
We denote u ( s, a ) = max i =1 ,...
N || Σ iθ ( s, a ) || F as Max Var , u ( s, a ) = max i =1 ,...
N || µ iθ ( s, a ) −
N (cid:80) Nj =1 µ jθ ( s, a ) || 2 as Ens Var , and average value of both uncertainty types as Average Both .
In practice, we found that Max Var achieves better performance compared to other types of uncertainty in mixed, expert dataset, while Ens Var achieves better at random dataset.
We hypothesize that the dynamics model is quite uncertain in the aspect of the Max Var on the random dataset due to the excessive randomness of the data.
Thus, it could induce excessive penalty on the predicted reward when Max Var is used, and the agent could become too conservative.
B.2.2 Rollout horizons
To validate the effectiveness of different rollout horizons, we conduct experiments with horizon length 1 ( +S2P (1step) ) and 5 ( +S2P (5step) ), while following the same rollout strategies in Section 4.2.
M ETHOD
[Table]
V AR
[Table]
For the 5 step case, as the proposed S2P generator is conditioned on the previous timestep’s image to synthesize the next image, we recurrently generate the image transition data.
That is, at the first timestep
[Table]
the ground truth image is conditioned on the image generator
[Table]
but after then, the generated image is conditioned on the image generator for generating the following timestep’s image.
The results shown in Table 8 represent that the augmentation with a longer horizon also has advantages in offline RL
[Table]
but a short horizon is more effective overall.
This is due to the uncertainty accumulation effect as shown in Figure 9.
The average uncertainty grows as the rollout horizon increases due to the model bias
[Table]
and it leads to more penalties on the predicted rewards, which can induce a too conservative agent.
[Table]
B.2.3 Results on policy constraint-based methods
We additionally test the proposed method on the policy constraint-based methods such as BEAR [26] and behavior cloning (BC) on the cheetah-run environment.
As shown in Table 9, the performance of these methods with augmented data is worse than the results of non-augmented data.
The poor performance is reasonable, because BEAR utilizes the action distribution’s support matching by
MMD, and BC is trained with maximizing the likelihood of the action.
As the augmented action distributions are totally different from the behavior policy that induces the offline dataset, these two methods perform poorly because these methods try to clone or match both types of actions.
That is, these methods try to clone or match the support of the given offline dataset and sampled actions that could have different distribution or support.
[Table]
B.3 Additional experiments on Dreamer and conventional image augmentation technique
Dreamer with larger dataset.
To validate whether the performance degradation by augmentation from Dreamer (Section 4.3) stems from the small dataset as the Dreamer requires more than 100k samples in online training, we collect an additional 250k dataset and augment the image transition amount of 250k by Dreamer (totally 500k datasets), and perform the same experiment. Despite the bigger size dataset, the performance does not increase overall (Table 10), and we could say that the inaccurate posture and quality of the images generated by the dreamer are attributed to the limited source of supervision from the uni-modal inputs rather than the size of the dataset, and it is not that proper for data augmentation in the offline setting. Comparison with the conventional image-augmentation technique. To validate why S2P is needed instead of the conventional image-augmentation method, we additionally experiment with random crop and reflection padding, which is frequently used in image representation learning and online image RL.
We perform the same experiment in Table 1, but replace the augmented images from S2P with randomly cropped images.
Slightly better performance than Dreamer could be interpreted as it just manipulates the given true images (while maintaining the accurate posture of the agent) rather than generating new images like Dreamer (Table 10).
But it still has difficulty in surpassing the S2P’s results as it cannot deviate from the state-action distribution of the offline dataset.
[Table]
+R EFLECT +D REAMER +D REAMER [Table] B.4 Visualization of the image distribution To validate whether the S2P really affects the distribution of the offline dataset
[Table]
we visualized the cheetah-run-expert dataset and the S2P-based augmented dataset by the random policy by applying the t-sne (Figure 10)
[Table]
As shown in Figure
[Table]
the augmented dataset not only occupies a similar area to the original dataset but also encloses the original dataset, even connecting the clusters of the original dataset.
It can be interpreted that S2P connects different modes of image distribution by deploying virtual exploration in the state space.
B.5 Model learning
In state space, we represent the dynamics, reward model as a probabilistic neural network that outputs a Gaussian distribution over the next state and reward given the current state and action.
[Equation]
We train an ensemble of N dynamics models { ˆ T θ = N ( µ θ ( s, a ) ,
Σ θ ( s, a )) } Ni =1 , with each model trained independently via maximum likelihood estimation on offline dataset D .
We set the number of dynamics model N as 7 following Yu et al.
During model rollouts, we randomly pick one model.
Each model in the ensemble is represented as 3 layers with 256 hidden units and relu activation function.
B.6 Representation learning
For image-based offline RL process, we follow Lee et al.
[33] that uses a variational model with the following components:
[Equation]
encoder
[Equation]
Posterior : z t
[Equation]
Prior : z t ∼ p ψ ( z t | z t
[Equation]
, a t
[Equation]
Image decoder : I t ∼ D
[Equation]
z
[Equation]
Policy : a t ∼
[Equation]
Full derivation of those equations are in Lee et al.
We train the representation model using the evidence lower bound :
[Equation]
where τ is the number of sequences, and z is the latent representation.
We set τ = 8 and the dimension of z is 288, same as the original implementation.
For image encoder, we use 6 convolutional neural network with kernel sizes [5, 3, 3, 3, 3, 4] and strides [2,2,2,2,2,2] respectively with leaky relu activation function.
The decoder is constructed in a symmetrical manner to the encoder.
We pre-train the image encoder by 300k steps and use the trained encoder’s weights as initial weights when offline
RL training proceeds.
B.7 Training Detail
We use the batch size of 128 and the Adam optimizer for training the value function Q and policy π with the learning rates 0.0003, 0.0001, respectively.
Each Q and π is represented as MLPs with hidden layer sizes (1024, 1024), relu activation function.
We apply tanh activation on the output of the policy with reparameterization trick same as Haarnoja et al.
Also, we use the uncertainty penalty coefficient λ = 2 for all environments except the finger-spin, which use λ = 1 .
The sampling ratio of the offline dataset is f = 0 .
For offline RL implementation, we referred to the original implementations of each work and follow default parameters.
B.8 Zero-shot Task Adaptation Detail
For the cheetah-jump task, we relabel the reward in the given offline dataset to the sparse reward that indicates whether the cheetah is jumping or not.
That is, if z-init z ≥ 0.3, the reward is relabeled as 1, otherwise 0, where z denotes the z-position of the cheetah and init z denotes the initial z-position.
For augmentation, we randomly select states whose z-position is greater than 0.2, and generate images from these states to encourage the jumping motion.
Then, we augment these generated image transition data in the same way of Table 1, and evaluate the offline RL.
For the walker-run task, we generate the dataset in the same manner as the mixed type in Section 4.1.
That is, we train the state-based SAC until convergence, and we randomly sample trajectories from the replay buffer.
Then, we generate ˆ I 1 from s 1 and I 0 , where I 0 is the initial image when the agent is reset, and recurrently generate images ˆ I t +1 from s t +1 , ˆ I t ( t = 0 , 1 , . . .
, N ) by the S2P generator trained with the walker-walk-mixed dataset.
After then, we apply offline RL on these generated image transition data.
B.9 Environment Detail
The DMControl’s environment details and the random and expert scores obtained by training the state-based SAC are shown in Table 11.
The normalized score is computed by 100 ∗ (return - random score)/(expert score - random score), which is proposed in Fu et al.
[Table]
[Figure]
B.10 Results with standard deviation
We included the results with standard deviation in the manuscript (Table 12, Table 13, Table 14,
Table 15)
[Figure]
[Table]
DATASET
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
28 ±
[Table]
[Table]
[Table]
[Table]
.84 ± 2.
95.88 ± 3.
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]