Update README.md

README.md

@@ -2,10 +2,12 @@
 datasets:
 - hellorobotinc/kitchen_cabinet_diagonal
 pipeline_tag: robotics
+tags:
+- Lerobot
 ---
 
 ## Model Card for Diffusion Policy / Kitchen cabinet opening - Diagonal configuration
-Diffusion Policy trained for [Stretch 3](https://hello-robot.com/stretch-3-product). Trained to open kitchen cabinet from diagonal orientation in Yi-Che (Jensen) Huang's home during his summer 2024 internship.
+Diffusion Policy trained for [Stretch 3](https://hello-robot.com/stretch-3-product) using [LeRobot](https://huggingface.co/lerobot). Trained to open kitchen cabinet from diagonal orientation in Yi-Che (Jensen) Huang's home during his Summer 2024 internship.
 ## Training Details
 Model trained on Dex-Teleop demonstrations collected with [stretch_ai](https://github.com/hello-robot/stretch_ai)
 
@@ -15,4 +17,15 @@ Model trained on Dex-Teleop demonstrations collected with [stretch_ai](https://g
 - Checkpoint: latest checkpoint
 - Training time: ~11 hours on RTX 3090
 
-Specific training configs can be found in config.json and config.yaml
+## Input/Output
+Inputs:
+
+- Head RGB (320x240)
+- Gripper RGB (240x320)
+- 9 dim joint state observations  (x, y, theta, lift, arm, roll, pitch, yaw, gripper)
+
+Output:
+- 9 dim joint state actions (x, y, theta, lift, arm, roll, pitch, yaw, gripper)
+- **optional extra dim for task progress (included in dataset, but not trained into this model)
+
+Specific training configs can be found in config.json and config.yaml