File size: 1,284 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 40 41 42 43 44 45 | 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.block_pushing.block_pushing_multimodal import BlockPushMultimodal
from gym.wrappers import FlattenObservation
from diffusion_policy.gym_util.multistep_wrapper import MultiStepWrapper
from diffusion_policy.gym_util.video_wrapper import VideoWrapper
def test():
env = MultiStepWrapper(
VideoWrapper(
FlattenObservation(
BlockPushMultimodal()
),
enabled=True,
steps_per_render=2
),
n_obs_steps=2,
n_action_steps=8,
max_episode_steps=16
)
env = BlockPushMultimodal()
obs = env.reset()
import pdb; pdb.set_trace()
env = FlattenObservation(BlockPushMultimodal())
obs = env.reset()
action = env.action_space.sample()
next_obs, reward, done, info = env.step(action)
print(obs[8:10] + action - next_obs[8:10])
import pdb; pdb.set_trace()
for i in range(3):
obs, reward, done, info = env.step(env.action_space.sample())
img = env.render()
import pdb; pdb.set_trace()
print("Done!", done)
if __name__ == '__main__':
test()
|