[Bug] Pets CartPoleEnv preprocess_fn changes the dimensions of an observation.

[MarkSelden]

Steps to reproduce

1. Run learning tasks with the pets cartpole mujoco environment with the preprocess function preprocessing model data

Observed Results

• What happened?

  File "/Users/MarkSelden/Work/Research/mbrl-lib/mbrl/util/math.py line 123, in update_stats assert data.ndim == 2 and data.shape[1] == self.mean.shape[1] AssertionError

Expected Results

• What did you expect to happen? I expected the code to run without throwing an error.

Relevant Code

  // TODO(you): code here to reproduce the problem

import hydra
import numpy as np
import omegaconf
import torch
import mbrl.env.mujoco_envs
import mbrl.env.termination_fns
import mbrl.env.reward_fns
import mbrl.algorithms.pets as pets

@hydra.main(config_path="exp/conf", config_name="main") #This config File is specific to my local machine. 
def run(cfg: omegaconf.DictConfig):
    env = mbrl.env.mujoco_envs.CartPoleEnv()
    term_fn = mbrl.env.termination_fns.no_termination
    reward_fn = mbrl.env.reward_fns.cartpole_pets
    cfg.device = 'cuda:0' if torch.cuda.is_available() else 'cpu'
    np.random.seed(cfg.seed)
    torch.manual_seed(cfg.seed)

    return pets.train(env, term_fn, reward_fn, cfg)

if __name__ == "__main__":
    run()

[MarkSelden]

I believe this error is caused by the preprocess function in mbrl/env/pets_cartpole CartPoleEnv.preprocess_fn.

This is the function's return value return np.concatenate( [ np.sin(state[..., 1:2]), np.cos(state[..., 1:2]), state[..., :1], state[..., 2:], ], axis=-1, ) the state argument for this function is of length 4, meaning this returns an array or tensor of length 5. I believe it is required that a preprocessing function does not change the observation's dimensions.

[luisenp]

Hi @MarkSelden. To use the cartpole env preprocess function you also need to set dynamics_model.in_size=6 (which overrides the default input size computation). The following command works for me

python -m mbrl.examples.main  algorithm=pets overrides=pets_cartpole_paper_version dynamics_model.in_size=6

I'll add a comment about this in the config file and also in the preprocess function.

Let me know if you keep having problems.

[luisenp]

Hi @MarkSelden, did this work for you? Let me know if I can close this issue now.

[MarkSelden]

@luisenp Thank you for the response! Feel free to close this out. Your help is very appreciated.