def format_frame()

in rlstructures/env_wrappers/gymenv.py [0:0]

• 28 lines of code
• 11 McCabe index (conditional complexity)

def format_frame(frame):
    if isinstance(frame, dict):
        r = {}
        for k in frame:
            r[k] = format_frame(frame[k])
        return r
    elif isinstance(frame, list):
        t = torch.tensor(frame).unsqueeze(0)
        if t.dtype != torch.float32:
            t = t.float()
        return t
    elif isinstance(frame, np.ndarray):
        t = torch.from_numpy(frame).unsqueeze(0)
        if t.dtype != torch.float32:
            t = t.float()
        return t
    elif isinstance(frame, torch.Tensor):
        return frame.unsqueeze(0)  # .float()
    elif isinstance(frame, int):
        return torch.tensor([frame]).unsqueeze(0).long()
    elif isinstance(frame, float):
        return torch.tensor([frame]).unsqueeze(0).float()

    else:
        try:
            # Check if its a LazyFrame from OpenAI Baselines
            o = torch.from_numpy(frame._force()).unsqueeze(0).float()
            return o
        except:
            assert False