code release

tests/test_robomimic_image_runner.py

@@ -0,0 +1,38 @@
+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()