| import sys | |
| import os | |
| ROOT_DIR = os.path.dirname(os.path.dirname(__file__)) | |
| sys.path.append(ROOT_DIR) | |
| os.chdir(ROOT_DIR) | |
| from diffusion_policy.env_runner.robomimic_image_runner import RobomimicImageRunner | |
| def test(): | |
| import os | |
| from omegaconf import OmegaConf | |
| cfg_path = os.path.expanduser('~/dev/diffusion_policy/diffusion_policy/config/task/lift_image.yaml') | |
| cfg = OmegaConf.load(cfg_path) | |
| cfg['n_obs_steps'] = 1 | |
| cfg['n_action_steps'] = 1 | |
| cfg['past_action_visible'] = False | |
| runner_cfg = cfg['env_runner'] | |
| runner_cfg['n_train'] = 1 | |
| runner_cfg['n_test'] = 1 | |
| del runner_cfg['_target_'] | |
| runner = RobomimicImageRunner( | |
| **runner_cfg, | |
| output_dir='/tmp/test') | |
| # import pdb; pdb.set_trace() | |
| self = runner | |
| env = self.env | |
| env.seed(seeds=self.env_seeds) | |
| obs = env.reset() | |
| for i in range(10): | |
| _ = env.step(env.action_space.sample()) | |
| imgs = env.render() | |
| if __name__ == '__main__': | |
| test() | |