Spaces:
Sleeping
Sleeping
File size: 4,911 Bytes
27cdb3e | 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 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 | """
SQLAlchemy ORM models for the Replay Buffer.
Stores episodes and steps in SQLite for querying, analysis, and training data.
"""
from __future__ import annotations
import json
from datetime import datetime, timezone
from typing import Optional
from sqlalchemy import (
Boolean, Column, DateTime, Float, ForeignKey, Index,
Integer, String, Text, create_engine,
)
from sqlalchemy.orm import DeclarativeBase, Session, relationship, sessionmaker
class Base(DeclarativeBase):
"""SQLAlchemy declarative base for all models."""
pass
class EpisodeRecord(Base):
"""Stores a complete episode of agent-environment interaction.
Attributes:
id: Auto-increment primary key.
episode_id: UUID string for the episode.
scenario_id: Which scenario was attempted.
level: Difficulty level (1, 2, or 3).
total_reward: Sum of all step rewards.
solved: Whether the scenario was solved.
total_steps: Number of steps taken.
timestamp: When the episode occurred.
training_episode: Which training episode number this was.
"""
__tablename__ = "episodes"
id = Column(Integer, primary_key=True, autoincrement=True)
episode_id = Column(String(36), unique=True, nullable=False, index=True)
scenario_id = Column(String(100), nullable=False, index=True)
level = Column(Integer, nullable=False, index=True)
total_reward = Column(Float, nullable=False, default=0.0)
solved = Column(Boolean, nullable=False, default=False)
total_steps = Column(Integer, nullable=False, default=0)
timestamp = Column(DateTime, nullable=False, default=lambda: datetime.now(timezone.utc))
training_episode = Column(Integer, nullable=True)
# Relationship to steps
steps = relationship("StepRecord", back_populates="episode", cascade="all, delete-orphan",
order_by="StepRecord.step_number")
# Indexes for common queries
__table_args__ = (
Index("idx_scenario_solved", "scenario_id", "solved"),
Index("idx_level_reward", "level", "total_reward"),
Index("idx_timestamp", "timestamp"),
)
def to_dict(self) -> dict:
"""Convert to dictionary for JSON serialization."""
return {
"episode_id": self.episode_id,
"scenario_id": self.scenario_id,
"level": self.level,
"total_reward": self.total_reward,
"solved": self.solved,
"total_steps": self.total_steps,
"timestamp": self.timestamp.isoformat() if self.timestamp else None,
"training_episode": self.training_episode,
"steps": [s.to_dict() for s in self.steps] if self.steps else [],
}
class StepRecord(Base):
"""Stores a single step within an episode.
Attributes:
id: Auto-increment primary key.
episode_db_id: Foreign key to EpisodeRecord.
step_number: Step index within the episode.
observation_json: JSON-serialized observation dict.
action: The shell command issued.
result_json: JSON-serialized ExecutionResult.
reward: Reward for this step.
reward_breakdown_json: JSON-serialized reward breakdown.
error_type: Classified error type at this step.
"""
__tablename__ = "steps"
id = Column(Integer, primary_key=True, autoincrement=True)
episode_db_id = Column(Integer, ForeignKey("episodes.id"), nullable=False)
step_number = Column(Integer, nullable=False)
observation_json = Column(Text, nullable=False, default="{}")
action = Column(Text, nullable=False, default="")
result_json = Column(Text, nullable=False, default="{}")
reward = Column(Float, nullable=False, default=0.0)
reward_breakdown_json = Column(Text, nullable=False, default="{}")
error_type = Column(String(50), nullable=False, default="unknown")
episode = relationship("EpisodeRecord", back_populates="steps")
def to_dict(self) -> dict:
"""Convert to dictionary for JSON serialization."""
return {
"step": self.step_number,
"observation": json.loads(self.observation_json) if self.observation_json else {},
"action": self.action,
"result": json.loads(self.result_json) if self.result_json else {},
"reward": self.reward,
"reward_breakdown": json.loads(self.reward_breakdown_json) if self.reward_breakdown_json else {},
"error_type": self.error_type,
}
def create_database(db_url: str = "sqlite:///replay_buffer.db") -> sessionmaker:
"""Create the database and return a session factory.
Args:
db_url: SQLAlchemy database URL.
Returns:
Configured sessionmaker instance.
"""
engine = create_engine(db_url, echo=False)
Base.metadata.create_all(engine)
return sessionmaker(engine, expire_on_commit=False)
|