File size: 986 Bytes
a1fc554 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 | 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()
|