import torch


def observationToCNNInput(observation, visual_width=320, visual_height=240):
    """
    Transforms the observation of the environment into
    an object with a suitable shape for the CNN
    """

    joint_pos     = torch.FloatTensor(observation["joint_positions"])
    touch_sensors = torch.FloatTensor(observation["touch_sensors"])
    retina        = torch.FloatTensor(observation["retina"])
    retina        = torch.reshape(retina, (-1, 3, visual_height, visual_width))
    goal          = torch.FloatTensor(observation["goal"])
    goal          = torch.reshape(goal, (-1, 3, visual_height, visual_width))

    return retina, torch.cat((joint_pos, touch_sensors)), goal