# AUTOGENERATED! DO NOT EDIT! File to edit: 01_model_utils.ipynb (unless otherwise specified).

__all__ = ['download_model', 'load_model', 'inference']

# Cell
def download_model(id, name):
    """Download a model.

    This function will download model from google drive.

    Args:
        id: The file id of google cloud.
        name: The output zip file name.

    Return:
        path: The zip model path.

    """
    dataset_url = 'https://drive.google.com/u/1/uc?id=' + id
    dataset_name = name + ".zip"
    path = os.path.join(os.getcwd(), "data", dataset_name)

    if not os.path.isdir(path):
        gdown.download(dataset_url, output = path, quiet=False)

    return path
def load_model(model_path):
    """Load the policy model.
    Args:
        model_path: The path of the policy weight.
    Return:
        path: model(policy network).
    """
    gpu = tf.config.experimental.list_physical_devices('GPU')
    tf.config.experimental.set_memory_growth(gpu[0], True)
    policy_network = tf.saved_model.load(model_path)
    return policy_network

def inference(model, state):
    """ Inference RL_vrx in gazebo or real world
    This function will use policy model and state to output the next action
    Args:
        model: policy model.
        state: observation.
    Return:
        action : The robot action[v,w].
    """
    state = tf.convert_to_tensor([state], dtype=tf.float32)
    action = model(state)[0].numpy()
    return action
