| """Oracle-flow MPC.""" | |
| import numpy as np | |
| from experiments.shared.src.control.geometric import goal_action | |
| from experiments.shared.src.vision.pose_from_image import estimate_pose_from_clean_image | |
| def plan(state, true_flow, task, config): | |
| pose = estimate_pose_from_clean_image(state) | |
| goal = np.asarray(task["goal"], dtype=np.float32) | |
| return goal_action( | |
| pose, | |
| goal, | |
| int(config["action_dim"]), | |
| float(config["gain"]), | |
| 0.6, | |
| drift=np.asarray(true_flow, dtype=np.float32), | |
| drift_gain=float(config["flow_gain"]), | |
| ) | |