Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[Test] Check dtypes of envs #666

Merged
merged 8 commits into from
Nov 11, 2022
Merged

[Test] Check dtypes of envs #666

Changes from 1 commit
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
b556d81
init
vmoens Nov 11, 2022
d95e79f
amend
vmoens Nov 11, 2022
0c6ed56
amend
vmoens Nov 11, 2022
b700620
lint
vmoens Nov 11, 2022
5b374cf
lint
vmoens Nov 11, 2022
31dba08
Merge remote-tracking branch 'origin/main' into check_dtype
vmoens Nov 11, 2022
6c1fb73
amend
vmoens Nov 11, 2022
5d2c037
lint
vmoens Nov 11, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Prev Previous commit
Next Next commit
amend
  • Loading branch information
@vmoens
vmoens committed Nov 11, 2022
commit d95e79f628b843f0933c51ae14f936c5ddff4336
263 changes: 147 additions & 116 deletions test/test_libs.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -82,66 +82,85 @@
[True, False],
],
)
def test_gym(env_name, frame_skip, from_pixels, pixels_only):
if env_name == PONG_VERSIONED and not from_pixels:
raise pytest.skip("already pixel")
elif (
env_name != PONG_VERSIONED
and from_pixels
and (not torch.has_cuda or not torch.cuda.device_count())
):
raise pytest.skip("no cuda device")
class TestGym:
def test_gym(self, env_name, frame_skip, from_pixels, pixels_only):
if env_name == PONG_VERSIONED and not from_pixels:
raise pytest.skip("already pixel")
elif (
env_name != PONG_VERSIONED
and from_pixels
and (not torch.has_cuda or not torch.cuda.device_count())
):
raise pytest.skip("no cuda device")

tdreset = []
tdrollout = []
final_seed = []
for _ in range(2):
env0 = GymEnv(
env_name,
frame_skip=frame_skip,
from_pixels=from_pixels,
pixels_only=pixels_only,
)
torch.manual_seed(0)
np.random.seed(0)
final_seed.append(env0.set_seed(0))
tdreset.append(env0.reset())
tdrollout.append(env0.rollout(max_steps=50))
assert env0.from_pixels is from_pixels
env0.close()
env_type = type(env0._env)
del env0

assert_allclose_td(*tdreset)
assert_allclose_td(*tdrollout)
final_seed0, final_seed1 = final_seed
assert final_seed0 == final_seed1

if env_name == PONG_VERSIONED:
base_env = gym.make(env_name, frameskip=frame_skip)
frame_skip = 1
else:
if gym_version < version.parse("0.26.0"):
base_env = gym.make(env_name)
else:
base_env = gym.make(env_name, render_mode="rgb_array")

if from_pixels and not _is_from_pixels(base_env):
base_env = PixelObservationWrapper(base_env, pixels_only=pixels_only)
assert type(base_env) is env_type
env1 = GymWrapper(base_env, frame_skip=frame_skip)
torch.manual_seed(0)
np.random.seed(0)
final_seed2 = env1.set_seed(0)
tdreset2 = env1.reset()
rollout2 = env1.rollout(max_steps=50)
assert env1.from_pixels is from_pixels
env1.close()
del env1, base_env

assert_allclose_td(tdreset[0], tdreset2, rtol=1e-4, atol=1e-4)
assert final_seed0 == final_seed2
assert_allclose_td(tdrollout[0], rollout2, rtol=1e-4, atol=1e-4)

def test_gym_fake_td(self, env_name, frame_skip, from_pixels, pixels_only):
if env_name == PONG_VERSIONED and not from_pixels:
raise pytest.skip("already pixel")
elif (
env_name != PONG_VERSIONED
and from_pixels
and (not torch.has_cuda or not torch.cuda.device_count())
):
raise pytest.skip("no cuda device")

tdreset = []
tdrollout = []
final_seed = []
for _ in range(2):
env0 = GymEnv(
env = GymEnv(
env_name,
frame_skip=frame_skip,
from_pixels=from_pixels,
pixels_only=pixels_only,
)
torch.manual_seed(0)
np.random.seed(0)
final_seed.append(env0.set_seed(0))
tdreset.append(env0.reset())
tdrollout.append(env0.rollout(max_steps=50))
assert env0.from_pixels is from_pixels
env0.close()
env_type = type(env0._env)
del env0

assert_allclose_td(*tdreset)
assert_allclose_td(*tdrollout)
final_seed0, final_seed1 = final_seed
assert final_seed0 == final_seed1

if env_name == PONG_VERSIONED:
base_env = gym.make(env_name, frameskip=frame_skip)
frame_skip = 1
else:
if gym_version < version.parse("0.26.0"):
base_env = gym.make(env_name)
else:
base_env = gym.make(env_name, render_mode="rgb_array")

if from_pixels and not _is_from_pixels(base_env):
base_env = PixelObservationWrapper(base_env, pixels_only=pixels_only)
assert type(base_env) is env_type
env1 = GymWrapper(base_env, frame_skip=frame_skip)
torch.manual_seed(0)
np.random.seed(0)
final_seed2 = env1.set_seed(0)
tdreset2 = env1.reset()
rollout2 = env1.rollout(max_steps=50)
assert env1.from_pixels is from_pixels
env1.close()
del env1, base_env

assert_allclose_td(tdreset[0], tdreset2, rtol=1e-4, atol=1e-4)
assert final_seed0 == final_seed2
assert_allclose_td(tdrollout[0], rollout2, rtol=1e-4, atol=1e-4)
_test_fake_tensordict(env)


@pytest.mark.skipif(not _has_dmc, reason="no dm_control library found")
Expand All @@ -155,77 +174,91 @@ def test_gym(env_name, frame_skip, from_pixels, pixels_only):
[False, False],
],
)
def test_dmcontrol(env_name, task, frame_skip, from_pixels, pixels_only):
if from_pixels and (not torch.has_cuda or not torch.cuda.device_count()):
raise pytest.skip("no cuda device")

tds = []
tds_reset = []
final_seed = []
for _ in range(2):
env0 = DMControlEnv(
class TestDMControl:
def test_dmcontrol(self, env_name, task, frame_skip, from_pixels, pixels_only):
if from_pixels and (not torch.has_cuda or not torch.cuda.device_count()):
raise pytest.skip("no cuda device")

tds = []
tds_reset = []
final_seed = []
for _ in range(2):
env0 = DMControlEnv(
env_name,
task,
frame_skip=frame_skip,
from_pixels=from_pixels,
pixels_only=pixels_only,
)
torch.manual_seed(0)
np.random.seed(0)
final_seed0 = env0.set_seed(0)
tdreset0 = env0.reset()
rollout0 = env0.rollout(max_steps=50)
env0.close()
del env0
tds_reset.append(tdreset0)
tds.append(rollout0)
final_seed.append(final_seed0)

tdreset1, tdreset0 = tds_reset
rollout0, rollout1 = tds
final_seed0, final_seed1 = final_seed

assert_allclose_td(tdreset1, tdreset0)
assert final_seed0 == final_seed1
assert_allclose_td(rollout0, rollout1)

env1 = DMControlEnv(
env_name,
task,
frame_skip=frame_skip,
from_pixels=from_pixels,
pixels_only=pixels_only,
)
torch.manual_seed(1)
np.random.seed(1)
final_seed1 = env1.set_seed(1)
tdreset1 = env1.reset()
rollout1 = env1.rollout(max_steps=50)
env1.close()
del env1

with pytest.raises(AssertionError):
assert_allclose_td(tdreset1, tdreset0)
assert final_seed0 == final_seed1
assert_allclose_td(rollout0, rollout1)

base_env = suite.load(env_name, task)
if from_pixels:
render_kwargs = {"camera_id": 0}
base_env = pixels.Wrapper(
base_env, pixels_only=pixels_only, render_kwargs=render_kwargs
)
env2 = DMControlWrapper(base_env, frame_skip=frame_skip)
torch.manual_seed(0)
np.random.seed(0)
final_seed0 = env0.set_seed(0)
tdreset0 = env0.reset()
rollout0 = env0.rollout(max_steps=50)
env0.close()
del env0
tds_reset.append(tdreset0)
tds.append(rollout0)
final_seed.append(final_seed0)

tdreset1, tdreset0 = tds_reset
rollout0, rollout1 = tds
final_seed0, final_seed1 = final_seed

assert_allclose_td(tdreset1, tdreset0)
assert final_seed0 == final_seed1
assert_allclose_td(rollout0, rollout1)

env1 = DMControlEnv(
env_name,
task,
frame_skip=frame_skip,
from_pixels=from_pixels,
pixels_only=pixels_only,
)
torch.manual_seed(1)
np.random.seed(1)
final_seed1 = env1.set_seed(1)
tdreset1 = env1.reset()
rollout1 = env1.rollout(max_steps=50)
env1.close()
del env1

with pytest.raises(AssertionError):
assert_allclose_td(tdreset1, tdreset0)
assert final_seed0 == final_seed1
assert_allclose_td(rollout0, rollout1)
final_seed2 = env2.set_seed(0)
tdreset2 = env2.reset()
rollout2 = env2.rollout(max_steps=50)

base_env = suite.load(env_name, task)
if from_pixels:
render_kwargs = {"camera_id": 0}
base_env = pixels.Wrapper(
base_env, pixels_only=pixels_only, render_kwargs=render_kwargs
)
env2 = DMControlWrapper(base_env, frame_skip=frame_skip)
torch.manual_seed(0)
np.random.seed(0)
final_seed2 = env2.set_seed(0)
tdreset2 = env2.reset()
rollout2 = env2.rollout(max_steps=50)
assert_allclose_td(tdreset0, tdreset2)
assert final_seed0 == final_seed2
assert_allclose_td(rollout0, rollout2)

assert_allclose_td(tdreset0, tdreset2)
assert final_seed0 == final_seed2
assert_allclose_td(rollout0, rollout2)
def test_faketd(self, env_name, task, frame_skip, from_pixels, pixels_only):
if from_pixels and (
not torch.has_cuda or not torch.cuda.device_count()):
raise pytest.skip("no cuda device")

env = DMControlEnv(
env_name,
task,
frame_skip=frame_skip,
from_pixels=from_pixels,
pixels_only=pixels_only,
)
_test_fake_tensordict(env)

@pytest.mark.skipif(
IS_OSX,
Expand Down Expand Up @@ -314,9 +347,7 @@ def test_collector_run(self, env_lib, env_args, env_kwargs, device):
class TestHabitat:
def test_habitat(self, envname):
env = HabitatEnv(envname)
print([_env for _env in env.available_envs if _env.startswith("Habitat")])
rollout = env.rollout(3)
print(rollout)
_test_fake_tensordict(env)


Expand Down