wehub-resource-sync
88 lines
2.3 KiB
Python
88 lines
2.3 KiB
Python
# Vendored from SkillRL (Apache-2.0 License)
|
|
# Original: agent_system/memory/base.py + agent_system/memory/memory.py
|
|
# Merged into a single file for simplicity.
|
|
|
|
from abc import ABC, abstractmethod
|
|
from typing import List, Dict, Any, Tuple
|
|
|
|
|
|
class BaseMemory(ABC):
|
|
"""Base class for memory management."""
|
|
|
|
@abstractmethod
|
|
def __len__(self):
|
|
pass
|
|
|
|
@abstractmethod
|
|
def __getitem__(self, idx: int):
|
|
pass
|
|
|
|
@abstractmethod
|
|
def reset(self, batch_size: int):
|
|
pass
|
|
|
|
@abstractmethod
|
|
def store(self, record: Dict[str, List[Any]]):
|
|
pass
|
|
|
|
@abstractmethod
|
|
def fetch(self, step: int):
|
|
pass
|
|
|
|
|
|
class SimpleMemory(BaseMemory):
|
|
"""Per-environment history buffer for storing observations and actions."""
|
|
|
|
def __init__(self):
|
|
self._data = None
|
|
self.keys = None
|
|
self.batch_size = 0
|
|
|
|
def __len__(self):
|
|
return len(self._data)
|
|
|
|
def __getitem__(self, idx):
|
|
return self._data[idx]
|
|
|
|
def reset(self, batch_size: int):
|
|
if self._data is not None:
|
|
self._data.clear()
|
|
self._data = [[] for _ in range(batch_size)]
|
|
self.batch_size = batch_size
|
|
self.keys = None
|
|
|
|
def store(self, record: Dict[str, List[Any]]):
|
|
if self.keys is None:
|
|
self.keys = list(record.keys())
|
|
assert self.keys == list(record.keys())
|
|
|
|
for env_idx in range(self.batch_size):
|
|
self._data[env_idx].append({k: record[k][env_idx] for k in self.keys})
|
|
|
|
def fetch(
|
|
self,
|
|
history_length: int,
|
|
obs_key: str = "text_obs",
|
|
action_key: str = "action",
|
|
) -> Tuple[List[str], List[int]]:
|
|
memory_contexts, valid_lengths = [], []
|
|
|
|
for env_idx in range(self.batch_size):
|
|
recent = self._data[env_idx][-history_length:]
|
|
valid_len = len(recent)
|
|
start_idx = len(self._data[env_idx]) - valid_len
|
|
|
|
lines = []
|
|
for j, rec in enumerate(recent):
|
|
step_num = start_idx + j + 1
|
|
act = rec[action_key]
|
|
obs = rec[obs_key]
|
|
lines.append(
|
|
f"[Observation {step_num}: '{obs}', Action {step_num}: '{act}']"
|
|
)
|
|
|
|
memory_contexts.append("\n".join(lines))
|
|
valid_lengths.append(valid_len)
|
|
|
|
return memory_contexts, valid_lengths
|