Fix NaN state input on gr.Examples click; use gr.JSON output with rounded values

app.py

@@ -216,13 +216,22 @@ def build_batch(images, prompt, state):
     return batch
 
 
+# Default robot state used whenever a caller doesn't provide one (e.g. the
+# gr.Examples rows below only populate camera_1 + instruction). This mirrors
+# the default values of the Robot State sliders in the UI, so results are
+# consistent regardless of entry point (example click, manual button click,
+# or a direct API call that omits the state args).
+DEFAULT_STATE_J = 0.0
+DEFAULT_GRIPPER = 0.04
+
+
 @spaces.GPU(duration=120)
 def predict_actions(
     camera_1: Image.Image,
     instruction: str,
-    state_j1: float, state_j2: float, state_j3: float,
-    state_j4: float, state_j5: float, state_j6: float,
-    state_j7: float, gripper: float,
+    state_j1: float = DEFAULT_STATE_J, state_j2: float = DEFAULT_STATE_J, state_j3: float = DEFAULT_STATE_J,
+    state_j4: float = DEFAULT_STATE_J, state_j5: float = DEFAULT_STATE_J, state_j6: float = DEFAULT_STATE_J,
+    state_j7: float = DEFAULT_STATE_J, gripper: float = DEFAULT_GRIPPER,
 ):
     """Predict a robot action chunk from laboratory camera views and a language instruction.
 
@@ -233,15 +242,15 @@ def predict_actions(
     Args:
         camera_1: Camera view of the laboratory workspace.
         instruction: Natural language task instruction (e.g. "Pick up the beaker").
-        state_j1..j7: Franka Panda 7-DOF arm joint angles (radians).
-        gripper: Gripper width in meters (0.0 = closed, 0.04 = fully open).
+        state_j1..j7: Franka Panda 7-DOF arm joint angles (radians). Defaults to 0.0.
+        gripper: Gripper width in meters (0.0 = closed, 0.04 = fully open). Defaults to 0.04.
 
     Returns:
         A matplotlib figure visualizing the predicted action trajectory,
-        and a JSON string with the raw action values.
+        and a JSON dict with the raw action values.
     """
     if not instruction or not instruction.strip():
-        return None, '{"error": "Please provide a task instruction."}'
+        return None, {"error": "Please provide a task instruction."}
 
     # Convert PIL image to numpy
     img1 = np.array(camera_1.convert("RGB"))
@@ -251,6 +260,20 @@ def predict_actions(
     # with only 1, slots 2/3 are masked out)
     images = [img1, None, None]
 
+    # Guard against any caller (or gr.Examples cache) passing an empty/missing
+    # value for a state component. NaN must never reach the model.
+    def _clean(value, default):
+        return default if value is None else value
+
+    state_j1 = _clean(state_j1, DEFAULT_STATE_J)
+    state_j2 = _clean(state_j2, DEFAULT_STATE_J)
+    state_j3 = _clean(state_j3, DEFAULT_STATE_J)
+    state_j4 = _clean(state_j4, DEFAULT_STATE_J)
+    state_j5 = _clean(state_j5, DEFAULT_STATE_J)
+    state_j6 = _clean(state_j6, DEFAULT_STATE_J)
+    state_j7 = _clean(state_j7, DEFAULT_STATE_J)
+    gripper = _clean(gripper, DEFAULT_GRIPPER)
+
     # Build state vector
     state = np.array([state_j1, state_j2, state_j3, state_j4,
                        state_j5, state_j6, state_j7, gripper], dtype=np.float32)
@@ -322,20 +345,23 @@ def predict_actions(
     fig.savefig(fig_path, dpi=150, bbox_inches="tight")
     plt.close(fig)
 
-    # Prepare JSON output
+    # Prepare JSON output (rounded for readability; gr.JSON renders the dict natively)
+    def _round_list(values, ndigits=4):
+        return [round(float(v), ndigits) for v in values]
+
     actions_list = actions.tolist()
     result = {
         "instruction": instruction,
-        "state_input": state.tolist(),
+        "state_input": _round_list(state.tolist()),
         "action_chunk_shape": list(actions.shape),
         "num_steps": int(actions.shape[0]),
         "action_dim": int(actions.shape[1]),
         "inference_time_s": round(infer_time, 3),
-        "first_action": actions_list[0] if len(actions_list) > 0 else None,
-        "last_action": actions_list[-1] if len(actions_list) > 0 else None,
+        "first_action": _round_list(actions_list[0]) if len(actions_list) > 0 else None,
+        "last_action": _round_list(actions_list[-1]) if len(actions_list) > 0 else None,
     }
 
-    return fig_path, json.dumps(result, indent=2)
+    return fig_path, result
 
 
 # ---- Gradio UI ----
@@ -390,7 +416,7 @@ with gr.Blocks(theme=gr.themes.Citrus(), css=CSS) as demo:
 
             with gr.Column(scale=1):
                 output_plot = gr.Image(label="Predicted Action Trajectory (50 steps)")
-                output_json = gr.Textbox(label="Action Details (JSON)", lines=12, max_lines=20)
+                output_json = gr.JSON(label="Action Details (JSON)")
 
     gr.Examples(
         examples=[