cccat6's picture
Initial FlowMo-WM public code release
604e535 verified
raw
history blame
617 Bytes
"""MPC with online current estimate."""
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, current_estimate, 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(current_estimate, dtype=np.float32),
drift_gain=float(config["current_gain"]),
)