From 97b69d1449ab8ee3dbc8db0dc50b41cb0dda2102 Mon Sep 17 00:00:00 2001 From: Vincent Moens Date: Mon, 8 Apr 2024 15:16:34 +0200 Subject: [PATCH] amend --- sota-implementations/dreamer/dreamer_utils.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/sota-implementations/dreamer/dreamer_utils.py b/sota-implementations/dreamer/dreamer_utils.py index 76e61cbb3f2..e3aad6afc9e 100644 --- a/sota-implementations/dreamer/dreamer_utils.py +++ b/sota-implementations/dreamer/dreamer_utils.py @@ -84,11 +84,13 @@ def _make_env(cfg, device): env = env.append_transform( TensorDictPrimer(random=False, default_value=0, **default_dict) ) + assert env is not None return env def transform_env(cfg, env): - env = TransformedEnv(env) + if not isinstance(env, TransformedEnv): + env = TransformedEnv(env) if cfg.env.from_pixels: # transforms pixel from 0-255 to 0-1 (uint8 to float32) env.append_transform(