import torch


def _gen_mask(valid_step: torch.Tensor, batch_size: int, seq_len: int):
    """
    Mask for dealing with different lengths of MDPs

    Example:
    valid_step = [[1], [2], [3]], batch_size=3, seq_len = 4
    mask = [
        [0, 0, 0, 1],
        [0, 0, 1, 1],
        [0, 1, 1, 1],
    ]
    """
    assert valid_step.shape == (batch_size, 1)
    assert (1 <= valid_step).all()
    assert (valid_step <= seq_len).all()
    device = valid_step.device
    mask = torch.arange(seq_len, device=device).repeat(batch_size, 1)
    mask = (mask >= (seq_len - valid_step)).float()
    return mask