ray-project--ray/rllib/algorithms/dreamerv3/utils/summaries.py

"""
[1] Mastering Diverse Domains through World Models - 2023
D. Hafner, J. Pasukonis, J. Ba, T. Lillicrap
https://arxiv.org/pdf/2301.04104v1.pdf

[2] Mastering Atari with Discrete World Models - 2021
D. Hafner, T. Lillicrap, M. Norouzi, J. Ba
https://arxiv.org/pdf/2010.02193.pdf
"""
import numpy as np

from ray.rllib.algorithms.dreamerv3.utils.debugging import (
    create_cartpole_dream_image,
    create_frozenlake_dream_image,
)
from ray.rllib.core import DEFAULT_MODULE_ID
from ray.rllib.core.columns import Columns
from ray.rllib.utils.framework import try_import_torch
from ray.rllib.utils.metrics import (
    LEARNER_RESULTS,
    REPLAY_BUFFER_RESULTS,
)
from ray.rllib.utils.torch_utils import inverse_symlog

torch, _ = try_import_torch()


def reconstruct_obs_from_h_and_z(
    h_t0_to_H,
    z_t0_to_H,
    dreamer_model,
    obs_dims_shape,
    framework="torch",
):
    """Returns"""
    shape = h_t0_to_H.shape
    T = shape[0]  # inputs are time-major
    B = shape[1]
    # Compute actual observations using h and z and the decoder net.
    # Note that the last h-state (T+1) is NOT used here as it's already part of
    # a new trajectory.
    # Use mean() of the Gaussian, no sample! -> No need to construct dist object here.
    if framework == "torch":
        device = next(iter(dreamer_model.world_model.decoder.parameters())).device
        reconstructed_obs_distr_means_TxB = (
            dreamer_model.world_model.decoder(
                # Fold time rank.
                h=torch.from_numpy(h_t0_to_H).reshape((T * B, -1)).to(device),
                z=torch.from_numpy(z_t0_to_H)
                .reshape((T * B,) + z_t0_to_H.shape[2:])
                .to(device),
            )
            .detach()
            .cpu()
            .numpy()
        )
    else:
        reconstructed_obs_distr_means_TxB = dreamer_model.world_model.decoder(
            # Fold time rank.
            h=h_t0_to_H.reshape((T * B, -1)),
            z=z_t0_to_H.reshape((T * B,) + z_t0_to_H.shape[2:]),
        )

    # Unfold time rank again.
    reconstructed_obs_T_B = np.reshape(
        reconstructed_obs_distr_means_TxB, (T, B) + obs_dims_shape
    )
    # Return inverse symlog'd (real env obs space) reconstructed observations.
    return reconstructed_obs_T_B


def report_dreamed_trajectory(
    *,
    results,
    env,
    dreamer_model,
    obs_dims_shape,
    batch_indices=(0,),
    desc=None,
    include_images=True,
    framework="torch",
):
    if not include_images:
        return

    dream_data = results["dream_data"]
    dreamed_obs_H_B = reconstruct_obs_from_h_and_z(
        h_t0_to_H=dream_data["h_states_t0_to_H_BxT"],
        z_t0_to_H=dream_data["z_states_prior_t0_to_H_BxT"],
        dreamer_model=dreamer_model,
        obs_dims_shape=obs_dims_shape,
        framework=framework,
    )
    func = (
        create_cartpole_dream_image
        if env.startswith("CartPole")
        else create_frozenlake_dream_image
    )
    # Take 0th dreamed trajectory and produce series of images.
    for b in batch_indices:
        images = []
        for t in range(len(dreamed_obs_H_B) - 1):
            images.append(
                func(
                    dreamed_obs=dreamed_obs_H_B[t][b],
                    dreamed_V=dream_data["values_dreamed_t0_to_H_BxT"][t][b],
                    dreamed_a=(dream_data["actions_ints_dreamed_t0_to_H_BxT"][t][b]),
                    dreamed_r_tp1=(dream_data["rewards_dreamed_t0_to_H_BxT"][t + 1][b]),
                    # `DISAGREE_intrinsic_rewards_H_B` are shifted by 1 already
                    # (from t1 to H, not t0 to H like all other data here).
                    dreamed_ri_tp1=(
                        results["DISAGREE_intrinsic_rewards_H_BxT"][t][b]
                        if "DISAGREE_intrinsic_rewards_H_BxT" in results
                        else None
                    ),
                    dreamed_c_tp1=(
                        dream_data["continues_dreamed_t0_to_H_BxT"][t + 1][b]
                    ),
                    value_target=results["VALUE_TARGETS_H_BxT"][t][b],
                    initial_h=dream_data["h_states_t0_to_H_BxT"][t][b],
                    as_tensor=True,
                ).numpy()
            )
        # Concat images along width-axis (so they show as a "film sequence" next to each
        # other).
        results.update(
            {
                f"dreamed_trajectories{('_'+desc) if desc else ''}_B{b}": (
                    np.concatenate(images, axis=1)
                ),
            }
        )


def report_predicted_vs_sampled_obs(
    *,
    metrics,
    sample,
    batch_size_B,
    batch_length_T,
    symlog_obs: bool = True,
    do_report: bool = True,
):
    """Summarizes sampled data (from the replay buffer) vs world-model predictions.

    World model predictions are based on the posterior states (z computed from actual
    observation encoder input + the current h-states).

    Observations: Computes MSE (sampled vs predicted/recreated) over all features.
    For image observations, also creates direct image comparisons (sampled images
    vs predicted (posterior) ones).
    Rewards: Compute MSE (sampled vs predicted).
    Continues: Compute MSE (sampled vs predicted).

    Args:
        metrics: The MetricsLogger object of the DreamerV3 algo.
        sample: The sampled data (dict) from the replay buffer. Already torch-tensor
            converted.
        batch_size_B: The batch size (B). This is the number of trajectories sampled
            from the buffer.
        batch_length_T: The batch length (T). This is the length of an individual
            trajectory sampled from the buffer.
        do_report: Whether to actually log the report (default). If this is set to
            False, this function serves as a clean-up on the given metrics, making sure
            they do NOT contain anymore any (spacious) data relevant for producing
            the report/videos.
    """
    fwd_output_key = (
        LEARNER_RESULTS,
        DEFAULT_MODULE_ID,
        "WORLD_MODEL_fwd_out_obs_distribution_means_b0xT",
    )
    # logged as a non-reduced item (still a list)
    predicted_observation_means_single_example = metrics.peek(
        fwd_output_key, default=[None]
    )[-1]
    metrics.delete(fwd_output_key, key_error=False)

    final_result_key = (
        f"WORLD_MODEL_sampled_vs_predicted_posterior_b0x{batch_length_T}_videos"
    )
    if not do_report:
        metrics.delete(final_result_key, key_error=False)
        return

    _report_obs(
        metrics=metrics,
        computed_float_obs_B_T_dims=np.reshape(
            predicted_observation_means_single_example,
            # WandB videos need to be channels first.
            (1, batch_length_T) + sample[Columns.OBS].shape[2:],
        ),
        sampled_obs_B_T_dims=sample[Columns.OBS][0:1],
        metrics_key=final_result_key,
        symlog_obs=symlog_obs,
    )


def report_dreamed_eval_trajectory_vs_samples(
    *,
    metrics,
    sample,
    burn_in_T,
    dreamed_T,
    dreamer_model,
    symlog_obs: bool = True,
    do_report: bool = True,
    framework="torch",
) -> None:
    """Logs dreamed observations, rewards, continues and compares them vs sampled data.

    For obs, we'll try to create videos (side-by-side comparison) of the dreamed,
    recreated-from-prior obs vs the sampled ones (over dreamed_T timesteps).

    Args:
        metrics: The MetricsLogger object of the DreamerV3 algo.
        sample: The sampled data (dict) from the replay buffer. Already torch-tensor
            converted.
        burn_in_T: The number of burn-in timesteps (these will be skipped over in the
            reported video comparisons and MSEs).
        dreamed_T: The number of timesteps to produce dreamed data for.
        dreamer_model: The DreamerModel to use to create observation vectors/images
            from dreamed h- and (prior) z-states.
        symlog_obs: Whether to inverse-symlog the computed observations or not. Set this
            to True for environments, in which we should symlog the observations.
        do_report: Whether to actually log the report (default). If this is set to
            False, this function serves as a clean-up on the given metrics, making sure
            they do NOT contain anymore any (spacious) data relevant for producing
            the report/videos.
    """
    dream_data = metrics.peek(
        (LEARNER_RESULTS, DEFAULT_MODULE_ID, "dream_data"),
        default={},
    )
    metrics.delete(LEARNER_RESULTS, DEFAULT_MODULE_ID, "dream_data", key_error=False)

    final_result_key_obs = f"EVALUATION_sampled_vs_dreamed_prior_H{dreamed_T}_obs"
    final_result_key_rew = (
        f"EVALUATION_sampled_vs_dreamed_prior_H{dreamed_T}_rewards_MSE"
    )
    final_result_key_cont = (
        f"EVALUATION_sampled_vs_dreamed_prior_H{dreamed_T}_continues_MSE"
    )
    if not do_report:
        metrics.delete(final_result_key_obs, key_error=False)
        metrics.delete(final_result_key_rew, key_error=False)
        metrics.delete(final_result_key_cont, key_error=False)
        return

    # Obs MSE.
    dreamed_obs_H_B = reconstruct_obs_from_h_and_z(
        h_t0_to_H=dream_data["h_states_t0_to_H_Bx1"][0],  # [0] b/c reduce=None (list)
        z_t0_to_H=dream_data["z_states_prior_t0_to_H_Bx1"][0],
        dreamer_model=dreamer_model,
        obs_dims_shape=sample[Columns.OBS].shape[2:],
        framework=framework,
    )
    t0 = burn_in_T
    tH = t0 + dreamed_T
    # Observation MSE and - if applicable - images comparisons.
    _report_obs(
        metrics=metrics,
        # WandB videos need to be 5D (B, L, c, h, w) -> transpose/swap H and B axes.
        computed_float_obs_B_T_dims=np.swapaxes(dreamed_obs_H_B, 0, 1)[
            0:1
        ],  # for now: only B=1
        sampled_obs_B_T_dims=sample[Columns.OBS][0:1, t0:tH],
        metrics_key=final_result_key_obs,
        symlog_obs=symlog_obs,
    )

    # Reward MSE.
    _report_rewards(
        metrics=metrics,
        computed_rewards=dream_data["rewards_dreamed_t0_to_H_Bx1"][0],
        sampled_rewards=sample[Columns.REWARDS][:, t0:tH],
        metrics_key=final_result_key_rew,
    )

    # Continues MSE.
    _report_continues(
        metrics=metrics,
        computed_continues=dream_data["continues_dreamed_t0_to_H_Bx1"][0],
        sampled_continues=(1.0 - sample["is_terminated"])[:, t0:tH],
        metrics_key=final_result_key_cont,
    )


def report_sampling_and_replay_buffer(*, metrics, replay_buffer):
    episodes_in_buffer = replay_buffer.get_num_episodes()
    ts_in_buffer = replay_buffer.get_num_timesteps()
    replayed_steps = replay_buffer.get_sampled_timesteps()
    added_steps = replay_buffer.get_added_timesteps()

    # Summarize buffer, sampling, and train ratio stats.
    metrics.log_dict(
        {
            "capacity": replay_buffer.capacity,
            "size_num_episodes": episodes_in_buffer,
            "size_timesteps": ts_in_buffer,
            "replayed_steps": replayed_steps,
            "added_steps": added_steps,
        },
        key=REPLAY_BUFFER_RESULTS,
        window=1,
    )  # window=1 b/c these are current (total count/state) values.


def _report_obs(
    *,
    metrics,
    computed_float_obs_B_T_dims,
    sampled_obs_B_T_dims,
    metrics_key,
    symlog_obs,
):
    """Summarizes computed- vs sampled observations: MSE and (if applicable) images.

    Args:
        metrics: The MetricsLogger object of the DreamerV3 algo.
        computed_float_obs_B_T_dims: Computed float observations
            (not clipped, not cast'd). Shape=(B, T, [dims ...]).
        sampled_obs_B_T_dims: Sampled observations (as-is from the environment, meaning
            this could be uint8, 0-255 clipped images). Shape=(B, T, [dims ...]).
        metrics_key: The metrics key (or key sequence) under which to log ths resulting
            video sequence.
        symlog_obs: Whether to inverse-symlog the computed observations or not. Set this
            to True for environments, in which we should symlog the observations.
    """
    # Videos: Create summary, comparing computed images with actual sampled ones.
    # 4=[B, T, w, h] grayscale image; 5=[B, T, w, h, C] RGB image.
    if len(sampled_obs_B_T_dims.shape) in [4, 5]:
        # WandB videos need to be channels first.
        transpose_axes = (
            (0, 1, 4, 2, 3) if len(sampled_obs_B_T_dims.shape) == 5 else (0, 3, 1, 2)
        )

        if symlog_obs:
            computed_float_obs_B_T_dims = inverse_symlog(computed_float_obs_B_T_dims)

        # Restore image pixels from normalized (non-symlog'd) data.
        if not symlog_obs:
            computed_float_obs_B_T_dims = (computed_float_obs_B_T_dims + 1.0) * 128
            sampled_obs_B_T_dims = (sampled_obs_B_T_dims + 1.0) * 128
            sampled_obs_B_T_dims = np.clip(sampled_obs_B_T_dims, 0.0, 255.0).astype(
                np.uint8
            )
            sampled_obs_B_T_dims = np.transpose(sampled_obs_B_T_dims, transpose_axes)
        computed_images = np.clip(computed_float_obs_B_T_dims, 0.0, 255.0).astype(
            np.uint8
        )
        computed_images = np.transpose(computed_images, transpose_axes)
        # Concat sampled and computed images along the height axis (3) such that
        # real images show below respective predicted ones.
        # (B, T, C, h, w)
        sampled_vs_computed_images = np.concatenate(
            [computed_images, sampled_obs_B_T_dims],
            axis=-1,  # concat on width axis (looks nicer)
        )
        # Add grayscale dim, if necessary.
        if len(sampled_obs_B_T_dims.shape) == 2 + 2:
            sampled_vs_computed_images = np.expand_dims(sampled_vs_computed_images, -1)

        metrics.log_value(
            metrics_key,
            sampled_vs_computed_images,
            reduce="item_series",  # No reduction, we want the obs tensor to stay in-tact.
            window=1,
        )


def _report_rewards(
    *,
    metrics,
    computed_rewards,
    sampled_rewards,
    metrics_key,
):
    mse_sampled_vs_computed_rewards = np.mean(
        np.square(computed_rewards - sampled_rewards)
    )
    mse_sampled_vs_computed_rewards = np.mean(mse_sampled_vs_computed_rewards)
    metrics.log_value(
        metrics_key,
        mse_sampled_vs_computed_rewards,
        window=1,
    )


def _report_continues(
    *,
    metrics,
    computed_continues,
    sampled_continues,
    metrics_key,
):
    # Continue MSE.
    mse_sampled_vs_computed_continues = np.mean(
        np.square(
            computed_continues - sampled_continues.astype(computed_continues.dtype)
        )
    )
    metrics.log_value(
        metrics_key,
        mse_sampled_vs_computed_continues,
        window=1,
    )