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)