We take the original PETS [1]} code (pytorch implementation used: https://github.com/quanvuong/handful-of-trials-pytorch) and modify the input and objective of the dynamics model which is done by adding the observation to group transformation for the input, and predicting coefficients for output. The input takes in group state and velocity (similar to PETS alone). Instead of predicting direct delta state, the probabilistic ensemble is trained to predict in coefficient (algebra) space as well as delta velocity. From there, we apply (REF ) to get the next predicted state from the coefficients. Though there is no two step network, just one ensemble network. This was done to not change the other structures in the original PETs code. The rest of the code is the same. GEM+PET experiment appears under the PETS folder in https://tinyurl.com/GEMMBRL. Figure  shows the comparison results for different environments.
