vlalab 0.1.0__py3-none-any.whl

vlalab/logging/jsonl_writer.py

@@ -0,0 +1,114 @@
+"""
+VLA-Lab JSONL Writer
+
+Efficient append-only writer for step records.
+"""
+
+import json
+from pathlib import Path
+from typing import Union, Dict, Any
+import threading
+
+
+class JsonlWriter:
+    """Thread-safe JSONL file writer for step records."""
+    
+    def __init__(self, file_path: Union[str, Path]):
+        """
+        Initialize JSONL writer.
+        
+        Args:
+            file_path: Path to the JSONL file
+        """
+        self.file_path = Path(file_path)
+        self.file_path.parent.mkdir(parents=True, exist_ok=True)
+        self._lock = threading.Lock()
+        self._file = None
+        self._line_count = 0
+        
+        # Open file in append mode
+        self._file = open(self.file_path, "a", encoding="utf-8")
+    
+    def write(self, record: Union[Dict[str, Any], "StepRecord"]) -> int:
+        """
+        Write a record to the JSONL file.
+        
+        Args:
+            record: Dictionary or StepRecord to write
+            
+        Returns:
+            Line number of the written record
+        """
+        if hasattr(record, "to_dict"):
+            record = record.to_dict()
+        
+        json_str = json.dumps(record, ensure_ascii=False)
+        
+        with self._lock:
+            self._file.write(json_str + "\n")
+            self._file.flush()
+            self._line_count += 1
+            return self._line_count - 1
+    
+    def close(self):
+        """Close the file."""
+        if self._file is not None:
+            self._file.close()
+            self._file = None
+    
+    @property
+    def line_count(self) -> int:
+        """Get the number of lines written."""
+        return self._line_count
+    
+    def __enter__(self):
+        return self
+    
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.close()
+        return False
+
+
+class JsonlReader:
+    """Reader for JSONL files."""
+    
+    def __init__(self, file_path: Union[str, Path]):
+        """
+        Initialize JSONL reader.
+        
+        Args:
+            file_path: Path to the JSONL file
+        """
+        self.file_path = Path(file_path)
+        
+        if not self.file_path.exists():
+            raise FileNotFoundError(f"JSONL file not found: {self.file_path}")
+    
+    def __iter__(self):
+        """Iterate over records in the file."""
+        with open(self.file_path, "r", encoding="utf-8") as f:
+            for line in f:
+                line = line.strip()
+                if line:
+                    yield json.loads(line)
+    
+    def read_all(self):
+        """Read all records into a list."""
+        return list(self)
+    
+    def count(self) -> int:
+        """Count the number of records."""
+        count = 0
+        with open(self.file_path, "r", encoding="utf-8") as f:
+            for line in f:
+                if line.strip():
+                    count += 1
+        return count
+    
+    def read_line(self, line_idx: int) -> Dict[str, Any]:
+        """Read a specific line by index."""
+        with open(self.file_path, "r", encoding="utf-8") as f:
+            for i, line in enumerate(f):
+                if i == line_idx:
+                    return json.loads(line.strip())
+        raise IndexError(f"Line {line_idx} not found")

vlalab/logging/run_loader.py

@@ -0,0 +1,216 @@
+"""
+VLA-Lab Run Loader
+
+Load and access run data.
+"""
+
+import json
+from pathlib import Path
+from typing import Dict, Any, List, Optional, Iterator
+import numpy as np
+
+from vlalab.schema.run import RunMeta
+from vlalab.schema.step import StepRecord
+from vlalab.logging.jsonl_writer import JsonlReader
+
+
+def load_run_info(run_dir: Path) -> Dict[str, Any]:
+    """
+    Load basic run information.
+    
+    Args:
+        run_dir: Path to run directory
+        
+    Returns:
+        Dictionary with run info
+    """
+    run_dir = Path(run_dir)
+    meta_path = run_dir / "meta.json"
+    steps_path = run_dir / "steps.jsonl"
+    
+    info = {
+        "run_dir": str(run_dir),
+        "run_name": run_dir.name,
+    }
+    
+    if meta_path.exists():
+        with open(meta_path, "r") as f:
+            meta = json.load(f)
+            info.update({
+                "model_name": meta.get("model_name", "unknown"),
+                "task_name": meta.get("task_name", "unknown"),
+                "robot_name": meta.get("robot_name", "unknown"),
+                "start_time": meta.get("start_time", "unknown"),
+                "total_steps": meta.get("total_steps", 0),
+            })
+    
+    if steps_path.exists():
+        reader = JsonlReader(steps_path)
+        info["actual_steps"] = reader.count()
+    
+    return info
+
+
+class RunLoader:
+    """
+    Load and access run data.
+    
+    Usage:
+        loader = RunLoader("runs/my_run")
+        
+        # Access metadata
+        print(loader.meta.model_name)
+        
+        # Iterate over steps
+        for step in loader.iter_steps():
+            print(step.step_idx, step.action.values)
+        
+        # Access specific step
+        step = loader.get_step(10)
+        
+        # Load image for a step
+        image = loader.load_image(step, camera_name="front")
+    """
+    
+    def __init__(self, run_dir: Path):
+        """
+        Initialize run loader.
+        
+        Args:
+            run_dir: Path to run directory
+        """
+        self.run_dir = Path(run_dir)
+        
+        if not self.run_dir.exists():
+            raise FileNotFoundError(f"Run directory not found: {self.run_dir}")
+        
+        # Load metadata
+        meta_path = self.run_dir / "meta.json"
+        if meta_path.exists():
+            self.meta = RunMeta.load(str(meta_path))
+        else:
+            # Create minimal metadata
+            self.meta = RunMeta(
+                run_name=self.run_dir.name,
+                start_time="unknown",
+            )
+        
+        # Initialize step reader
+        self._steps_path = self.run_dir / "steps.jsonl"
+        self._steps_cache: Optional[List[StepRecord]] = None
+    
+    @property
+    def step_count(self) -> int:
+        """Get the number of steps."""
+        if self._steps_path.exists():
+            return JsonlReader(self._steps_path).count()
+        return 0
+    
+    def iter_steps(self) -> Iterator[StepRecord]:
+        """Iterate over all steps."""
+        if not self._steps_path.exists():
+            return
+        
+        reader = JsonlReader(self._steps_path)
+        for data in reader:
+            yield StepRecord.from_dict(data)
+    
+    def get_steps(self) -> List[StepRecord]:
+        """Get all steps as a list (cached)."""
+        if self._steps_cache is None:
+            self._steps_cache = list(self.iter_steps())
+        return self._steps_cache
+    
+    def get_step(self, step_idx: int) -> Optional[StepRecord]:
+        """
+        Get a specific step by index.
+        
+        Note: This loads steps sequentially, so it's O(n).
+        For random access, use get_steps() first.
+        """
+        for step in self.iter_steps():
+            if step.step_idx == step_idx:
+                return step
+        return None
+    
+    def load_image(
+        self,
+        step: StepRecord,
+        camera_name: Optional[str] = None,
+    ) -> Optional[np.ndarray]:
+        """
+        Load image for a step.
+        
+        Args:
+            step: StepRecord
+            camera_name: Camera name (None = first camera)
+            
+        Returns:
+            Image array (RGB) or None if not found
+        """
+        import cv2
+        
+        if not step.obs.images:
+            return None
+        
+        # Find image ref
+        image_ref = None
+        for ref in step.obs.images:
+            if camera_name is None or ref.camera_name == camera_name:
+                image_ref = ref
+                break
+        
+        if image_ref is None:
+            return None
+        
+        # Load image
+        image_path = self.run_dir / image_ref.path
+        if not image_path.exists():
+            return None
+        
+        image = cv2.imread(str(image_path))
+        if image is None:
+            return None
+        
+        return cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+    
+    def get_all_states(self) -> np.ndarray:
+        """Get all states as a numpy array."""
+        steps = self.get_steps()
+        if not steps:
+            return np.array([])
+        
+        states = [step.obs.state for step in steps if step.obs.state]
+        if not states:
+            return np.array([])
+        
+        return np.array(states)
+    
+    def get_all_actions(self) -> List[np.ndarray]:
+        """Get all actions as a list of arrays."""
+        steps = self.get_steps()
+        actions = []
+        for step in steps:
+            if step.action.values:
+                actions.append(np.array(step.action.values))
+        return actions
+    
+    def get_timing_series(self) -> Dict[str, np.ndarray]:
+        """Get timing data as arrays."""
+        steps = self.get_steps()
+        
+        timing_keys = [
+            "transport_latency_ms",
+            "inference_latency_ms",
+            "total_latency_ms",
+            "message_interval_ms",
+        ]
+        
+        result = {key: [] for key in timing_keys}
+        
+        for step in steps:
+            for key in timing_keys:
+                value = getattr(step.timing, key, None)
+                result[key].append(value if value is not None else np.nan)
+        
+        return {key: np.array(values) for key, values in result.items()}

vlalab/logging/run_logger.py

@@ -0,0 +1,343 @@
+"""
+VLA-Lab Run Logger
+
+Unified logging interface for VLA real-world deployment.
+"""
+
+import os
+import time
+import base64
+from pathlib import Path
+from datetime import datetime
+from typing import Optional, List, Dict, Any, Union
+import numpy as np
+
+from vlalab.schema.step import StepRecord, ObsData, ActionData, TimingData, ImageRef
+from vlalab.schema.run import RunMeta, CameraConfig
+from vlalab.logging.jsonl_writer import JsonlWriter
+
+
+class RunLogger:
+    """
+    Unified logger for VLA deployment runs.
+    
+    Creates a run directory with:
+    - meta.json: Run metadata
+    - steps.jsonl: Step records (one per line)
+    - artifacts/images/: Image files
+    
+    Usage:
+        logger = RunLogger(
+            run_dir="runs/my_run",
+            model_name="diffusion_policy",
+            task_name="pick_and_place",
+        )
+        
+        # Log a step
+        logger.log_step(
+            step_idx=0,
+            state=[0.5, 0.2, 0.3, 0, 0, 0, 1, 1.0],
+            action=[[0.51, 0.21, 0.31, 0, 0, 0, 1, 1.0]],
+            images={"front": image_array},
+            timing={...},
+        )
+        
+        # Close when done
+        logger.close()
+    """
+    
+    def __init__(
+        self,
+        run_dir: Union[str, Path],
+        model_name: str = "unknown",
+        model_path: Optional[str] = None,
+        model_type: Optional[str] = None,
+        task_name: str = "unknown",
+        task_prompt: Optional[str] = None,
+        robot_name: str = "unknown",
+        robot_type: Optional[str] = None,
+        cameras: Optional[List[Dict[str, Any]]] = None,
+        inference_freq: Optional[float] = None,
+        action_dim: Optional[int] = None,
+        action_horizon: Optional[int] = None,
+        server_config: Optional[Dict[str, Any]] = None,
+        client_config: Optional[Dict[str, Any]] = None,
+        auto_create: bool = True,
+        image_quality: int = 85,
+    ):
+        """
+        Initialize the run logger.
+        
+        Args:
+            run_dir: Directory to store run data
+            model_name: Name of the model
+            model_path: Path to the model checkpoint
+            model_type: Type of model (diffusion_policy, groot, etc.)
+            task_name: Name of the task
+            task_prompt: Language prompt for the task
+            robot_name: Name of the robot
+            robot_type: Type of robot (franka, ur5, etc.)
+            cameras: List of camera configurations
+            inference_freq: Inference frequency in Hz
+            action_dim: Action dimension
+            action_horizon: Action horizon (chunk size)
+            server_config: Server configuration dict
+            client_config: Client configuration dict
+            auto_create: Whether to create the run directory
+            image_quality: JPEG quality for saved images (1-100)
+        """
+        self.run_dir = Path(run_dir)
+        self.image_quality = image_quality
+        self._step_count = 0
+        self._closed = False
+        
+        if auto_create:
+            self.run_dir.mkdir(parents=True, exist_ok=True)
+            (self.run_dir / "artifacts" / "images").mkdir(parents=True, exist_ok=True)
+        
+        # Create run name from directory name
+        run_name = self.run_dir.name
+        start_time = datetime.now().isoformat()
+        
+        # Build camera configs
+        camera_configs = []
+        if cameras:
+            for cam in cameras:
+                if isinstance(cam, dict):
+                    camera_configs.append(CameraConfig(**cam))
+                else:
+                    camera_configs.append(cam)
+        
+        # Create metadata
+        self.meta = RunMeta(
+            run_name=run_name,
+            start_time=start_time,
+            model_name=model_name,
+            model_path=model_path,
+            model_type=model_type,
+            task_name=task_name,
+            task_prompt=task_prompt,
+            robot_name=robot_name,
+            robot_type=robot_type,
+            cameras=camera_configs,
+            inference_freq=inference_freq,
+            action_dim=action_dim,
+            action_horizon=action_horizon,
+            server_config=server_config or {},
+            client_config=client_config or {},
+        )
+        
+        # Save initial metadata
+        self._save_meta()
+        
+        # Initialize JSONL writer
+        self._jsonl_writer = JsonlWriter(self.run_dir / "steps.jsonl")
+    
+    def _save_meta(self):
+        """Save metadata to meta.json."""
+        self.meta.save(str(self.run_dir / "meta.json"))
+    
+    def _save_image(
+        self,
+        image: np.ndarray,
+        step_idx: int,
+        camera_name: str = "default",
+    ) -> ImageRef:
+        """
+        Save an image to the artifacts directory.
+        
+        Args:
+            image: Image array (H, W, C) in RGB or BGR format
+            step_idx: Step index
+            camera_name: Camera name
+            
+        Returns:
+            ImageRef with path information
+        """
+        # Import cv2 here to avoid hard dependency
+        import cv2
+        
+        # Create filename
+        filename = f"step_{step_idx:06d}_{camera_name}.jpg"
+        rel_path = f"artifacts/images/{filename}"
+        abs_path = self.run_dir / rel_path
+        
+        # Ensure image is uint8
+        if image.dtype != np.uint8:
+            image = (image * 255).astype(np.uint8)
+        
+        # Convert RGB to BGR for cv2
+        if len(image.shape) == 3 and image.shape[2] == 3:
+            image_bgr = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
+        else:
+            image_bgr = image
+        
+        # Save image
+        cv2.imwrite(str(abs_path), image_bgr, [cv2.IMWRITE_JPEG_QUALITY, self.image_quality])
+        
+        return ImageRef(
+            path=rel_path,
+            camera_name=camera_name,
+            shape=list(image.shape),
+            encoding="jpeg",
+        )
+    
+    def _decode_base64_image(self, b64_str: str) -> Optional[np.ndarray]:
+        """Decode a base64-encoded image."""
+        try:
+            import cv2
+            img_data = base64.b64decode(b64_str)
+            img_array = np.frombuffer(img_data, dtype=np.uint8)
+            image = cv2.imdecode(img_array, cv2.IMREAD_COLOR)
+            return cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+        except Exception:
+            return None
+    
+    def log_step(
+        self,
+        step_idx: int,
+        state: Optional[List[float]] = None,
+        pose: Optional[List[float]] = None,
+        gripper: Optional[float] = None,
+        action: Optional[Union[List[float], List[List[float]]]] = None,
+        images: Optional[Dict[str, Union[np.ndarray, str]]] = None,
+        timing: Optional[Dict[str, Any]] = None,
+        prompt: Optional[str] = None,
+        tags: Optional[Dict[str, Any]] = None,
+    ):
+        """
+        Log a single step.
+        
+        Args:
+            step_idx: Step index
+            state: Low-dim state vector
+            pose: Robot pose [x, y, z, qx, qy, qz, qw] or similar
+            gripper: Gripper state (0-1)
+            action: Action values (single or chunk)
+            images: Dict of camera_name -> image_array or base64_string
+            timing: Timing information dict
+            prompt: Language prompt
+            tags: Additional tags/metadata
+        """
+        if self._closed:
+            raise RuntimeError("Logger is closed")
+        
+        # Process images
+        image_refs = []
+        if images:
+            for camera_name, image_data in images.items():
+                # Handle base64 string
+                if isinstance(image_data, str):
+                    image_array = self._decode_base64_image(image_data)
+                    if image_array is None:
+                        continue
+                else:
+                    image_array = image_data
+                
+                # Save image and get reference
+                ref = self._save_image(image_array, step_idx, camera_name)
+                image_refs.append(ref)
+        
+        # Build observation
+        obs = ObsData(
+            state=state or [],
+            images=image_refs,
+            pose=pose,
+            gripper=gripper,
+        )
+        
+        # Build action
+        action_data = ActionData()
+        if action is not None:
+            # Normalize to list of lists (chunk format)
+            if action and not isinstance(action[0], (list, tuple)):
+                action = [action]  # Single action -> chunk of 1
+            action_data.values = [list(a) for a in action]
+            if action:
+                action_data.action_dim = len(action[0])
+                action_data.chunk_size = len(action)
+        
+        # Build timing
+        timing_data = TimingData()
+        if timing:
+            timing_data = TimingData.from_dict(timing)
+            timing_data.compute_latencies()
+        
+        # Create step record
+        record = StepRecord(
+            step_idx=step_idx,
+            obs=obs,
+            action=action_data,
+            timing=timing_data,
+            prompt=prompt,
+            tags=tags or {},
+        )
+        
+        # Write to JSONL
+        self._jsonl_writer.write(record)
+        self._step_count += 1
+    
+    def log_step_raw(
+        self,
+        step_idx: int,
+        obs_dict: Dict[str, Any],
+        action_dict: Dict[str, Any],
+        timing_dict: Dict[str, Any],
+        **kwargs,
+    ):
+        """
+        Log a step from raw dictionaries (for adapters).
+        
+        Args:
+            step_idx: Step index
+            obs_dict: Observation dictionary
+            action_dict: Action dictionary
+            timing_dict: Timing dictionary
+            **kwargs: Additional fields
+        """
+        record = StepRecord(
+            step_idx=step_idx,
+            obs=ObsData.from_dict(obs_dict),
+            action=ActionData.from_dict(action_dict),
+            timing=TimingData.from_dict(timing_dict),
+            **kwargs,
+        )
+        self._jsonl_writer.write(record)
+        self._step_count += 1
+    
+    def update_meta(self, **kwargs):
+        """Update metadata fields."""
+        for key, value in kwargs.items():
+            if hasattr(self.meta, key):
+                setattr(self.meta, key, value)
+        self._save_meta()
+    
+    def close(self):
+        """Close the logger and finalize metadata."""
+        if self._closed:
+            return
+        
+        # Update final metadata
+        self.meta.end_time = datetime.now().isoformat()
+        self.meta.total_steps = self._step_count
+        
+        # Calculate duration
+        start = datetime.fromisoformat(self.meta.start_time)
+        end = datetime.fromisoformat(self.meta.end_time)
+        self.meta.total_duration_s = (end - start).total_seconds()
+        
+        self._save_meta()
+        self._jsonl_writer.close()
+        self._closed = True
+    
+    @property
+    def step_count(self) -> int:
+        """Get the number of logged steps."""
+        return self._step_count
+    
+    def __enter__(self):
+        return self
+    
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.close()
+        return False

vlalab/schema/__init__.py

@@ -0,0 +1,17 @@
+"""
+VLA-Lab Schema Module
+
+Defines standardized data structures for VLA deployment logging.
+"""
+
+from vlalab.schema.step import StepRecord, ObsData, ActionData, TimingData, ImageRef
+from vlalab.schema.run import RunMeta
+
+__all__ = [
+    "StepRecord",
+    "ObsData",
+    "ActionData", 
+    "TimingData",
+    "ImageRef",
+    "RunMeta",
+]