chore: import upstream snapshot with attribution

rllib/algorithms/dreamerv3/utils/debugging.py

@@ -0,0 +1,189 @@
+import gymnasium as gym
+import numpy as np
+from gymnasium.envs.classic_control.cartpole import CartPoleEnv
+from PIL import Image, ImageDraw
+
+from ray.rllib.utils.framework import try_import_torch
+
+torch, _ = try_import_torch()
+
+
+class CartPoleDebug(CartPoleEnv):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+        low = np.concatenate([np.array([0.0]), self.observation_space.low])
+        high = np.concatenate([np.array([1000.0]), self.observation_space.high])
+
+        self.observation_space = gym.spaces.Box(low, high, shape=(5,), dtype=np.float32)
+
+        self.timesteps_ = 0
+        self._next_action = 0
+        self._seed = 1
+
+    def reset(self, *, seed=None, options=None):
+        ret = super().reset(seed=self._seed)
+        self._seed += 1
+        self.timesteps_ = 0
+        self._next_action = 0
+        obs = np.concatenate([np.array([self.timesteps_]), ret[0]])
+        return obs, ret[1]
+
+    def step(self, action):
+        ret = super().step(self._next_action)
+
+        self.timesteps_ += 1
+        self._next_action = 0 if self._next_action else 1
+
+        obs = np.concatenate([np.array([self.timesteps_]), ret[0]])
+        reward = 0.1 * self.timesteps_
+        return (obs, reward) + ret[2:]
+
+
+gym.register("CartPoleDebug-v0", CartPoleDebug)
+cartpole_env = gym.make("CartPoleDebug-v0", render_mode="rgb_array")
+cartpole_env.reset()
+
+frozenlake_env = gym.make(
+    "FrozenLake-v1", render_mode="rgb_array", is_slippery=False, map_name="4x4"
+)  # desc=["SF", "HG"])
+frozenlake_env.reset()
+
+
+def create_cartpole_dream_image(
+    dreamed_obs,  # real space (not symlog'd)
+    dreamed_V,  # real space (not symlog'd)
+    dreamed_a,
+    dreamed_r_tp1,  # real space (not symlog'd)
+    dreamed_ri_tp1,  # intrinsic reward
+    dreamed_c_tp1,  # continue flag
+    value_target,  # real space (not symlog'd)
+    initial_h,
+    as_tensor=False,
+):
+    # CartPoleDebug
+    if dreamed_obs.shape == (5,):
+        # Set the state of our env to the given observation.
+        cartpole_env.unwrapped.state = np.array(dreamed_obs[1:], dtype=np.float32)
+    # Normal CartPole-v1
+    else:
+        cartpole_env.unwrapped.state = np.array(dreamed_obs, dtype=np.float32)
+
+    # Produce an RGB-image of the current state.
+    rgb_array = cartpole_env.render()
+
+    # Add value-, action-, reward-, and continue-prediction information.
+    image = Image.fromarray(rgb_array)
+    draw_obj = ImageDraw.Draw(image)
+
+    # fnt = ImageFont.load_default(size=40)
+
+    draw_obj.text(
+        (5, 6), f"Vt={dreamed_V:.2f} (Rt={value_target:.2f})", fill=(0, 0, 0)
+    )  # , font=fnt.font, size=30)
+    draw_obj.text(
+        (5, 18),
+        f"at={'<--' if dreamed_a == 0 else '-->'} ({dreamed_a})",
+        fill=(0, 0, 0),
+    )
+    draw_obj.text((5, 30), f"rt+1={dreamed_r_tp1:.2f}", fill=(0, 0, 0))
+    if dreamed_ri_tp1 is not None:
+        draw_obj.text((5, 42), f"rit+1={dreamed_ri_tp1:.6f}", fill=(0, 0, 0))
+    draw_obj.text((5, 54), f"ct+1={dreamed_c_tp1}", fill=(0, 0, 0))
+    draw_obj.text((5, 66), f"|h|t={np.mean(np.abs(initial_h)):.5f}", fill=(0, 0, 0))
+
+    if dreamed_obs.shape == (5,):
+        draw_obj.text((20, 100), f"t={dreamed_obs[0]}", fill=(0, 0, 0))
+
+    # Return image.
+    np_img = np.asarray(image)
+    if as_tensor:
+        return torch.from_numpy(np_img, dtype=torch.uint8)
+    return np_img
+
+
+def create_frozenlake_dream_image(
+    dreamed_obs,  # real space (not symlog'd)
+    dreamed_V,  # real space (not symlog'd)
+    dreamed_a,
+    dreamed_r_tp1,  # real space (not symlog'd)
+    dreamed_ri_tp1,  # intrinsic reward
+    dreamed_c_tp1,  # continue flag
+    value_target,  # real space (not symlog'd)
+    initial_h,
+    as_tensor=False,
+):
+    frozenlake_env.unwrapped.s = np.argmax(dreamed_obs, axis=0)
+
+    # Produce an RGB-image of the current state.
+    rgb_array = frozenlake_env.render()
+
+    # Add value-, action-, reward-, and continue-prediction information.
+    image = Image.fromarray(rgb_array)
+    draw_obj = ImageDraw.Draw(image)
+
+    draw_obj.text((5, 6), f"Vt={dreamed_V:.2f} (Rt={value_target:.2f})", fill=(0, 0, 0))
+    action_arrow = (
+        "<--"
+        if dreamed_a == 0
+        else "v"
+        if dreamed_a == 1
+        else "-->"
+        if dreamed_a == 2
+        else "^"
+    )
+    draw_obj.text((5, 18), f"at={action_arrow} ({dreamed_a})", fill=(0, 0, 0))
+    draw_obj.text((5, 30), f"rt+1={dreamed_r_tp1:.2f}", fill=(0, 0, 0))
+    if dreamed_ri_tp1 is not None:
+        draw_obj.text((5, 42), f"rit+1={dreamed_ri_tp1:.6f}", fill=(0, 0, 0))
+    draw_obj.text((5, 54), f"ct+1={dreamed_c_tp1}", fill=(0, 0, 0))
+    draw_obj.text((5, 66), f"|h|t={np.mean(np.abs(initial_h)):.5f}", fill=(0, 0, 0))
+
+    # Return image.
+    np_img = np.asarray(image)
+    if as_tensor:
+        return torch.from_numpy(np_img, dtype=torch.uint8)
+    return np_img
+
+
+if __name__ == "__main__":
+    # CartPole debug.
+    rgb_array = create_cartpole_dream_image(
+        dreamed_obs=np.array([100.0, 1.0, -0.01, 1.5, 0.02]),
+        dreamed_V=4.3,
+        dreamed_a=1,
+        dreamed_r_tp1=1.0,
+        dreamed_c_tp1=True,
+        initial_h=0.0,
+        value_target=8.0,
+    )
+    # ImageFont.load("arial.pil")
+    image = Image.fromarray(rgb_array)
+    image.show()
+
+    # Normal CartPole.
+    rgb_array = create_cartpole_dream_image(
+        dreamed_obs=np.array([1.0, -0.01, 1.5, 0.02]),
+        dreamed_V=4.3,
+        dreamed_a=1,
+        dreamed_r_tp1=1.0,
+        dreamed_c_tp1=True,
+        initial_h=0.1,
+        value_target=8.0,
+    )
+    # ImageFont.load("arial.pil")
+    image = Image.fromarray(rgb_array)
+    image.show()
+
+    # Frozenlake
+    rgb_array = create_frozenlake_dream_image(
+        dreamed_obs=np.array([1.0] + [0.0] * (frozenlake_env.observation_space.n - 1)),
+        dreamed_V=4.3,
+        dreamed_a=1,
+        dreamed_r_tp1=1.0,
+        dreamed_c_tp1=True,
+        initial_h=0.1,
+        value_target=8.0,
+    )
+    image = Image.fromarray(rgb_array)
+    image.show()