controller.py

import numpy as np


def predict(obs, gain=20.0):
    error = obs["desired_goal"] - obs["achieved_goal"]
    return np.clip(gain * error, -1.0, 1.0).astype(np.float32)