From 0155bd7e1bcc63f9601ffc9cb7463d5ee1bb8934 Mon Sep 17 00:00:00 2001 From: firestar5683 <168790843+firestar5683@users.noreply.github.com> Date: Tue, 2 Dec 2025 11:53:45 -0600 Subject: [PATCH] New planplus --- frogpilot/tinygrad_modeld/fill_model_msg.py | 7 +++---- frogpilot/tinygrad_modeld/parse_model_outputs.py | 12 ++++++------ frogpilot/tinygrad_modeld/tinygrad_modeld.py | 5 ++++- 3 files changed, 13 insertions(+), 11 deletions(-) diff --git a/frogpilot/tinygrad_modeld/fill_model_msg.py b/frogpilot/tinygrad_modeld/fill_model_msg.py index 2cb5b0aa3..9a3352971 100644 --- a/frogpilot/tinygrad_modeld/fill_model_msg.py +++ b/frogpilot/tinygrad_modeld/fill_model_msg.py @@ -11,16 +11,15 @@ ConfidenceClass = log.ModelDataV2.ConfidenceClass # Return curvature for lateral action. If the model outputs desired_curvature and we're not in mlsim mode, # use it directly; otherwise derive from the plan using yaw and yaw-rate. -def get_curvature_from_output(output: dict, v_ego: float, lat_action_t: float, mlsim: bool) -> float: +def get_curvature_from_output(output: dict, plan: np.ndarray, v_ego: float, lat_action_t: float, mlsim: bool) -> float: if not mlsim: desired = output.get('desired_curvature') if desired is not None: return float(desired[0, 0]) - plan_out = output['plan'][0] # Use yaw (index 2) and yaw_rate (index 2) - theta = plan_out[:, Plan.T_FROM_CURRENT_EULER][:, 2] - theta_dot = plan_out[:, Plan.ORIENTATION_RATE][:, 2] + theta = plan[:, Plan.T_FROM_CURRENT_EULER][:, 2] + theta_dot = plan[:, Plan.ORIENTATION_RATE][:, 2] return float(get_curvature_from_plan(theta, theta_dot, ModelConstants.T_IDXS, v_ego, lat_action_t)) diff --git a/frogpilot/tinygrad_modeld/parse_model_outputs.py b/frogpilot/tinygrad_modeld/parse_model_outputs.py index 836f49046..3514524b1 100644 --- a/frogpilot/tinygrad_modeld/parse_model_outputs.py +++ b/frogpilot/tinygrad_modeld/parse_model_outputs.py @@ -97,12 +97,12 @@ class Parser: self.parse_mdn('lead', outs, in_N=ModelConstants.LEAD_MHP_N, out_N=ModelConstants.LEAD_MHP_SELECTION, out_shape=(ModelConstants.LEAD_TRAJ_LEN, ModelConstants.LEAD_WIDTH)) if 'plan' in outs: - if outs['plan'].shape[1] == 2 * ModelConstants.IDX_N * ModelConstants.PLAN_WIDTH: - self.parse_mdn('plan', outs, in_N=0, out_N=0, - out_shape=(ModelConstants.IDX_N, ModelConstants.PLAN_WIDTH)) - else: - self.parse_mdn('plan', outs, in_N=ModelConstants.PLAN_MHP_N, out_N=ModelConstants.PLAN_MHP_SELECTION, - out_shape=(ModelConstants.IDX_N, ModelConstants.PLAN_WIDTH)) + plan_mhp = outs['plan'].shape[1] != 2 * ModelConstants.IDX_N * ModelConstants.PLAN_WIDTH + plan_in_N, plan_out_N = (ModelConstants.PLAN_MHP_N, ModelConstants.PLAN_MHP_SELECTION) if plan_mhp else (0, 0) + self.parse_mdn('plan', outs, in_N=plan_in_N, out_N=plan_out_N, + out_shape=(ModelConstants.IDX_N, ModelConstants.PLAN_WIDTH)) + if 'planplus' in outs: + self.parse_mdn('planplus', outs, in_N=plan_in_N, out_N=plan_out_N, out_shape=(ModelConstants.IDX_N, ModelConstants.PLAN_WIDTH)) if 'lane_lines' in outs: self.parse_mdn('lane_lines', outs, in_N=0, out_N=0, out_shape=(ModelConstants.NUM_LANE_LINES, ModelConstants.IDX_N, ModelConstants.LANE_LINES_WIDTH)) diff --git a/frogpilot/tinygrad_modeld/tinygrad_modeld.py b/frogpilot/tinygrad_modeld/tinygrad_modeld.py index 6e62b5587..76b6f3b54 100755 --- a/frogpilot/tinygrad_modeld/tinygrad_modeld.py +++ b/frogpilot/tinygrad_modeld/tinygrad_modeld.py @@ -33,6 +33,7 @@ from openpilot.frogpilot.common.frogpilot_variables import get_frogpilot_toggles PROCESS_NAME = "frogpilot.tinygrad_modeld.tinygrad_modeld" SEND_RAW_PRED = os.getenv('SEND_RAW_PRED') +RECOVERY_POWER = 1.0 # The higher this number the more aggressively the model will recover to lanecenter, too high and it will ping-pong LAT_SMOOTH_SECONDS = 0.1 @@ -43,6 +44,8 @@ MIN_LAT_CONTROL_SPEED = 0.3 def get_action_from_model(model_output: dict[str, np.ndarray], prev_action: log.ModelDataV2.Action, lat_action_t: float, long_action_t: float, v_ego: float, mlsim: bool, is_v9: bool) -> log.ModelDataV2.Action: plan = model_output['plan'][0] + if 'planplus' in model_output: + plan = plan + RECOVERY_POWER*model_output['planplus'][0] desired_accel, should_stop = get_accel_from_plan_tomb_raider(plan[:,Plan.VELOCITY][:,0], plan[:,Plan.ACCELERATION][:,0], ModelConstants.T_IDXS, @@ -56,7 +59,7 @@ def get_action_from_model(model_output: dict[str, np.ndarray], prev_action: log. else: desired_curvature = prev_action.desiredCurvature else: - desired_curvature = get_curvature_from_output(model_output, v_ego, lat_action_t, mlsim=mlsim) + desired_curvature = get_curvature_from_output(model_output, plan, v_ego, lat_action_t, mlsim=mlsim) if v_ego > MIN_LAT_CONTROL_SPEED: desired_curvature = smooth_value(desired_curvature, prev_action.desiredCurvature, LAT_SMOOTH_SECONDS) else: