Abstract: While deep neural networks can attain good accuracy on in-distribution test points, many applications require robustness even in the face of unexpected perturbations in the input, changes in the domain, or other sources of distribution shift. We study the problem of test time robustification, i.e., using the test input to improve model robustness. Recent prior works have proposed methods for test time adaptation, however, they each introduce additional assumptions, such as access to multiple test points, that prevent widespread adoption. In this work, we aim to study and devise methods that make no assumptions about the model training process and are broadly applicable at test time. We propose a simple approach that can be used in any test setting where the model is probabilistic and adaptable: when presented with a test example, perform different data augmentations on the data point, and then adapt (all of) the model parameters by minimizing the entropy of the model's average, or marginal, output distribution across the augmentations. Intuitively, this objective encourages the model to make the same prediction across different augmentations, thus enforcing the invariances encoded in these augmentations, while also maintaining confidence in its predictions. In our experiments, we evaluate two baseline ResNet models, two robust ResNet-50 models, and a robust vision transformer model, and we demonstrate that this approach achieves accuracy gains of 1-8% over standard model evaluation and also generally outperforms prior augmentation and adaptation strategies. For the setting in which only one test point is available, we achieve state-of-the-art results on the ImageNet-C, ImageNet-R, and, among ResNet-50 models, ImageNet-A distribution shift benchmarks.

Under review as a conference paper at ICLR 2022 TEST TIME ROBUSTIFICATION OF DEEP MODELS VIA ADAPTATION AND AUGMENTATION

INTRODUCTION Deep neural network models have achieved excellent performance on many machine learning problems, such as image classification, but are often brittle and susceptible to issues stemming from distribution shift. For example, deep image classifiers may degrade precipitously in accuracy when encountering input perturbations, such as noise or changes in lighting  (Hendrycks & Dietterich, 2019)  or domain shifts which occur naturally in real world applications  (Koh et al., 2021) . Therefore, robustification of deep models against these test shifts is an important and active area of study. Most prior works in this area have focused on techniques for training time robustification, including utilizing larger models and datasets  (Orhan, 2019) , various forms of adversarial training  (Sagawa et al., 2020; Wong et al., 2020) , and aggressive data augmentation  (Yin et al., 2019; Hendrycks et al., 2020; Li et al., 2021; Hendrycks et al., 2021a) . Employing these techniques requires modifying the training process, which may not be feasible if, e.g., it involves heavy computation or non-public data. Furthermore, these techniques do not rely on any information about the test points that the model must predict on, even though these test points may provide significant information for improving model robustness. Recently, several works have proposed methods for improving accuracy via adaptation after seeing the test data, typically by updating a subset of the model's weights  (Sun et al., 2020) , normalization statistics  (Schneider et al., 2020) , or both  (Wang et al., 2021; Zhang et al., 2021) . Though effective at handling test shifts, these methods sometimes still require specialized training procedures, and they typically rely on extracting distributional information via batches or even entire sets of test inputs, thus introducing additional assumptions. Figure  1 : A schematic of our overall approach. Left: at test time, as detailed in Section 3, we have a trained model that outputs a probabilistic predictive distribution and has adaptable parameters θ, a single test input x, and a set of data augmentation functions {a 1 , . . . , a M }. Note that we do not assume access to the model training procedure or multiple test inputs for adaptation. We perform different augmentations to x and pass these augmented inputs to the model in order to estimate the marginal output distribution averaged over augmentations. Right: rather than using this distribution to make the final prediction, we instead perform a gradient update on the model to minimize the entropy of this marginal distribution, thus encouraging the model predictions to be invariant across different augmentations while maintaining confident predictions. The final prediction is then made on the original data point, i.e., the predictive distribution in the top right of the schematic. We are interested in studying and devising methods for improving model robustness that are "plug and play", i.e., they can be readily used with a wide variety of pretrained models and test settings. Such methods may simply constitute a different, more robust way to perform inference with preexisting models, under virtually the same assumptions as standard test time inference. In this work, we focus on methods for test time robustness, in which the specific test input may be leveraged in order to improve the model's prediction on that point. Though broad applicability is our primary goal, we also want methods that synergize with other robustification techniques, in order to achieve greater performance than using either set of techniques in isolation. To satisfy both of these desiderata, we devise a novel test time robustness method based on adaptation and augmentation. As illustrated in Figure  1 , when presented with a test point, we propose to adapt the model by augmenting the test point in different ways and ensuring that the model makes the same predictions across these augmentations, thus respecting the invariances encoded in the data augmentations. We further encourage the model to make confident predictions, thus arriving at the proposed method: minimize the marginal entropy of the model's predictions across the augmented versions of the test point. We refer to the proposed method as marginal entropy minimization with ensembled augmentations (MEME), and this is the primary contribution of our work. MEME makes direct use of pretrained models without any assumptions about their particular training procedure or architecture, while requiring only a single test input for adaptation. In Section 4, we demonstrate empirically that MEME consistently improves the performance of ResNet  (He et al., 2016)  and vision transformer  (Dosovitskiy et al., 2021)  models on several challenging ImageNet distribution shift benchmarks, achieving several new state-of-the-art results for these models in the setting in which only one test point is available. In particular, MEME consistently outperforms non adaptive marginal distribution predictions (between 1-10% improvement) on corruption and rendition shifts -tested by the ImageNet-C  (Hendrycks & Dietterich, 2019)  and ImageNet-R  (Hendrycks et al., 2021a)  datasets, respectively -indicating that adaptation plays a crucial role in improving predictive accuracy. MEME encourages both invariance across augmentations and confident predictions, and an ablation study in Section 4 shows that both components are important for maximal performance gains. Also, MEME is, to the best of our knowledge, the first adaptation method to improve performance (by 1-4% over standard model evaluation) on ImageNet-A  (Hendrycks et al., 2021b) , demonstrating that MEME is more broadly applicable on a wide range of distribution shifts.
 RELATED WORK The general problem of distribution shift has been studied under a number of frameworks  (Quiñonero Candela et al., 2009) , including domain adaptation  (Shimodaira, 2000; Csurka, 2017; Wilson & Cook, 2020) , domain generalization  (Blanchard et al., 2011; Muandet et al., 2013; Gulrajani & Lopez-Paz, 2021) , and distributionally robust optimization  (Ben-Tal et al., 2013; Hu et al., 2018; Sagawa et al., 2020) , to name just a few. These frameworks typically leverage additional training or test assumptions in order to make the distribution shift problem more tractable. Largely separate from these frameworks, various empirical methods have also been proposed for dealing with shift, such as increasing the model and training dataset size or using heavy training augmentations  (Orhan, 2019; Yin et al., 2019; Hendrycks et al., 2021a) . The focus of this work is complementary to these efforts: the proposed MEME method is applicable to a wide range of pretrained models, including those trained via robustness methods, and can achieve further performance gains via test time adaptation. Prior test time adaptation methods generally either make significant training or test time assumptions. Some methods update the model using batches or even entire datasets of test inputs, such as by computing batch normalization (BN) statistics on the test set  (Li et al., 2017; Kaku et al., 2020; Nado et al., 2020; Schneider et al., 2020) , or minimizing the (conditional) entropy of model predictions across a batch of test data  (Wang et al., 2021) . The latter approach is closely related to MEME. The differences are that MEME minimizes marginal entropy using single test points and data augmentation and adapts all of the model parameters rather than just those associated with normalization layers, thus not requiring multiple test points or specific model architectures. Other test time adaptation methods can be applied to single test points but require specific training procedures or models  (Sun et al., 2020; Huang et al., 2020; Schneider et al., 2020) . Test time training (TTT)  (Sun et al., 2020)  requires a specialized model with a rotation prediction head, as well as a different procedure for training this model.  Schneider et al. (2020)  show that BN adaptation can be effective even with only one test point, and we refer to this approach as "single point" BN adaptation. As we discuss in Section 3, MEME synergizes well with single point BN adaptation. A number of works have noted that varying forms of strong data augmentation on the training set can improve the resulting model's robustness  (Yin et al., 2019; Hendrycks et al., 2020; Li et al., 2021; Hendrycks et al., 2021a) . Data augmentations are also sometimes used on the test data directly by averaging the model's outputs across augmented copies of the test point  (Krizhevsky et al., 2012; Shorten & Khoshgoftaar, 2019) , i.e., predicting according to the model's marginal output distribution. This technique, which we refer to as test time augmentation (TTA), has been shown to be useful both for improving model accuracy and calibration  (Ashukha et al., 2020)  as well as handling distribution shift  (Molchanov et al., 2020) . We take this idea one step further by explicitly adapting the model such that its marginal output distribution has low entropy. This extracts an additional learning signal for improving the model, and furthermore, the adapted model can then make its final prediction on the clean test point rather than the augmented copies. We empirically show in Section 4 that these differences lead to improved performance over this non adaptive TTA baseline.
 ROBUSTNESS VIA ADAPTATION AND AUGMENTATION Data augmentations are typically used to train the model to respect certain invariances -e.g., changes in lighting or viewpoint do not change the underlying class label -but, especially when faced with distribution shift, the model is not guaranteed to obey the same invariances at test time. In this section, we introduce MEME, a method for test time robustness that adapts the model such that it respects these invariances on the test input. We use "test time robustness" specifically to refer to techniques that operate directly on pretrained models and single test inputs -single point BN adaptation and TTA, as described in Section 2, are examples of prior test time robustness methods. In the test time robustness setting, we are given a trained model f θ : X → Y with parameters θ ∈ Θ. We do not require any special training procedure and do not make any assumptions about the model, except that θ is adaptable and that f θ produces a conditional output distribution p θ (y|x) that is differentiable with respect to θ. 1 All standard deep neural network models satisfy these assumptions. A single point x ∈ X is presented to f θ , for which it must predict a label ŷ ∈ Y immediately. Note that this is precisely identical to the standard test time inference procedure for regular supervised learning models -in effect, we are simply modifying how inference is done, without any additional assumptions on the training process or on test time data availability. This makes test time robustness methods a simple "slot-in" replacement for the ubiquitous and standard test time inference process. We assume sampling access to a set of augmentation functions A {a 1 , . . . , a M } that can be applied to the test point x. We use these augmentations and the self-supervised objective detailed below to adapt the model before it predicts on x. When given a set of test inputs, the model adapts and predicts on each test point independently. We do not assume access to any ground truth labels.
 MARGINAL ENTROPY MINIMIZATION WITH ENSEMBLED AUGMENTATIONS Given a test point x and set of augmentation functions A, we sample B augmentations from A and apply them to x in order to produce a batch of augmented data x1 , . . . , xB . The model's average, or marginal, output distribution with respect to the augmented points is given by pθ (y|x) E U (A) [p θ (y|a(x))] ≈ 1 B B i=1 p θ (y|x i ) , (1) where the expectation is with respect to uniformly sampled augmentations a ∼ U(A). What properties do we desire from this marginal distribution? To answer this question, consider the role that data augmentation typically serves during training. For each training point (x train , y train ), the model f θ is trained using multiple augmented forms of the input xtrain 1 , . . . , xtrain E . f is trained to obey the invariances between the augmentations and the label -no matter the augmentation on x train , f should predict the same label y train , and it should do so confidently. We seek to devise a similar learning signal during test time, when no ground truth labels are available. That is, after adapting: (1) the model f θ predictions should be invariant across augmented versions of the test point, and (2) the model f θ should be confident in its predictions, even for heavily augmented versions of the test point, due to the additional knowledge that all versions have the same underlying label. Optimizing the model for more confident predictions can be justified from the assumption that the true underlying decision boundaries between classes lie in low density regions of the data space  (Grandvalet & Bengio, 2005) . With these two goals in mind, we propose to adapt the model using the entropy of its marginal output distribution over augmentations (Equation  1 ), i.e., (θ; x) H (p θ (•|x)) = − y∈Y pθ (y|x) log pθ (y|x) . (2) Optimizing this objective encourages both confidence and invariance to augmentations, since the entropy of pθ (•|x) is minimized when the model outputs the same (confident) prediction regardless of the augmentation. Given that θ is adaptable and p θ (y|x) is differentiable with respect to θ, we can directly use gradient based optimization to adapt θ according to this objective. We use only one gradient step per test point, because empirically we found this to be sufficient for improved performance while being more computationally efficient. After this step, we can use the adapted model, which we denote f θ , to predict on the original test input x. Algorithm 1 presents the overall method MEME for test time adaptation. Though prior test time adaptation methods must carefully choose which parameters to adapt in order to avoid degenerate solutions  (Wang et al., 2021) , our adaptation procedure simply adapts all of the model's parameters θ (line 3). Note that, as discussed above, the model f θ adapts using augmented data but makes its final prediction on the original point (line 4), which may be easier to predict on.
 COMPOSING MEME WITH PRIOR METHODS An additional benefit of MEME is that it synergizes with other approaches for handling distribution shift. In particular, MEME can be composed with prior methods for training robust models and adapting model statistics, thus leveraging the performance improvements of each technique. Pretrained robust models. Since MEME makes no assumptions about, or modifications to, the model training procedure, performing adaptation on top of pretrained robust models, such as those trained with heavy data augmentations, is as simple as using any other pretrained model. Crucially, we find that, in practice, the set of augmentations that we use at test time A does not have to match the augmentations that were used to train the model. This is important as we require a few properties from the test time augmentations: that they can be easily sampled and are applied directly to the model input x. These properties do not hold for, e.g., data augmentation techniques based on image translation models, such as DeepAugment  (Hendrycks et al., 2021a) , or feature mixing, such as moment exchange  (Li et al., 2021) . However, we can still use models trained with these data augmentation techniques as our starting point for adaptation, thus allowing us to improve upon their state-of-the-art results. As noted above, using pretrained models is not as easily accomplished for adaptation methods which require complicated or specialized training procedures and model architectures, such as TTT  (Sun et al., 2020)  or ARM  (Zhang et al., 2021) . In our experiments, we use AugMix as our set of augmentations  (Hendrycks et al., 2020) , as it satisfies the above properties and still yields significant diversity when applied, as depicted in Figure  2 . Adapting BN statistics.  Schneider et al. (2020)  showed that, even when presented with just a single test point, partially adapting the estimated mean and variance of the activations in each batch normalization (BN) layer of the model can still be effective in some cases for handling distribution shift. In this setting, to prevent overfitting to the test point, the channelwise mean and variance [µ test , σ 2 test ] estimated from this point are mixed with the the mean and variance [µ train , σ 2 train ] computed during training according to a prior strength N , i.e., µ N N + 1 µ train + 1 N + 1 µ test , σ 2 N N + 1 σ 2 train + 1 N + 1 σ 2 test . This technique is also straightforward to combine with MEME: we simply use the adapted BN statistics whenever computing the model's output distribution. That is, we adapt the BN statistics alongside all of the model parameters for MEME. Following the suggestion in  Schneider et al. (2020) , we set N = 16 for all of our experiments in the next section.
 EXPERIMENTS Our experiments aim to answer the following questions: (1) How does MEME compare to prior methods for test time adaptation, which make additional training and test assumptions, and test time robustness? (2) Can MEME be successfully combined with a wide range of pretrained models? (3) Which aspects of MEME are the most important for strong performance? We conduct experiments on two distribution shift benchmarks for CIFAR-10 (  Krizhevsky, 2009)  and three distribution shift benchmarks for ImageNet  (Russakovsky et al., 2015) . Specifically, for CIFAR-10, we evaluate on the CIFAR-10-C  (Hendrycks & Dietterich, 2019)  and CIFAR-10.1  (Recht et al., 2018)     2020 ), set the prior strength N = 256 accordingly. For Tent, we use test batch sizes of 64 and, for ResNet-50 models, test both "online" adaptation -where the model adapts continually through the entire evaluation -and "episodic" adaptationwhere the model is reset after each test batch.  (Wang et al., 2021) . Note that the evaluation protocols are different for these two methods: whereas MEME is tasked with predicting on each test point immediately after adaptation, BN adaptation predicts on a batch of 256 test points after computing BN statistics on the batch, and Tent predicts on a batch of 64 inputs after adaptation but also, in the online setting, continually adapts throughout evaluation. In all experiments, we further compare to single point BN adaptation  (Schneider et al., 2020)  and the TTA baseline that simply predicts according to the model's marginal output distribution over augmentations pθ (y|x) (Equation  1 )  (Krizhevsky et al., 2012; Ashukha et al., 2020) . Full details on our experimental protocol are provided in Appendix A. To answer question (2), we apply MEME on top of multiple pretrained models with different architectures, trained via several different procedures. For CIFAR-10, we train our own ResNet-26  (He et al., 2016)  models. For ImageNet, we use the best performing ResNet-50 robust models from prior work, which includes those trained with DeepAugment and AugMix augmentations  (Hendrycks et al., 2021a)  as well as those trained with moment exchange and CutMix  (Li et al., 2021) . To evaluate the generality of prior test time robustness methods and MEME, we also evaluate the small robust vision transformer (RVT * -small), which provides superior performance on all three ImageNet distribution shift benchmarks compared to the robust ResNet-50 models  (Mao et al., 2021) . Finally, to answer (3), we conduct ablative studies in subsection 4.2: first to determine the relative importance of maximizing confidence (via entropy minimization) versus enforcing invariant predictions across augmented copies of each test point, and second to determine the importance of the particular augmentation functions used. The comparison to the non adaptive TTA baseline also helps determine whether simply augmenting the test point is sufficient or if adaptation is additionally helpful.
 MAIN RESULTS We summarize results for CIFAR-10, CIFAR-10.1, and CIFAR-10-C in MEME also provides a larger performance gain on CIFAR-10.1 compared to TTT. We find that the non adaptive TTA baseline is competitive for these relatively simple test sets, though it is worse than MEME for CIFAR-10-C. Of these three test sets, CIFAR-10-C is the only benchmark that explicitly introduces distribution shift, which suggests that adaptation is useful when the test shifts are 7.3 (−1.9) 14.7 (−3.7) 19.6 (−2.9) + Joint training  (Sun et al., 2020)  8.1 16.7 22.8 + TTT  (Sun et al., 2020)  7.9 (−0.2) 15.9 (−0.8) 21.5 (−1.3) more prominent. Both TTA and MEME are also effective at improving performance for the original CIFAR-10 test set where there is no distribution shift, providing further support for the widespread use of augmentations in standard evaluation protocols  (Krizhevsky et al., 2012; Ashukha et al., 2020) . We summarize results for ImageNet-C, ImageNet-R, and ImageNet-A in   (Schneider et al., 2020; Wang et al., 2021) , accessing multiple test points can be powerful for benchmarks such as ImageNet-C and ImageNet-R, in which inferred statistics from the test input distribution may aid in prediction. In the case of Tent, which adapts online, the model has adapted using the entire test set by the end of evaluation. However, these methods do not help, and oftentimes even hurt, for ImageNet-A. Furthermore, we find that these methods are less effective with the RVT * -small model, which may indicate their sensitivity to model architecture choices. Therefore, for this model, we also test a modification of Tent which adapts all parameters, and we find that this version of Tent works better for ImageNet-C but is significantly worse for ImageNet-R. MEME results in substantial improvement for ImageNet-A and is competitive with TTA on this problem. No prior test time adaptation methods have reported improvements on ImageNet-A, and some have reported explicit negative results  (Schneider et al., 2020) . As discussed, it is reasonable for adaptation methods that rely on multiple test points to achieve greater success on other benchmarks such as ImageNet-C, in which a batch of inputs provides significant information about the specific corruption that must be dealt with. In contrast, ImageNet-A does not have such obvious characteristics associated with the input distribution, as it is simply a collection of images that are difficult to classify. As MEME instead extracts a learning signal from single test points, it is, to the best of our knowledge, the first test time adaptation method to report successful results on this testbed. When used on top of a model trained with moment exchange and CutMix  (Li et al., 2021) , MEME achieves stateof-the-art performance among ResNet-50 models and single test point methods. TTA generally offers larger performance gains on ImageNet-A and also results in the highest overall accuracy when combined with the RVT * -small model; however, TTA performs worse than MEME on ImageNet-R and consistently decreases accuracy on ImageNet-C compared to standard evaluation. We view the consistency with which MEME outperforms the best prior methods, which change across different test sets, to be a major advantage of the proposed method.
 ABLATIVE STUDY MEME increases model robustness at test time via adaptation and augmentation. In this section, we ablate the adaptation procedure, and in Appendix B we ablate the use and choice of augmentations. From the results above, we conclude that adaptation generally provides additional benefits beyond simply using TTA to predict via the marginal output distribution pθ (y|x). However, we can disentangle two distinct self-supervised learning signals that may be effective for adaptation: encouraging invariant predictions across different augmentations of the test point, and encouraging confidence via entropy minimization. The marginal entropy objective in Equation 2 encapsulates both of these learning signals, but it cannot easily be decomposed into these pieces. Thus, we instead use two ablative adaptation methods that each only make use of one of these learning signals. First, we consider optimizing the pairwise cross entropy between each pair of augmented points, i.e., PCE (θ; x) 1 B × (B − 1) B i=1 j =i y∈Y p θ (y|x i ) log p θ (y|x j ) , Where xi again refers to the i-th sampled augmentation applied to x. Intuitively, this loss function encourages the model to adapt such that it produces the same predictive distribution for all augmentations of the test point, but it does not encourage the model to produce confident predictions. Conversely, as an objective that encourages confidence but not invariance, we also consider optimizing conditional entropy on the batch of augmented points, i.e., CE (θ; x) 1 B B i=1 H(p θ (•|x i )) . This ablation is effectively a version of the episodic variant of Tent  (Wang et al., 2021)  that produces augmented copies of a single test point rather than assuming access to a test batch. We first evaluate these ablations on the CIFAR-10 test sets. We use the same adaptation procedure outlined in Algorithm 1, with replaced with the above objectives, and we keep the same hyperparameter values. The results are presented in Table  3 . We see that MEME, i.e., marginal entropy minimization, generally performs better than adaptation with either of the alternative objectives. This supports the hypothesis that both invariance across augmentations and confidence are important learning signals for self-supervised adaptation. When faced with CIFAR-10.1, we see poor performance from the pairwise cross entropy based adaptation method. On the original CIFAR-10 test set and CIFAR-10-C, the ablations perform nearly identically and uniformly worse than MEME. To further test the CE ablation, we also evaluate it on the ImageNet test sets for the RVT * -small model. We find that, similarly, minimizing conditional entropy generally improves performance compared to the baseline evaluation. MEME is more performant for ImageNet-C and ImageNet-R, again indicating the benefits of encouraging invariance to augmentations. Adaptation via CE performs slightly better for ImageNet-A, though for this problem, TTA is still the best method.
 DISCUSSION We presented MEME, a method for test time robustification again distribution shift via adaptation and augmentation. MEME does not require access or changes to the model training procedure and is thus broadly applicable for a wide range of pretrained models. Furthermore, MEME adapts at test time using single test inputs, thus it does not assume access to multiple test points as in several recent methods for test time adaptation  (Schneider et al., 2020; Wang et al., 2021; Zhang et al., 2021) . On a range of distribution shift benchmarks for CIFAR-10 and ImageNet classification, and for both ResNet and vision transformer models, MEME consistently improves performance at test time and achieves several new state-of-the-art results for these models in the single test point setting. Inference via MEME is more computationally expensive than standard model inference, primarily because adaptation is performed per test point and thus inference cannot be batched. When deployed in the real world, it is natural to expect that test points will arrive one at a time and batched inference will not be possible. However, MEME is also more computationally expensive due to its augmentation and adaptation procedure. One interesting direction for future work is to develop techniques for selectively determining when to adapt the model in order to achieve more efficient inference. For example, with well calibrated models  (Guo et al., 2017) , we may run simple "feedforward" inference when the prediction confidence is over a certain threshold, thus achieving better efficiency. Additionally, it would be interesting to explore MEME in the test setting where the model is allowed to continually adapt as more test data is observed. In our preliminary experiments in this setting, MEME tended to lead to degenerate solutions, e.g., the model predicting a constant label with maximal confidence, and this may potentially be rectified by carefully choosing which parameters to adapt  (Wang et al., 2021)  or regularizing the model such that it does not change too drastically from the pretrained model. Distribution shift, in general, lies at the heart of many ethical concerns in machine learning. When machine learning models and research do not adhere to ethical principles such as "contribute to society and to human well-being", "avoid harm", and "be fair and take action to avoid discrimination", it can oftentimes be attributed at least partially to issues stemming from distribution shift. As such, research into ways to combat and mitigate shift contribute to furthering the ethical goals and principles laid out by this conference and the broader community. Related to this work, we may imagine that models that are more robust or adaptable to shift may produce more fair results when considering underrepresented subpopulations and could be more trustworthy in safety critical applications. However, as mentioned in Section 5, the proposed method is also more computationally intensive, and care must be taken in general to not place the most powerful machine learning tools exclusively in the hands of the most privileged and resource rich individuals and organizations. Beyond this, we do not see any immediate ethical concerns regarding the content, presentation, or methodology of this work.
 REPRODUCIBILITY STATEMENT The proposed method is relatively simple and, to the best of our ability, explained fully in Section 3 and Algorithm 1. This explanation is bolstered by a complete description of all hyperparameters in Appendix A as well as the example code provided in the supplementary materials. As briefly explained in the README in the example code, all datasets were downloaded from publicly accessible links and preprocessed only following the standard protocols. Upon publication, we will include a link to the full (non anonymous) code release containing full instructions for setting up the code, downloading the datasets, and reproducing the results.
 A EXPERIMENTAL PROTOCOL We select hyperparameters using the four disjoint validation corruptions provided with CIFAR-10-C and ImageNet-C  (Hendrycks & Dietterich, 2019) . As the other benchmarks are only test sets and do not provide validation sets, we use the same hyperparameters found using the corruption validation sets and do not perform any additional tuning. For the ResNet models that we evaluate, we use stochastic gradients as the update rule G; for ResNet-26 models, we set the number of augmentations B = 32 and the learning rate η = 0.005; and for ResNet-50 models, we set B = 64 and η = 0.00025. For the robust vision transformer, we use AdamW  (Loshchilov & Hutter, 2019)  as the update rule G, with learning rate η = 0.00001 and weight decay 0.01, and B = 64. In the CIFAR evaluation, we compare to TTT, which, as noted, can also be applied to single test inputs but requires a specialized training procedure  (Sun et al., 2020) . Thus, the ResNet-26 model we use for our method closely follows the modifications that  Sun et al. (2020)  propose, in order to provide a fair point of comparison. In particular,  Sun et al. (2020)  elect to use group normalization  (Wu & He, 2018)  rather than BN, thus single point BN adaptation is not applicable for this model architecture. As noted before, TTT also requires the joint training of a separate rotation prediction head, thus further changing the model architecture, while MEME directly adapts the standard pretrained model. The TTA results are obtained using the same AugMix augmentations as for MEME. The single point BN adaptation results use N = 16, as suggested by  Schneider et al. (2020) . As noted, the BN adaptation results (using multiple test points) are obtained using N = 256 as the prior strength and batches of 256 test inputs for adaptation. For Tent, we use the hyperparameters suggested in  Wang et al. (2021) : we use stochastic gradients with learning rate 0.00025 and momentum 0.9, the adaptation is performed with test batches of 64 inputs, and the method is run online, i.e., prediction and adaptation occur simultaneously and the model is allowed to continuously adapt through the entire test epoch. Since  Wang et al. (2021)  did not experiment with transformer models, we also attempted to run Tent with Adam  (Kingma & Ba, 2015)  and AdamW  (Loshchilov & Hutter, 2019)  and various hyperparameters for the RVT * -small model; however, we found that this generally resulted in worse performance than using stochastic gradient updates with the aforementioned hyperparameters. We obtain the baseline ResNet-50 parameters directly from the torchvision library. The parameters for the ResNet-50 trained with DeepAugment and AugMix are obtained from https://drive.google.com/file/d/1QKmc_p6-qDkh51WvsaS9HKFv8bX5jLnP. The parameters for the ResNet-50 trained with moment exchange and CutMix are obtained from https://drive.google.com/file/d/1cCvhQKV93pY-jj8f5jITywkB9EabiQDA. The parameters for the small robust vision transformer (RVT * -small) model are obtained from https://drive.google.com/file/d/1g40huqDVthjS2H5sQV3ppcfcWEzn9ekv.
 B ADDITIONAL EXPERIMENTS In this section, we analyze the importance of using augmentations during adaptation, study the tradeoffs between efficiency and accuracy for MEME, and present results with ResNext101 models  (Xie et al., 2017; Mahajan et al., 2018; Orhan, 2019)  on ImageNet-A.
 B.1 ANALYSIS ON AUGMENTATIONS One may first wonder: are augmentations needed in the first place? In the test time robustness setting when only one test point is available, how would simple entropy minimization fare? We answer this question in Table  4  by evaluating the episodic variant of Tent (i.e., with model resetting after each batch) with a test batch size of 1. This approach is also analogous to a variant of MEME that does not use augmentations, since for one test point and no augmented copies, conditional and marginal entropy are the same. Similar to MEME, we also incorporate single point BN adaptation with N = 16, in place of the standard BN adaptation that Tent typically employs using batches of test inputs. The results in Table  4  indicate that entropy minimization on a single test point generally provides no additional performance gains beyond just single point BN adaptation. This empirically shows that using augmentations is important for achieving the reported results. We also wish to understand the importance of the choice of augmentation functions A. As mentioned, we used AugMix  (Hendrycks et al., 2020)  in the previous experiments as it best fit our criteria: 69.9 (−6.8) 58.8 (−5.1) 99.1 (−0.9) + Tent (episodic, batch size 1)  (Wang et al., 2021)   69.1 (−5.7) 59.4 (−3.3) 89.0 (−2.9) + Tent (episodic, batch size 1)  (Wang et al., 2021)  70.9 (−3.9) 62.6 (−1.9) 91.1 (−0.8) RVT * -small  (Mao et al., 2021)  49  AugMix requires only the input x, and randomly sampled augmentations lead to diverse augmented data points. A simple alternative is to instead use the "standard" set of augmentations commonly used in ImageNet training, i.e., random resized cropping and random horizontal flipping. We evaluate this ablation of using MEME with standard augmentations also on the CIFAR-10 test sets, again with the same hyperparameter values. From the results in Table  5 , we can see that MEME is still effective with simpler augmentation functions. This is true particularly for the cases where there is no test shift, as in the original CIFAR-10 test set, or subtle shifts as in CIFAR-10.1; however, for the more severe and systematic CIFAR-10-C shifts, using heavier AugMix data augmentations leads to greater performance gains over the standard augmentations. Furthermore, this ablation was conducted using the ResNet-26 model, which was trained with standard augmentations -for robust models such as those in Table  2 , AugMix may offer greater advantages at test time since these models were exposed to heavy augmentations during training.
 B.2 ANALYZING THE TRADEOFF BETWEEN EFFICIENCY AND ACCURACY In Figure  3 , we analyze the % test error of MEME adaptation on ImageNet-R as a function of the efficiency of adaptation, measured in seconds per evaluation. We achieve various tradeoffs by varying the number of augmented copies B = {1, 2, 4, 8, 16, 32, 64, 128}. We note that small values of B such as 4 and 8 can already provide significant performance gains, indicating that a practical tradeoff between efficiency and accuracy is possible. For large B, the wall clock time is dominated by computing the augmentations -in our implementation, we do not compute augmentations in
 B.3 EVALUATING RESNEXT101 MODELS ON IMAGENET-A ResNext-101 models  (Xie et al., 2017)  have been found to achieve higher accuracies on ImageNet-A  (Hendrycks et al., 2021b) , particularly when trained with massive scale weakly supervised pretraining  (Mahajan et al., 2018; Hendrycks et al., 2021a) . In this section, we evaluate whether MEME can successfully adapt these models and further improve performance on this challenging test set. We use the same hyperparameters as for the robust vision transformer with no additional tuning: AdamW  (Loshchilov & Hutter, 2019)  as the update rule G, learning rate η = 0.00001, weight decay 0.01, and B = 64. We obtain the baseline ResNext-101 (32x8d) parameters, pretrained on ImageNet, directly from the torchvision library. We also evaluate a ResNext-101 (32x8d) pretrained with weakly supervised learning (WSL) on billions of Instagram images  (Mahajan et al., 2018) , and we obtained the parameters from https://download.pytorch.org/models/ig_resnext101_32x8-c38310e5.pth. For the WSL model, we did not use single point BN adaptation as we found this technique to be actually harmful to performance, and this corroborates previous findings  (Schneider et al., 2020) . Table  6  summarizes the results. We can see that, similar to Table  2 , Both TTA and MEME significantly improve upon the baseline model evaluation. TTA performs best for the baseline ResNext-101 model. However, MEME ultimately achieves the best accuracy by a significant margin, as it is more successful at adapting the WSL model, which has a much higher accuracy. This suggests that combining MEME with other large pretrained models may be an interesting direction for future work. C FULL CIFAR-10-C AND IMAGENET-C RESULTS In the following tables, we present test results broken down by corruption and level for CIFAR-10-C for the methods evaluated in Table  1 . We omit joint training and TTT because these results are available from  Sun et al. (2020) . Our test results for ImageNet-C are provided in the CSV files in the supplementary material.   25.3 16.9 15.9 8.4 33.5 14.6 13.0 17.7 14.4 8.5 7.7 9.6 10.7 11.9 16.2  Table  11 : Test error (%) on CIFAR-10-C level 1 corruptions. gauss shot impul defoc glass motn zoom snow frost fog brit contr elast pixel jpeg ResNet-26 20.8 16.5 15.8 9.2 38.9 11.8 12.8 13.9 13.4 9.7 9.4 9.6 13.1 12.0 16.4 + TTA 15.8 12.8 11.8 7.3 35.1 10.8 12.5 11.0 10.8 7.4 7.4 7.7 11.4 9.2 12.5 + MEME (ours) 16.1 12.9 11.9 7.4 34.7 10.4 12.1 11.0 10.7 7.4 7.3 7.5 10.9 9.2 12.5 Figure 2 : 2 Figure 2: We visualize augmentations of a randomly chosen data point from the "Gaussian Noise level 3" ImageNet-C test set. Even for a robust model trained with heavy data augmentations (Hendrycks et al., 2021a), both its predictive accuracy and confidence drop sharply when encountering test distribution shift. As shown in the bottom two rows, these drops can be remedied via MEME.
 Figure 3 : 3 Figure 3: Plotting MEME efficiency as seconds per evaluation (x axis) and % test error on ImageNet-R (y axis) for the ResNet-50 models (left) and RVT * -small (right) while varying B = {1, 2, 4, 8, 16, 32, 64, 128}. Note the log scale on the x axis.

 Algorithm 1 Test time robustness via MEME Require: trained model f θ , test point x, # augmentations B, learning rate η, update rule G 1: Sample a 1 , . . . , a B i.i.d. ∼ U(A) and produce augmented points xi = a i (x) for i ∈ {1, . . . , B} 2: Compute Monte Carlo estimate p = 1 B 3: Adapt model parameters via update rule θ ← G(θ, η, ˜ ) B i=1 p θ (y|x i ) ≈ pθ (y|x) and ˜ = H(p) ≈ (θ; x) 4: Predict ŷ arg max y p θ (y|x)
 test sets, and for ImageNet, we evaluate on the ImageNet-C (Hendrycks & Dietterich, 2019) , ImageNet-R (Hendrycks et al., 2021a) , and ImageNet-A (Hendrycks et al., 2021b)  test sets.To answer question (1), we compare to test time training (TTT)(Sun et al., 2020)  in the CIFAR-10 experiments, for which we train ResNet-26 models following their protocol and specialized architecture. We do not compare to TTT for the ImageNet experiments due to the computational demands of training state-of-the-art models and becauseSun et al. (2020)  do not report competitive ImageNet results. For the ImageNet experiments, we compare to Tent(Wang et al., 2021)  and BN adaptation, which can be used with pretrained models but require multiple test inputs (or even the entire test set) for adaptation. We provide BN adaptation with 256 test inputs at a time and, followingSchneider et al. (
 Table 1, with full CIFAR-10-C results in Appendix C. We use indentations to indicate composition, e.g., TTT is performed at test time on top of their specialized joint training procedure. Across all corruption types in CIFAR-10-C, MEME consistently improves test error compared to the baselines, non adaptive TTA, and TTT.
 Table 1 : 1 Results for the original CIFAR-10 test set, CIFAR-10.1, and CIFAR-10-C. MEME outperforms TTT despite not making any training assumptions. Results from Sun et al. (2020). CIFAR-10 CIFAR-10.1 CIFAR-10-C Error (%) Error (%) Average Error (%) ResNet-26 (He et al., 2016) 9.2 18.4 22.5 + TTA 7.3 (−1.9) 14.8 (−3.6) 19.9 (−2.6) + MEME (ours)
 Table 2 : 2 Test results for ImageNet-C, ImageNet-R, and ImageNet-A. MEME achieves new state-ofthe-art performance on each benchmark for ResNet-50 models for the single test point setting. For RVT * -small, MEME substantially improves performance across all benchmarks and reaches a new state of the art for ImageNet-C and ImageNet-R. ImageNet-C ImageNet-R ImageNet-A mCE ↓ Error (%) Error (%)
 Table 2 2 , with complete ImageNet-C results in Appendix C. We again use indentations to indicate composition, e.g., the best results on ImageNet-C for our setting are attained through a combination of starting from a model trained with DeepAugment and AugMix (Hendrycks et al., 2021a)  and using MEME on
 Table 3 : 3 Ablating the adaptation objective to test pairwise cross entropy and conditional entropy (CE) based adaptation. MEME generally performs the best, indicating that both encouraging invariance across augmentations and confidence are helpful in adapting the model. top. For both ImageNet-C and ImageNet-R, and for both the ResNet-50 and RVT * -small models, combining MEME with robust training techniques leads to new state of the art performance among methods that observe only one test point at a time. We highlight in gray the methods that require multiple test points for adaptation, and we list in bold the best results from these methods which outperform the test time robustness methods. As Table2and prior work both show CIFAR-10 CIFAR-10.1 CIFAR-10-C Error (%) Error (%) Average Error (%) ResNet-26 (He et al., 2016) 9.2 18.4 22.5 + MEME (ours) 7.3 (−1.9) 14.7 (−3.7) 19.6 (−2.9) − (Equation 2) + PCE 7.6 (−1.6) 15.3 (−3.1) 20.0 (−2.5) − (Equation 2) + CE 7.6 (−1.6) 14.7 (−3.7) 20.0 (−2.5) ImageNet-C ImageNet-R ImageNet-A mCE ↓ Error (%) Error (%) RVT  *  -small (Mao et al., 2021) 49.4 52.3 73.9 + MEME (ours) 40.6 (−8.8) 43.8 (−8.5) 69.8 (−4.1) − (Equation 2) + CE 41.2 (−8.2) 44.2 (−8.1) 69.7 (−4.2)
 Table 4 : 4 Evaluating the episodic version of Tent with a batch size of 1, which corresponds to a simple entropy minimization approach for the test time robustness setting. This approach also uses single point BN adaptation, and entropy minimization does not provide much, if any, additional gains. ImageNet-C ImageNet-R ImageNet-A mCE ↓ Error (%) Error (%)
 Table 5 : 5 Ablating the augmentation functions to test standard augmentations (random resized cropping and horizontal flips). When changing the augmentations used, the post-adaptation performance generally does not change much, though it suffers the most on CIFAR-10-C. CIFAR-10 CIFAR-10.1 CIFAR-10-C Error (%) Error (%) Average Error (%)
 Table 6 : 6 ImageNet-A results for the ResNext-101 models. parallel, though in principle this is possible for AugMix and should improve efficiency overall. These experiments used four Intel Xeon Skylake 6130 CPUs and one NVIDIA TITAN RTX GPU. ImageNet-A Error (%)
 Table 7 : 7 Test error (%) on CIFAR-10-C level 5 corruptions. gauss shot impul defoc glass motn zoom snow frost fog brit contr elast pixel jpeg ResNet-26 48.4 44.8 50.3 24.1 47.7 24.5 24.1 24.1 33.1 28.0 14.1 29.7 25.6 43.7 28.3 + TTA 43.4 39.6 42.9 28.3 44.7 26.3 26.3 21.4 28.5 23.3 12.1 32.9 21.7 43.2 21.7 + MEME (ours) 43.5 39.8 43.3 26.4 44.4 25.1 25.0 20.9 28.3 22.8 11.9 28.3 21.1 42.8 21.7
 Table 8 : 8 Test error (%) on CIFAR-10-C level 4 corruptions. gauss shot impul defoc glass motn zoom snow frost fog brit contr elast pixel jpeg ResNet-26 43.8 37.2 39.3 14.8 48.0 19.9 18.7 22.0 24.9 15.1 11.4 16.8 19.1 27.9 24.9 + TTA 39.5 32.0 31.8 15.4 45.0 20.9 20.2 18.9 21.7 12.9 9.3 16.8 17.7 25.7 18.9 + MEME (ours) 39.7 32.3 32.2 14.7 45.0 20.0 19.2 18.7 21.1 12.5 9.3 15.2 16.9 25.2 18.9
 Table 9 : 9 Test error (%) on CIFAR-10-C level 3 corruptions. gauss shot impul defoc glass motn zoom snow frost fog brit contr elast pixel jpeg ResNet-26 40.0 33.8 26.4 11.5 37.3 20.0 16.6 20.0 24.7 12.2 10.5 13.6 15.0 18.4 22.7 + TTA 34.3 27.7 20.3 11.3 32.9 20.7 16.7 16.3 21.1 9.8 8.5 12.5 13.7 14.5 17.2 + MEME (ours) 34.4 27.9 20.5 10.8 32.8 19.8 16.1 16.1 20.9 9.6 8.6 11.7 13.2 14.5 17.2
 Table 10 : 10 Test error (%) on CIFAR-10-C level 2 corruptions. gauss shot impul defoc glass motn zoom snow frost fog brit contr elast pixel jpeg ResNet-26 30.1 21.8 21.1 9.7 38.3 15.3 13.8 21.2 17.6 10.5 9.7 11.6 12.9 15.4 21.3 + TTA 25.3 16.9 15.8 8.5 33.8 15.3 13.7 17.8 14.5 8.7 7.8 10.0 11.0 12.0 16.1 + MEME (ours)
			 Single point BN adaptation also assumes that the model has batch normalization layers, and, as shown empirically in Section 4, this is an assumption that we do not require but can also benefit from.


Abstract: This talk proposes a very simple "baseline architecture" for a learning agent that can handle stochastic, partially observable environments. The architecture uses reinforcement learning together with a method for representing temporal processes as graphical models. I will discuss methods for leaming the parameters and structure of such representations from sensory inputs, and for computing posterior probabilities. Some open problems remain before we can try out the complete agent; more arise when we consider scaling up. A second theme of the talk will be whether reinforcement learning can provide a good model of animal and human learning. To answer this question, we must do inverse reinforcement learning: given the observed behaviour, what reward signal, if any, is being optimized? This seems to be a very interesting problem for the COLT, UAI, and ML communities, and has been addressed in econometrics under the heading of structural estimation of Markov decision processes. 1 Learning in uncertain environments AI is about the construction of intelligent agents, i.e., systems that perceive and act effectively (according to some performance measure) in an environment. I have argued elsewhere  Russell and Norvig (1995)  that most AI research has focused on environments that are static, deterministic, discrete, and fully observable. What is to be done when, as in the real world, the environment is dynamic, stochastic, continuous, and partially observable? 'This paper draws on a variety of research efforts supported

Learning agents for uncertain environments (extended abstract)

In recent years, reinforcement learning (also called neurodynamic programming) has made rapid progress as an approachfor building agents automatically  (Sutton, 1988; Kaelbling et al., 1996; Bertsekas & Tsitsiklis, 1996) . The basic idea is that the performance measure is made available to the agent in the form of a rewardfunction specifying the reward for each state that the agent passes through. The performance measure is then the sum of the rewards obtained. For example, when a bumble bee forages, the reward function at each time step might be some combination of the distance flown (weighted negatively) and the nectar ingested. Reinforcement learning (RL) methods are essentially online algorithmd for solving Markovdecisionprocesses (MDPs). An MDP is defined by the reward function and a model, that is, the state transition probabilities conditioned on each possible action. RL algorithms can be model-based, where the agent learns a model, or model-free-e.g., Q-learning cite-Watkins: 1989, which learns just a function Q(s, a) specifying the long-term value of taking action a in state s and acting optimally thereafter. Despite their successes, RL methods have been restricted largely tofully observable MDPs, in which the sensory input at each state is sufficient to identify the state. Obviously, in the real world, we must often deal with partially observable MDPs (POMDPs).  Astrom (1965)  proved that optimal decisions in POMDPs depend on the belief state b at each point in time, i.e., the posterior probability distribution over all possible actual states, given all evidence to date. The functions V and Q then become functions of b instead of s.  Parr and Russell (1995)  describes a very simple POMDP RL algorithm using an explicit representation of b as a vector of probabilities, and  McCallum (1993)  shows a way to approximate the belief state using recent percept sequences. Neither approach is likely to scale up to situations with large numbers of state variables and long-term temporal dependencies. What is needed is a way of representing the model compactly and updating the belief state efficiently given the model and each new observation. Dynamic Bayesian networks  (Dean & Kanazawa, 1989)  seem to have some of the required properties; in particular, they have significant advantages over other approaches such as Kalman filters and hidden Markov models. Our baseline architecture, shown in Figure  1 , uses DBNs to represent and update the belief state as new sensor information arrives. Given a representation for b, the reward signal is used to learn a Q-function represented by some "black-box" function approximator such as a neural network. Provided we can handle hybrid (dis-Figure  1 : A baseline architecture for learning agents in uncertain environments crete+continuous) DPNs, and provided we have a learning algorithm that can construct an approximately correct DBN model from scratch, then this baseline architecture has the capacity, in principle, to be thrown into more or less any environment and to learn to behave reasonably.* The talk will cover a variety of research topics arising from this proposal: Parametric learning in DBNs  (Binder, Koller, Russell, & Kanazawa, 1997a) . Structural learning in DBNs  (Friedman, Murphy, & Russell, 1998) . Approximate inference in DBNs  (Kanazawa, Koller, & Russell, 1995; Boyen & Koller, 1998) . Space-efficient inference in DBNs  (Binder, Murphy, & Russell, 1997b) . Reinforcement learning with DBN models-that is, how to do Q-learning with the belief state information provided by the DBN. Some tentative ideas will be presented but as yet there are no convincing solutions. Scaling up the environment will inevitably overtax the resources of the baseline architecture. There are several obvious directions for improvement, including hierarchical and first-order models, hierarchical representations of behaviour  (Parr & Russell, 1998) , and model-based lookahead methods for decision making. Which of these is important in any particular class of environments can only be ascertained by experiment.
 Inverse reinforcement learning Reinforcement learning is a powerful method for adaptive control in real tasks, so it is natural to seek analogous mechanisms in nature. Connections have been made between reinforcement learning and operant conditioning models of animal learning (see, e.g.,  Schmajuk & Zanutto, 1997; Touretzky & Saksida, 1997) . There is also neurophysiological evidence that reinforcement learning occurs in bee foraging  (Montague et al., 1995)  and in songbird vocalization  (Doya & Sejnowski, 1995) . In this work, it is generally assumed that the reward function is fixed and known. For example, in experiments on bees it is assumed to be the rate of nectar ingestion: Montague et al. (1995) cite evidence of a "neuron with widespread projections to odour processing regions of the honeybee brain 'We say "more or less" because full generality require dealing with game-theoretic issues requiring stochastic decision making. whose activity represents the reward value of gustatory stimuli." It seems clear, however, that in examining animal and human behaviour we must consider the reward function as an unknown to be ascertained. The reasons for this are straightforward: l The specification of a given reward function is an empirical hypothesis and may turn out to be wrong. For example, it was assumed initially that horses' gait selection for a given speed was determined by energetic economy  (Hoyt & Taylor, 1981) ; this turns out not to be the case  (Farley & Taylor, 1991) . l The parameters of a multiattribute reward function can surely not be determined a priori; e.g., for running, attributes might be speed, efficiency, stability against perturbations, wear and tear on muscles, tendons, and bones, etc. How are these to be weighted and combined? Therefore, to model natural learning using reinforcement learning ideas, we must first solve the following computational task, which we call inverse reinforcement learning: Given 1) measurements of an agent's behaviour over time, in a variety of circumstances, 2) measurements of the sensory inputs to that agent; 3) a model of the physical environment (including the agent's body). Determine the reward function that the agent is optimizing. Given an assumption of optimization, this computational task is well-defined. Notice that is the dual of unsupervised reinforcement learning, where the task is to determine optimal behaviour given the reward inputs. To our knowledge, this computational task has not been studied in any generality in computer science, control theory, psychology, or biology. The closest work is in economics, where the task of multiattribute utility assessment has been studiedin depth-that is, how does a person actually combine the various attributes of each available choice when making a decision. The theory is well-developed  (Keeney & Raiffa, 1976) , and the applications numerous. However, this field studies only one-shot decisions where a single action is taken and the outcome is immediate. The sequential case was not considered until a seminal paper by  Sargent (1978)  tried to ascertain the effective hiring cost for labor by examining a firm's hiring behaviour over time, assuming it to be rational. In the last decade, the area of structural estimation of Markov decision processes has grown rapidly in econometrics  (Rust, 1994) . Many of the basic results carry over to our setting, although virtually nothing has been done on computational aspects, experimentation, or control-type applications. The open research problems are many: What are efficient algorithms for solving the inverse reinforcement learning problem? What is its computational complexity? Are there closed-form solutions for some parametric forms? Under what circumstances can we determine the existence of a consistent reward function? To what extent is the reward function uniquely recoverable? What effect do sensor and process noise have on robustness of the determination? What are appropriate error metrics for fitting? If behaviour is strongly inconsistent with optimality, can we identify "locally consistent" reward functions for specific regions in state space? Can we determine the reward function by observation during rather than after learning? How much observation is required to determine an estimated reward function that is within E of the true reward function? How can experiments be designed to maximize the identifiability of the reward function? Considering the design of possible algorithms, one can take maximum-likelihood approach to fit a parametric form for the reward functionas is commonly done in econometrics. That is, one defines a function L,(w)(B), the likelihood of observing behaviour B if the true reward function is r(w). From this, one can compute dWdw. One important question will be how to compute this gradient efficiently; presumably, it can be done in an obvious way by carrying the differential operator through the optimization algorithm for the behaviour. More elegant closed-form solutions may exist in special cases (e.g., linear-quadratic regulators). One may also be able to show that in some cases (e.g., linear reward functions) a globally optimal estimate can always be found. The solution of inverse reinforcement learning problems may also be an effective way to learn from observing experts. For tasks such as walking, diving, and driving, the designer of an artificial system may have only an intuitive idea of the appropriate reward function to be supplied to an RL algorithm in order to achieve "desirable" behavior. Instead of learning direct control functions from observation of experts (as in Pomerleau's ALVINN driging system), it may be better to solve the inverse reinforcement learning problem. The reward function should usually be a simple monotonic function of the current sensory inputs, and thus may be much simpler than the direct decision mapping itself. That is, the most compact and hence robustly learnable representation of expert behavior may be the reward function.


Abstract: Cooperation in settings where agents have different but overlapping preferences (mixed-motive settings) has recently received considerable attention in multi-agent learning. However, the mixed-motive environments typically studied are simplistic in that they have a single cooperative outcome on which all agents can agree. Multi-agent systems in general may exhibit many payoff profiles which might be called cooperative, but which agents have different preferences over. This causes problems for independently trained agents that do not arise in the case of that there is a unique cooperative payoff profile. In this note, we illustrate this problem with a class of games called mixed-motive coordination problems (MCPs). We demonstrate the failure of several methods for achieving cooperation in sequential social dilemmas when used to independently train policies in a simple MCP. We discuss some possible directions for ameliorating MCPs.

Julian Stastny
Johannes Treutlein
Maxime Riché
Jesse Clifton

Multi-agent learning in mixed-motive coordination problems

Introduction In An Essay on  Bargaining [1956] , Thomas Schelling describes a negotiation over the price of a house. Considering different strategies the negotiators might use, he points out that If the buyer is an agent authorized by a board of directors to buy at $16000 but not a cent more, and the directors cannot constitutionally meet again for several months and the buyer cannot exceed his authority, and if all this can be made known to the seller, then the buyer 'wins' -if, again, the seller has not tied himself up with a commitment to  $19000. (pg. 284)  Imagine various actors with different preferences -a mixed-motive setting  [Schelling, 1958]  -developing machine learning systems to act on their behalf. These actors are like Schelling's bargainers, in that their systems implicitly encode a commitment to pursuing a particular tradeoff between their own preferences and those of other players. Just as the price negotiation in Schelling's example may fail if buyer and seller are each committed to insisting on incompatible prices, machine learning systems may fail to cooperate if they encode commitments to incompatible tradeoffs between their principals' goals. Multi-agent learning (MAL) has in recent years paid considerable attention to problems of cooperation between agents with different, but not zero-sum, preferences. However, much of this work has been in the setting of (sequential) social dilemmas (SSDs) (e.g.,  Peysakhovich and Lerer 2017a , Lerer and Peysakhovich 2017 , Peysakhovich and Lerer 2017b , Foerster et al. 2018 , Wang et al. 2018 , Eccles et al. 2019 . The classic example of a social dilemma is the Prisoner's Dilemma (PD), and the SSDs studied in this literature are similar to the Prisoner's Dilemma in that there is a single salient notion of "cooperation". This means that it is relatively easy for actors to coordinate in their selection of policies to deploy in these settings. This situation is quite unlike that faced by Schelling's bargainers, for whom there are many possible agreements to commit to, and who thus may fail to reach an agreement if they make incompatible commitments 1 .  Cao et al. [2018]  are an exception in that they look at negotiation between deep reinforcement learners, such that there are multiple incompatible ways in which agents could choose among possible agreements. However, they do not evaluate the performance of independently trained policies when played against one Working draft. Last updated March 8, 2021. 1 The inadequacy of the Prisoner's Dilemma as a model for cooperation problems, and the need to study games that exhibit competing "cooperative" solutions, has been pointed out in other fields. See  Snidal [1985]  for a discussion in international relations,  McAdams [2008]  in law, and  Binmore [1998]  in evolutionary game theory. another ("cross-play"  [Hu et al., 2020] ), which is where the existence of multiple incompatible cooperative solutions creates problems. In this note, we argue that a focus on sequential social dilemmas and the lack of cross-play evaluation obscures a fundamental problem for decentralized MAL in mixed-motive settings. We discuss a class of games that we call mixed-motive coordination problems (MCPs), which highlight this problem. We compare MCPs to pure coordination problems and SSDs. We then illustrate how MAL methods which achieve cooperation in other social dilemmas fail in some simple MCPs, and argue that future methods for promoting cooperation in mixed-motive settings should be benchmarked in these more difficult cases. We close with some tentative positive contributions, by sketching some partial solutions to MCPs.
 Preliminaries We work in the setting of partially observable stochastic games (POSGs). For simplicity we assume two players, i = 1, 2. We will index player i's counterpart by −i. Each player has an action space A i , there is an underlying state space S which evolves according to a Markovian transition function P (S t+1 | S t , A t 1 , A t 2 ). At each time step, each player sees an observation O t i which depends on S t . Thus each player has an accumulating history of observations H t i = {O v 1 , A v 1 } t v=1 . Call the set of all histories H i = ∪ ∞ t=1 H t i . Assume for simplicity that the initial observation history is fixed and common knowledge, h 0 1 = h 0 2 ≡ h 0 . Finally, principals choose among policies π i : H → ∆(A i ), which we imagine as artificial agents deployed by the principals. We will refer to policy profiles generically as π ∈ Π := Π 1 × Π 2 . Each player has a reward function r i , such that r i (S t , A t 1 , A t 2 ) is their reward at time t. Define the value to player i of policy profile π starting at history h t i as V i (h t i , π) = E π [ ∞ v=t γ v−t r i (S v , A v 1 , A v 2 ) | H t i = h t i ], where γ < 1 is a discount factor. Define the value of a policy profile to player i as V i (π) = V i (h 0 , π). A payoff profile is then a tuple (V 1 (π), V 2 (π)). We say that π is a (Nash) equilibrium of a POSG if π i ∈ arg max π i ∈Πi V i (h 0 , π i , π −i ) for i = 1, 2. We say that π is Pareto-optimal if for i = 1, 2 and π ∈ Π, we have that V i (π ) > V i (π) implies V −i (π ) < V −i (π).
 Mixed-motive coordination problems We envision multiple actors ("principals") deploying machine learning systems, defined by policies, into the world. The principals may be human developers, or they may be machine learning systems themselves. The principals simultaneously deploy policies, without communication (in Section 6 we discuss ways in which communication may change the analysis). Let W be a set of welfare functions, which are supposed to encode normative properties (total welfare, fairness, etc.) of different payoff profiles. Two uncontroversial properties of a welfare function are Pareto-optimality (i.e., its optimizer should be Pareto-optimal) and symmetry (the welfare of a policy profile should be invariant to permutations of player indices). Beyond that, there are several appealing properties which cannot all be satisfied by the same welfare function (see Table  1 ). The experimental evidence regarding people's preferences over welfare functions is inconclusive  [Felsenthal and Diskin, 1982] . We will say that a game is a mixed-motive coordination problem (MCP) if it has several equilibria 2 which are optimal with respect to some welfare function, but are incompatible with each other. Definition 3.1 (Mixed-motive coordination problem). Let G be a POSG. Let W be a set of welfare functions, and let W E ⊆ W be the set of welfare functions in W which are optimized by some equilibrium of G. Then G is a mixed-motive coordination problem (with respect to W) if there exist distinct w, w ∈ W E 2 such that (π w 1 , π w 2 ) has a Pareto-suboptimal policy profile for all π w ∈ arg max π w(π ), π w ∈ arg max π w (π ). An example of a mixed-motive coordination problem is the Bach or Stravinsky game. A variant of this game with asymmetric payoffs is given in Figure  1  (left). The folk theorem  [Friedman, 1971]  tells us that any payoff profile in the convex hull of these payoffs is attainable in an equilibrium of the repeated game, for players with sufficiently high discount factors. And, as is suggested by the shape of the Pareto frontier for Name of welfare function w
 Form of w(π) Nash  [Nash, 1950]  [ V 1 (π) − d 1 ] • [V 2 (π) − d 2 ] Kalai-Smorodinsky  [Kalai and Smorodinsky, 1975]   V1(π)−d1 V2(π)−d2 − sup π V1(π)−d1 sup π V2(π)−d2 + ι {π Pareto-optimal } Egalitarian  [Kalai, 1977]  min {V 1 (π) − d 1 , V 2 (π) − d 2 , } Utilitarian (e.g. Harsanyi 1955) V 1 (π) + V 2 (π) Table  1 : Welfare functions, adapted to the multi-agent RL setting where two agents with value functions V 1 , V 2 are bargaining over the policy profile π. to deploy. The function ι in the definition of the Kalai-Smorodinsky welfare is the ∞-0 indicator, used to enforce the constraint in its argument. V d i = V i (π d ) are the players' disagreement values. Note that in bargaining problems in which the set of feasible payoffs is convex, the Nash welfare function uniquely satisfies the properties of (1) Pareto optimality, (2) symmetry, (3) invariance to affine transformations, and (4) independence of irrelevant alternatives. See  Zhou [1997]  for a characterization of the Nash welfare in non-convex bargaining problems. The Nash welfare can also be obtained as the subgame perfect equilibrium of an alternating-offers game as the "patience" of the players goes to infinity  [Binmore et al., 1986] . On the other hand, Kalai-Smorodinsky uniquely satisfies (1)-(  3 ) plus (  5 ) resource monotonicity, which means that all players are weakly better off when there are more resources to go around. The egalitarian solution instead satisfies (1), (  2 ), (4), and (5). The utilitarian welfare function is implicitly used in the work of  Peysakhovich and Lerer [2017a] ,  Lerer and Peysakhovich [2017] ,  Wang et al. [2018]  on cooperation in sequential social dilemmas. BoS B S B 4, 1 0, 0 S 0, 0 2, 2 PD C D C -1, -1 -3, 0 D 0, -3 -2, -2 Figure  1 : Payoffs for asymmetric Bach or Stravinsky (BoS) and the Prisoner's Dilemma (PD). BoS (Figure  2 , left), different welfare functions pick out different Pareto-optimal payoff profiles. For instance, taking the disagreement point to be d 1 , d 2 = (0, 0), the Nash welfare is optimized at (3.0, 1.5), whereas the egalitarian welfare is optimized at (2.0, 2.0). On the other hand, all welfare functions in Table  1  agree on the payoffs corresponding to (C, C) in the iterated Prisoner's Dilemma (Figure  2 , right). In the following, we compare mixed-motive coordination problems to the more general equilibrium selection problem and to sequential social dilemmas.
 Relation to other cooperation problems Equilibrium and Pareto selection problems It is well-known that achieving cooperation between independently-trained agents this is a challenging problem for MAL in general (e.g.,  Boutilier 1999 , Matignon et al. 2012 , Hu et al. 2020 ). One reason is that there are often multiple equilibria, such that (to the extent that one wants to deploy a policy which is a best-response to that of one's counterpart) the choice of a policy constitutes an equilibrium selection problem. Such problems even arise in the fully cooperative case, i.e., when V 1 = V 2 ≡ V , and are then also called Pareto selection problems  [Matignon et al., 2012]  or pure coordination problems. In a pure coordination problem, any π ∈ arg max π V (π) has an optimal value and is thus in equilibrium. There is an equilibrium selection problem if there exist π, π ∈ arg max π V (π) such that (π 1 , π 2 ) has Pareto-suboptimal payoff profile. Mixed-motive coordination problems that occur between policy profiles in equilibrium are a subclass of equilibrium selection problems. They only occur in general-sum, non-fully cooperative cases and pose a challenge which is qualitatively different from the fully cooperative case: players have to coordinate among options over which they have conflicting preferences. The pure coordination problem can be solved by merely giving the players a chance to converge upon a coordinated policy profile (as in  Boutilier [1999] , for instance) and any such policy will do equally well. This is not so in mixed-motive coordination problems, since agents that are open to converging upon any Pareto-optimal outcome would be exploitable by opponents that insist on an outcome they prefer. (See Example 6.0.1 for a concrete illustration of this difference.) Sequential social dilemmas Sequential social dilemmas (SSDs), as defined by  Leibo et al. [2017] , are POSGs whose policy space contains subsets corresponding to cooperation and defection, called Π C i and Π D i , respectively. At each state, the normal form game whose payoffs are value functions under profiles of policies in Π C i , Π D i satisfies the social dilemma inequalities  [Macy and Flache, 2002] , where atomic "cooperate" and "defect" actions are replaced with elements of Π C i , Π D i . An SSD is not necessarily an equilibrium selection problem. For instance, any Pareto-optimal and symmetric policy profile will select an equilibrium profile that leads to constant cooperation in the iterated Prisoner's Dilemma, and which is compatible with other such profiles. But other SSDs may exhibit multiple incompatible ways of implementing cooperation. SSDs with symmetric Pareto-optimal equilibria are in any case not mixed-motive coordination problems. This includes the Gathering, Wolfpack  [Leibo et al., 2017] , and Coin Game  [Peysakhovich and Lerer, 2017a]  environments. For instance, in the Coin Game, the single cooperative payoff profile is attained when agents get only the coin of their own color, rather than attempting to take their counterpart's coin as well.
 Learning algorithms Play between policies trained by different runs or variations of an algorithm, or "cross-play"  [Hu et al., 2020]  is problematic for MCPs. To illustrate the need for evaluating policies in cross-play, we train policies using two multi-agent learning algorithms which have been shown to achieve cooperation in SSDs. The first is a variant of  Lerer and Peysakhovich [2017] 's approximate Markov tit-for-tat (amTFT). The original amTFT algorithm trains a cooperative policy on the utilitarian welfare (Table  1 ), as well as a punishment policy, and switches from the cooperative to the punishment policy when it detects that the other player is defecting from the cooperative policy. We consider the more general class of algorithms in which a cooperative policy is constructed by optimizing some welfare function. Call the version of amTFT in which the cooperative policy optimizes welfare function w as amTFT(w). In our experiments, we use the utilitarian welfare function w Util and an inequity averse welfare function w IA . Following  Hughes et al. [2018] , this welfare function is constructed using modified reward functions which penalize inequitable outcomes: V IA i (a t 1 , a t 2 ; α, β) = r 1 (a t 1 , a t 2 ) + r 1 (a t 2 , a t 2 ) − α[e t 1 (a t 1 , a t 2 ) − e t 2 (a t 1 , a t 2 )] + − β[e 2 (a t 1 , a t 2 ) − e 1 (a t 1 , a t 2 )] + where [x] + = max{x, 0} and e t i are smoothed cumulative rewards computed as e t i (a t 1 , a t 2 ) = γλe t−1 i (a t−1 1 , a t−1 2 )+ r i (a t 1 , a t 2 ) with hyperparameter λ. and discount factor γ α and β control how much unequal outcomes are penalized.  3  The second learning algorithm is learning with opponent-learning awareness (LOLA;  Foerster et al. 2018 ) a policy gradient method in which learners attempt to shape the others' learning through the choice of their own policy updates. Write V i (θ 1 , θ 2 ) as the value to player i under a profile of policies with parameters θ 1 , θ 2 . Then, the LOLA update for player 1 at time t with parameters δ, η > 0 is θ t 1 = θ t−1 1 + δ∇ θ1 V i (θ 1 , θ 2 ) + δη [∇ θ2 V 1 (θ 1 , θ 2 )] T ∇ θ1 ∇ θ2 V 2 (θ 1 , θ 2 ). We use the discount factor γ = 0.96 throughout.
 Results Here we illustrate the special problem posed by MCPs by comparing the performance of the learning algorithms above in iterated BoS (IBoS) and in the iterated Prisoner's Dilemma (IPD). The folk theorems  [Mailath et al., 2006]  tell us that, for players with sufficiently high discount factors, any payoff profile of the stage game in which players get at least as much as they can unilaterally guarantee themselves can be attained in an equilibrium of the repeated game. IBoS is an MCP with respect to the set of welfare functions W = {w Util , w IA }, since these welfare functions are optimized by distinct payoff profiles for sufficiently large values of α, β. We measure the performance of these algorithms in self-play and cross-play. Self-play payoffs are the average payoffs attained by a profile of policies produced by the same training run. Cross-play payoffs are the average payoffs attained by a profile of two policies produced by two different training runs. In the case of amTFT, we also consider cross-play between policies trained with amTFT(w Util ) and those trained with amTFT(w IA ).
 LOLA Following  Foerster et al. [2018] 's parameterization of policies in the iterated Prisoner's Dilemma, we used policies which condition on the previous pair of actions played, with one parameter for every possible previous action profile. To learn policy profiles, we used a discount factor of γ = 0.96 and applied LOLA to the closed-form value functions (which can be computed for each policy profile using the Bellman equations; this is called exact LOLA in  Foerster et al. [2018] ). Figure  2 : LOLA average payoff profiles for 40 training runs in IBoS (left) and IPD (right). Cross-play payoffs were obtained by randomly sampling 40 pairs (π 1 , π 2 ), where π 1 and π 2 come from different runs. A small amount of random noise is applied to each point to enhance visibility. On average, self-play yields payoffs (2.69, 1.34), while cross-play yields (1.30, 0.70). As Figure  2  shows, the self-play payoffs almost all lie close to the empirical Pareto frontier, which consists of payoff profiles roughly equal to {(4, 1), (3, 1.5), (2, 2)}. This clustering is caused by the policies depending only on actions taken at the previous timestep, which allows for always playing (S,S), alternating between them, or always playing (B,B). These coincide with the maxima for w Util , the Nash welfare  [Nash, 1950]  under disagreement profile (0, 0), or w IA . On the other hand, the cross-play points are much more often Pareto-suboptimal. This is not surprising, as if LOLA produces tit-for-tat-like strategies which punish deviations from a particular strategy profile, then playing policies from separate runs is liable to lead to incompatible policies and therefore punishments.
 amTFT We evaluated self-and cross-play for amTFT in IBoS and IPD, where each episode lasted 100 of steps. Payoff profiles from individual runs are displayed in Figure  3 , and average payoffs under self-and cross-play for each combination of welfare functions are displayed in Table  2 . Table 2: Self-play ("•") and cross-play ("×") payoffs in IBoS for amTFT(w Util ) and amTFT(w IA ). Payoffs are averages over 40 training runs. w Util × w IA leads to the same payoff regardless of which player operates under which welfare function. All standard errors were below 0.01. Of course, the goal of developing benchmarks is not only to point out flaws in existing approaches, but also to help develop better ones. In this section we discuss several research directions for ameliorating MCPs, and give a toy illustration of one approach.
 Policy profile w Util • w Util w IA • w IA w Util × w Util w IA × w IA w Util × Coordination by the principals MCPs can be solved if the principals can agree on a way of selecting among policy profiles to deploy. For instance, principals might agree on a welfare function, and choose learning algorithms which converge to an optimizer of that welfare function. Coordination on a welfare function (and among equilibria which all maximize that welfare function) may be easy in toy environments such as the ones we have considered. But as systems become more complex, it will likely be intractable for principals to fully coordinate so as to train welfare-optimal agents. Thus there is a need to develop methods for low-bandwidth coordination between principals training agents in MCPs which captures most of the benefits of full coordination. Bandwidth-constrained coordination has already been studied in fully cooperative games (e.g.,  Stone and Veloso 1999, Zhang et al. 2019) . Compromise between welfare functions We have not discussed how giving agents the opportunity to communicate post-deployment might ameliorate MCPs. One possibility is to have them reach a compromise over different welfare-optimal policy profiles before deciding which to deploy. On possibility is for the train to several policies corresponding to several welfare functions according to which they are willing to play. When given the chance to communicate, their agents can then apply a bargaining protocol to this reduced set of policy profiles. Here is a simple example of such a protocol. Agents first announce the sets of policy profiles they are willing to play. If these sets overlap, a policy profile is taken at random from the intersection (compare with  Chakravorti and Conley [2004] 's discussion of randomization between the Nash and Kalai-Smorodinsky bargaining solutions in the asymmetric iterated Prisoner's Dilemma). Only when these sets do not overlap do agents revert to a default policy. We give an illustration in Example 6.0.1. Example 6.0.1 (Bargaining over amTFT policies). An example of a protocol for bargaining over welfareoptimal policy profiles is given in Algorithm 1. Policies are obtained by applying amTFT to a welfare function w, returning a policy amTFT(w) i for player i. Refer to the strategy which reports welfare functions W i and default policy π d i as (W i , π d i ).  {(W 1 , π d 1 ), (W 2 , π d 2 )} in this bargaining protocol applied to IBoS. This payoff matrix clearly shows the tradeoff between exploitability and robustness to independentlytrained policies. The strategies which play {w Util , w IA } are of more robust, in the sense that they avoid the mutually dispreferred profile of (1.0, 0.3). However, these strategies are not played in equilibrium, because they are exploitable: an agent can deviate from this profile by instead playing only their preferred welfare function. Nevertheless, one has the intuition that it is reasonable to play {w Util , w IA } in view of its robustness, despite the fact that it is exploitable. On the other hand, the situation becomes more complicated when the set of welfare functions under consideration is much larger. Identifying reasonable solution concepts for this setting is an important direction for future work. Besides the fact that avoiding bargaining failure is still not guaranteed, there are additional problems for this protocol. First, it may be ex-ante Pareto suboptimal, because randomization between Pareto optimal payoff profiles may be ex-ante Pareto suboptimal. Second, while it assumes that agents agree to uniform randomization over the intersection of their sets, agents will in fact have differing preferences over different distributions with which to randomize, generating another bargaining problem. Less ad-hoc solutions might be found by using a dynamic bargaining protocol (as in the classical Rubinstein bargaining game  Rubinstein 1982) , which may narrow down the set of equilibria. As a final note, compare this protocol with an analogous protocol for a pure coordination problem (Figure  4 ). In this protocol, agents play a default action but also announce a set of actions that they are willing to play. Again, actions are taken uniformly at random from the overap of the sets if it is nonempty, and otherwise the default action is taken. Clearly playing {A, B} is a dominant strategy for each player, unlike the analogous strategy in the protocol for IBoS, where greater flexibility comes at the price of greater exploitability.    Figure 3 : 3 Figure 3: amTFT payoff profiles for 40 training runs in IBoS (left) and IPD (right). A small amount of random noise is applied to each point to enhance visibility.
 Figure 4 : 4 Figure 4: Payoffs for a pure coordination problem.
 Payoff matrix for agents playing Pure Coordination and following a variant of the bargaining protocol in 1 where sets of rather than welfare functions are announced.
 Table 3 displays the payoff table of the game with strategy profiles
 Algorithm 1: Welfare function selection protocol Input: Sets of welfare functions W i , default policies π d i // Independently train policies to maximize welfare functionsfor i = 1, 2 do Π i ← ∅ for w ∈ W i do π w i ← amTFT(w) i Π i ← Π i ∪ {π w i } //Communicate and decide what policies to deployPlayers announce W 1 , W 2 if W 1 ∩ W 2 = ∅ then w ∼ Unif(W 1 ∩ W 2 ) π 1 , π 2 ← π w 1 , π w 2 else π 1 , π 2 ← π d 1 , π d 2 return π 1 , π 2
 Table 3 : 3 Empirical payoff matrix for agents playing IBoS and following bargaining protocol in Algorithm 1 with welfare functions in {w Util , w IA }.{w Util }, π IA 2 {w IA }, π IA 2 {w Util , w IA }, π IA 2 {w Util }, π Util
			 The spirit of an MCP is that independent training predictably leads to policy profiles which optimize different welfare functions, but are incompatible. So one could modify the requirement that welfare-optimal policies be equilibria and retain this spirit. One could replace the requirement of Nash equilibrium with a different or a broader set of solution concepts, e.g., the recently-proposed Markov-Conley chains used in the α-rank evaluation methodology [Omidshafiei et al., 2019] .
			 We choose the utilitarian and inequity-averse welfare functions largely for convenience, and not because they have particularly compelling normative properties. Specifically, these welfare functions admit a Bellman equation, allowing us to straightforwardly apply reinforcement learning methods in order to maximize them. Methods for optimizing other welfare functions (Table1), which may be less amenable to off-the-shelf reinforcement learning methods, are a direction for future work.