Add Prod transform, Add test_factory

This commit is contained in:
Cadene
2024-02-20 14:22:16 +00:00
parent 3da6ffb2cb
commit 3dc14b5576
5 changed files with 56 additions and 12 deletions

View File

@@ -48,23 +48,16 @@ class PushtEnv(EnvBase):
if not _has_gym:
raise ImportError("Cannot import gym.")
import gym
from diffusion_policy.env.pusht.pusht_env import PushTEnv
if not from_pixels:
raise NotImplementedError("Use PushTEnv, instead of PushTImageEnv")
from diffusion_policy.env.pusht.pusht_image_env import PushTImageEnv
from gym.wrappers import TimeLimit
self._env = PushTImageEnv(render_size=self.image_size)
self._env = TimeLimit(self._env, self.max_episode_length)
# MAX_NUM_ACTIONS = 4
# num_actions = len(TASKS[self.task]["action_space"])
# self._action_space = gym.spaces.Box(low=-1.0, high=1.0, shape=(num_actions,))
# self._action_padding = np.zeros(
# (MAX_NUM_ACTIONS - num_actions), dtype=np.float32
# )
# if "w" not in TASKS[self.task]["action_space"]:
# self._action_padding[-1] = 1.0
self._make_spec()
self.set_seed(seed)