from typing import Iterable
import torch


def concat_state_action_pairs(iterator):
    """
    Overview:
        Concatenate state and action pairs from input.
    Arguments:
        - iterator (:obj:`Iterable`): Iterables with at least ``obs`` and ``action`` tensor keys.
    Returns:
        - res (:obj:`Torch.tensor`): State and action pairs.
    """
    assert isinstance(iterator, Iterable)
    res = []
    for item in iterator:
        state = item['obs'].flatten()  # to allow 3d obs and actions concatenation
        action = item['action']
        s_a = torch.cat([state, action.float()], dim=-1)
        res.append(s_a)
    return res