Skip to content

Commit

Permalink
Merge pull request #48 from MichaelTMatthews/main
Browse files Browse the repository at this point in the history
Fix craftax imports
  • Loading branch information
yuanmingqi committed May 14, 2024
2 parents b812d2a + bad7661 commit 42737c7
Showing 1 changed file with 6 additions and 6 deletions.
12 changes: 6 additions & 6 deletions rllte/env/craftax/__init__.py
Original file line number Diff line number Diff line change
@@ -1,27 +1,27 @@
from craftax.envs.craftax_pixels_env import CraftaxPixelsEnv
from craftax_classic.envs.craftax_pixels_env import CraftaxClassicPixelsEnv
from environment_base.wrappers import (
from craftax.craftax.envs.craftax_pixels_env import CraftaxPixelsEnv
from craftax.craftax_classic.envs.craftax_pixels_env import CraftaxClassicPixelsEnv
from craftax.environment_base.wrappers import (
LogWrapper,
BatchEnvWrapper,
OptimisticResetVecEnvWrapper,
)

from rllte.env.craftax.wrappers import TorchWrapper, ResizeTorchWrapper, RecordEpisodeStatistics4Craftax


def make_craftax_env(
env_id: str = "Craftax-Classic",
num_envs: int = 32,
reset_ratio: int = 16,
device: str = "cpu",
):

):
if env_id == "Craftax-Classic":
env = CraftaxClassicPixelsEnv()
elif env_id == "Craftax":
env = CraftaxPixelsEnv()
else:
raise ValueError(f"Unknown environment: {env_id}")

env = LogWrapper(env)
env = OptimisticResetVecEnvWrapper(env, num_envs=num_envs, reset_ratio=reset_ratio)
env = TorchWrapper(env, device=device)
Expand Down

0 comments on commit 42737c7

Please sign in to comment.