39 lines
986 B
Python
39 lines
986 B
Python
|
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()
|