Memory manager: responsible for storing & fetching per‑environment history records.
| 2 | from .base import BaseMemory |
| 3 | |
| 4 | class SimpleMemory(BaseMemory): |
| 5 | """ |
| 6 | Memory manager: responsible for storing & fetching per‑environment history records. |
| 7 | """ |
| 8 | def __init__(self): |
| 9 | self._data = None |
| 10 | self.keys = None |
| 11 | self.batch_size = 0 |
| 12 | |
| 13 | def __len__(self): |
| 14 | return len(self._data) |
| 15 | |
| 16 | def __getitem__(self, idx): |
| 17 | return self._data[idx] |
| 18 | |
| 19 | def reset(self, batch_size: int): |
| 20 | if self._data is not None: |
| 21 | self._data.clear() |
| 22 | self._data = [[] for _ in range(batch_size)] |
| 23 | self.batch_size = batch_size |
| 24 | self.keys = None |
| 25 | |
| 26 | def store(self, record: Dict[str, List[Any]]): |
| 27 | """ |
| 28 | Store a new record (one step of history) for each environment instance. |
| 29 | |
| 30 | Args: |
| 31 | record (Dict[str, List[Any]]): |
| 32 | A dictionary where each key corresponds to a type of data |
| 33 | (e.g., 'text_obs', 'action'), and each value is a list of |
| 34 | length `batch_size`, containing the data for each environment. |
| 35 | """ |
| 36 | if self.keys is None: |
| 37 | self.keys = list(record.keys()) |
| 38 | assert self.keys == list(record.keys()) |
| 39 | |
| 40 | for env_idx in range(self.batch_size): |
| 41 | self._data[env_idx].append({k: record[k][env_idx] for k in self.keys}) |
| 42 | |
| 43 | def fetch( |
| 44 | self, |
| 45 | history_length: int, |
| 46 | obs_key: str = "text_obs", |
| 47 | action_key: str = "action", |
| 48 | ) -> Tuple[List[str], List[int]]: |
| 49 | """ |
| 50 | Fetch and format recent interaction history for each environment instance. |
| 51 | Args: |
| 52 | history_length (int): |
| 53 | Maximum number of past steps to retrieve per environment. |
| 54 | obs_key (str, default="text_obs"): |
| 55 | The key name used to access the observation in stored records. |
| 56 | For example: "text_obs" or "Observation", depending on the environment. |
| 57 | action_key (str, default="action"): |
| 58 | The key name used to access the action in stored records. |
| 59 | For example: "action" or "Action". |
| 60 | Returns: |
| 61 | memory_contexts : List[str] |