Upload model card for best_model.pt

README.md

@@ -0,0 +1,300 @@
+---
+license: mit
+library_name: pytorch
+tags:
+- robotics
+- progress-estimation
+- behavior-cloning
+---
+
+# SARM Progress Prediction
+
+Stage-aware progress prediction model for robot manipulation tasks
+
+## Model Description
+
+SARM predicts:
+- **Progress**: How far through the task (0.0 to 1.0)
+- **Stage**: Which stage of the task is being executed
+
+The model uses a transformer architecture to process sequences of RGB images and robot states.
+
+**Task**: clearing_food_from_table_into_fridge
+**Dataset**: IliaLarchenko/behavior_224_rgb
+
+## Model Details
+
+### Architecture
+- **Type**: Transformer with dual prediction heads (stage classification + progress regression)
+- **Model dimension**: 768
+- **Attention heads**: 12
+- **Transformer layers**: 8
+- **MLP dimension**: 512
+- **Number of stages**: 100
+- **Number of tasks**: 50
+
+### Training Details
+- **Checkpoint**: `best_model.pt`
+- **Training step**: 4800
+- **Epoch**: unknown
+- **Training loss**: unknown
+- **Validation loss**: 1.0865614609792829
+- **Batch size**: 16
+- **Learning rate**: 0.0001
+- **Max sequence length**: 13
+
+## Usage
+
+### Download and Load Model
+
+```python
+from hf_model_hub import download_model_from_hub
+from model import SARM
+import torch
+import json
+
+# Download model and config
+files = download_model_from_hub(
+    repo_id="YOUR_USERNAME/YOUR_REPO",
+    checkpoint_name="best_model.pt",
+    output_dir="./downloaded_model"
+)
+
+# Load config
+with open(files["config"], "r") as f:
+    config = json.load(f)
+
+# Create model
+model_config = config["model"]
+model = SARM(
+    d_model=model_config["d_model"],
+    n_heads=model_config["n_heads"],
+    n_layers=model_config["n_layers"],
+    d_mlp=model_config["d_mlp"],
+    num_stages=model_config["num_stages"],
+    d_state=model_config["d_state"],
+    num_tasks=model_config["num_tasks"],
+)
+
+# Load checkpoint
+checkpoint = torch.load(files["checkpoint"])
+model.load_state_dict(checkpoint["model_state_dict"])
+model.eval()
+```
+
+### Run Inference
+
+```python
+# Assuming you have images and states prepared
+with torch.no_grad():
+    stage_logits, progress = model(images, states, tasks, padding_mask)
+    
+    # Get predictions for the last frame
+    predicted_stage = stage_logits[:, -1].argmax(dim=-1)
+    predicted_progress = progress[:, -1]
+```
+
+## Training Data
+
+This model was trained on the **IliaLarchenko/behavior_224_rgb** for robot manipulation tasks.
+
+Training episodes: 90 episodes
+Validation episodes: 15 episodes
+
+## Intended Use
+
+- Progress estimation for robot manipulation tasks
+- Stage classification for multi-step tasks  
+- Adaptive window sampling for VLA training
+- Task monitoring and intervention detection
+
+## Limitations
+
+- Trained on specific tasks from BEHAVIOR dataset
+- Requires RGB images (224x224) and robot state information
+- Fixed sequence length input
+
+## Citation
+
+If you use this model, please cite:
+
+```bibtex
+@misc{sarm-model,
+  author = {Your Name},
+  title = {SARM Progress Prediction},
+  year = {2025},
+  publisher = {HuggingFace},
+  url = {https://huggingface.co/YOUR_USERNAME/YOUR_REPO}
+}
+```
+
+## Training Configuration
+
+<details>
+<summary>Click to expand full training configuration</summary>
+
+```json
+{
+  "metadata": {
+    "model_name": "SARM Progress Prediction",
+    "description": "Stage-aware progress prediction model for robot manipulation tasks",
+    "task": "clearing_food_from_table_into_fridge",
+    "task_number": 25,
+    "dataset": "IliaLarchenko/behavior_224_rgb",
+    "version": "1.0",
+    "author": "Your Name",
+    "tags": [
+      "robotics",
+      "progress-estimation",
+      "behavior-cloning"
+    ]
+  },
+  "model": {
+    "d_model": 768,
+    "n_heads": 12,
+    "n_layers": 8,
+    "d_mlp": 512,
+    "num_stages": 100,
+    "d_state": 256,
+    "num_tasks": 50
+  },
+  "training": {
+    "max_steps": 10000,
+    "learning_rate": 0.0001,
+    "weight_decay": 0.0001,
+    "batch_size": 16,
+    "gradient_accumulation_steps": 4,
+    "max_grad_norm": 1.0,
+    "scheduler": "cosine",
+    "stage_loss_weight": 1.0,
+    "progress_loss_weight": 1.0,
+    "validation_steps": 100,
+    "save_steps": 200
+  },
+  "data": {
+    "max_sequence_length": 13,
+    "image_size": 224,
+    "num_workers": 10,
+    "val_workers": 10,
+    "val_samples": 500,
+    "train_episodes": [
+      1,
+      2,
+      3,
+      4,
+      5,
+      6,
+      7,
+      8,
+      9,
+      10,
+      11,
+      12,
+      13,
+      14,
+      15,
+      16,
+      17,
+      18,
+      19,
+      20,
+      21,
+      22,
+      23,
+      24,
+      25,
+      26,
+      27,
+      28,
+      29,
+      30,
+      31,
+      32,
+      33,
+      34,
+      35,
+      36,
+      37,
+      38,
+      39,
+      40,
+      41,
+      42,
+      43,
+      44,
+      45,
+      46,
+      47,
+      48,
+      49,
+      50,
+      51,
+      52,
+      53,
+      54,
+      55,
+      56,
+      57,
+      58,
+      59,
+      60,
+      61,
+      62,
+      63,
+      64,
+      65,
+      66,
+      67,
+      68,
+      69,
+      70,
+      71,
+      72,
+      73,
+      74,
+      75,
+      76,
+      77,
+      78,
+      79,
+      80,
+      81,
+      82,
+      83,
+      84,
+      85,
+      86,
+      87,
+      88,
+      89,
+      90
+    ],
+    "val_episodes": [
+      91,
+      92,
+      93,
+      94,
+      95,
+      96,
+      97,
+      98,
+      99,
+      100,
+      101,
+      102,
+      103,
+      104,
+      105
+    ],
+    "seed": 42
+  },
+  "logging": {
+    "project_name": "sarm-training",
+    "run_name": null,
+    "log_freq": 10,
+    "checkpoint_dir": "checkpoints_sarm_25_2"
+  }
+}
+```
+
+</details>