def normalize_obs(obs_dict, obs_normalization_stats):
    """
    Normalize observations using the provided "mean" and "std" entries 
    for each observation key. The observation dictionary will be
    modified in-place.

    Args:
        obs_dict (dict): dictionary mapping observation key to np.array or
            torch.Tensor. Leading batch dimensions are optional.

        obs_normalization_stats (dict): this should map observation keys to dicts
            with a "mean" and "std" of shape (1, ...) where ... is the default
            shape for the observation.

    Returns:
        obs_dict (dict): obs dict with normalized observation arrays
    """

    # ensure we have statistics for each modality key in the observation
    assert set(obs_dict.keys()).issubset(obs_normalization_stats)

    for m in obs_dict:
        mean = obs_normalization_stats[m]["mean"]
        std = obs_normalization_stats[m]["std"]

        # check shape consistency
        shape_len_diff = len(mean.shape) - len(obs_dict[m].shape)
        assert shape_len_diff in [0, 1], "shape length mismatch in @normalize_obs"
        assert mean.shape[shape_len_diff:] == obs_dict[m].shape, "shape mismatch in @normalize obs"

        # handle case where obs dict is not batched by removing stats batch dimension
        if shape_len_diff == 1:
            mean = mean[0]
            std = std[0]

        obs_dict[m] = (obs_dict[m] - mean) / std

    return obs_dict