Spaces:
Sleeping
Sleeping
File size: 12,882 Bytes
7c117ed |
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 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 378 379 380 381 382 383 384 385 386 387 388 389 390 391 392 393 394 395 396 397 398 399 400 401 402 403 404 405 406 407 408 409 410 |
import random
import uuid
from dataclasses import dataclass, field
from typing import Dict, List, TypeVar
from abc import ABC, abstractmethod
from math import ceil
from aworld.core.common import ActionModel, Observation
from aworld.replay_buffer.query_filter import QueryCondition, QueryFilter
from aworld.logs.util import logger
T = TypeVar('T')
@dataclass
class Experience:
'''
Experience of agent.
'''
state: Observation
actions: List[ActionModel]
reward_t: float = None
adv_t: float = None
v_t: float = None
messages: List[Dict] = None
def to_dict(self):
return {
"state": self.state,
"actions": self.actions,
"reward_t": self.reward_t,
"adv_t": self.adv_t,
"v_t": self.v_t,
"messages": self.messages
}
@dataclass
class ExpMeta:
'''
Experience meta data.
'''
task_id: str
task_name: str
agent_id: str
step: int
execute_time: float
pre_agent: str
def to_dict(self):
return {
"task_id": self.task_id,
"task_name": self.task_name,
"agent_id": self.agent_id,
"step": self.step,
"execute_time": self.execute_time,
"pre_agent": self.pre_agent
}
@dataclass
class DataRow:
'''
Data row for storing data.
'''
exp_meta: ExpMeta
exp_data: Experience
id: str = field(default_factory=lambda: str(uuid.uuid4()))
def to_dict(self):
return {
"exp_meta": self.exp_meta.to_dict(),
"exp_data": self.exp_data.to_dict(),
"id": self.id
}
class Storage(ABC):
'''
Storage for storing and sampling data.
'''
@abstractmethod
def add(self, data: DataRow):
'''
Add data to the storage.
Args:
data (DataRow): Data to add.
'''
@abstractmethod
def add_batch(self, data_batch: List[DataRow]):
'''
Add batch of data to the storage.
Args:
data_batch (List[DataRow]): List of data to add.
'''
@abstractmethod
def size(self, query_condition: QueryCondition = None) -> int:
'''
Get the size of the storage.
Returns:
int: Size of the storage.
'''
@abstractmethod
def get_paginated(self, page: int, page_size: int, query_condition: QueryCondition = None) -> List[DataRow]:
'''
Get paginated data from the storage.
Args:
page (int): Page number.
page_size (int): Number of data per page.
Returns:
List[DataRow]: List of data.
'''
@abstractmethod
def get_all(self, query_condition: QueryCondition = None) -> List[DataRow]:
'''
Get all data from the storage.
Returns:
List[DataRow]: List of data.
'''
@abstractmethod
def get_by_task_id(self, task_id: str) -> List[DataRow]:
'''
Get data by task_id from the storage.
Args:
task_id (str): Task id.
Returns:
List[DataRow]: List of data.
'''
@abstractmethod
def get_bacth_by_task_ids(self, task_ids: List[str]) -> Dict[str, List[DataRow]]:
'''
Get batch of data by task_ids from the storage.
Args:
task_ids (List[str]): List of task ids.
Returns:
Dict[str, List[DataRow]]: Dictionary of data.
The key is the task_id and the value is the list of data.
The list of data is sorted by step.
'''
class Sampler(ABC):
'''
Sample data from the storage.
'''
def sample(self,
storage: Storage,
batch_size: int,
query_condition: QueryCondition = None) -> List[DataRow]:
'''
Sample data from the storage.
Args:
storage (Storage): Storage to sample from.
batch_size (int): Number of data to sample.
query_condition (QueryCondition, optional): Query condition. Defaults to None.
Returns:
List[DataRow]
'''
class TaskSampler(Sampler):
'''
Sample task data from storage, returns Dict[str, List[DataRow]] where:
- key is task_id
- value is list of task all data rows
'''
def sorted_by_step(self, task_experience: List[DataRow]) -> List[DataRow]:
'''
Sort the task experience by step and execute_time.
Args:
task_experience (List[DataRow]): List of task experience.
Returns:
List[DataRow]: List of task experience sorted by step and execute_time.
'''
return sorted(task_experience, key=lambda x: (x.exp_meta.step, x.exp_meta.execute_time))
def sample(self,
storage: Storage,
batch_size: int,
query_condition: QueryCondition = None) -> List[DataRow]:
task_ids = self.sample_task_ids(storage, batch_size, query_condition)
return storage.get_bacth_by_task_ids(task_ids)
def sample_tasks(self,
storage: Storage,
batch_size: int,
query_condition: QueryCondition = None) -> Dict[str, List[DataRow]]:
'''
Sample data from the storage.
Args:
storage (Storage): Storage to sample from.
batch_size (int): Number of data to sample.
query_condition (QueryCondition, optional): Query condition. Defaults to None.
Returns:
Dict[str, List[DataRow]]: Dictionary of sampled data.
The key is the task_id and the value is the list of data.
The list of data is sorted by step.
'''
task_ids = self.sample_task_ids(storage, batch_size, query_condition)
raws = storage.get_bacth_by_task_ids(task_ids)
return {task_id: self.sorted_by_step(raws) for task_id, raws in raws.items()}
@abstractmethod
def sample_task_ids(self,
storage: Storage,
batch_size: int,
query_condition: QueryCondition = None) -> List[str]:
'''
Sample task_ids from the storage.
Args:
storage (Storage): Storage to sample from.
batch_size (int): Number of task_ids to sample.
query_condition (QueryCondition, optional): Query condition. Defaults to None.
Returns:
List[str]: List of task_ids.
'''
class Converter(ABC):
'''
Convert data to dataset row.
'''
@abstractmethod
def to_dataset_row(self, task_experience: List[DataRow]) -> T:
'''
Convert task experience to dataset row.
Args:
task_experience (List[DataRow]): List of task experience.
Returns:
T: type of dataset row.
'''
class InMemoryStorage(Storage):
'''
In-memory storage for storing and sampling data.
'''
def __init__(self, max_capacity: int = 10000):
self._data: Dict[str, List[DataRow]] = {}
self._max_capacity = max_capacity
self._fifo_queue = [] # (task_id)
def add(self, data: DataRow):
if not data:
raise ValueError("Data is required")
if not data.exp_meta:
raise ValueError("exp_meta is required")
while self.size() >= self._max_capacity and self._fifo_queue:
oldest_task_id = self._fifo_queue.pop(0)
if oldest_task_id in self._data:
del self._data[oldest_task_id]
if data.exp_meta.task_id not in self._data:
self._data[data.exp_meta.task_id] = []
self._data[data.exp_meta.task_id].append(data)
self._fifo_queue.append(data.exp_meta.task_id)
if data.exp_meta.task_id not in self._data:
self._data[data.exp_meta.task_id] = []
self._data[data.exp_meta.task_id].append(data)
def add_batch(self, data_batch: List[DataRow]):
for data in data_batch:
self.add(data)
def size(self, query_condition: QueryCondition = None) -> int:
return len(self.get_all(query_condition))
def get_paginated(self, page: int, page_size: int, query_condition: QueryCondition = None) -> List[DataRow]:
if page < 1:
raise ValueError("Page must be greater than 0")
if page_size < 1:
raise ValueError("Page size must be greater than 0")
all_data = self.get_all(query_condition)
start_index = (page - 1) * page_size
end_index = start_index + page_size
return all_data[start_index:end_index]
def get_all(self, query_condition: QueryCondition = None) -> List[DataRow]:
all_data = []
query_filter = None
if query_condition:
query_filter = QueryFilter(query_condition)
for data in self._data.values():
if query_filter:
all_data.extend(query_filter.filter(data))
else:
all_data.extend(data)
return all_data
def get_by_task_id(self, task_id: str) -> List[DataRow]:
return self._data.get(task_id, [])
def get_bacth_by_task_ids(self, task_ids: List[str]) -> Dict[str, List[DataRow]]:
return {task_id: self._data.get(task_id, []) for task_id in task_ids}
def clear(self):
self._data = {}
self._fifo_queue = []
class RandomTaskSample(TaskSampler):
'''
Randomly sample data from the storage.
'''
def sample_task_ids(self,
storage: Storage,
batch_size: int,
query_condition: QueryCondition = None) -> List[str]:
total_size = storage.size(query_condition)
if total_size <= batch_size:
return storage.get_all(query_condition)
sampled_task_ids = set()
page_size = min(100, batch_size * 2)
total_pages = ceil(total_size/page_size)
visited_pages = set()
while len(sampled_task_ids) < batch_size and len(visited_pages) < total_pages:
page = random.choice(
[p for p in range(1, total_pages+1) if p not in visited_pages])
visited_pages.add(page)
current_page = storage.get_paginated(
page, page_size, query_condition)
if not current_page:
continue
current_page_task_ids = set(
[data.exp_meta.task_id for data in current_page if data.exp_meta.task_id not in sampled_task_ids])
sample_count = min(len(current_page_task_ids),
batch_size - len(sampled_task_ids))
sampled_task_ids.update(random.sample(
list(current_page_task_ids), sample_count))
return list(sampled_task_ids)
class DefaultConverter(Converter):
'''
Default converter do nothing.
'''
def to_dataset_row(self, task_experience: List[DataRow]) -> List[DataRow]:
return task_experience
class ReplayBuffer:
'''
Replay buffer for storing and sampling data.
'''
def __init__(
self,
storage: Storage = InMemoryStorage()
):
self._storage = storage
def store(self, data: DataRow):
'''
Store data in the replay buffer.
'''
if not data:
raise ValueError("Data is required")
self._storage.add(data)
def store_batch(self, data_batch: List[DataRow]):
'''
Store batch of data in the replay buffer.
'''
if not data_batch:
raise ValueError("Data batch is required")
self._storage.add_batch(data_batch)
def sample_task(self,
sampler: TaskSampler = RandomTaskSample(),
query_condition: QueryCondition = None,
converter: Converter = DefaultConverter(),
batch_size: int = 1000) -> List[T]:
'''
Sample Task from the replay buffer and convert to dataset row.
DefaultConverter return List[DataRow]
'''
sampled_task = sampler.sample_tasks(
self._storage, batch_size, query_condition)
return [converter.to_dataset_row(task_experiences) for task_experiences in sampled_task.values()]
def sample(self,
sampler: Sampler = RandomTaskSample(),
query_condition: QueryCondition = None,
converter: Converter = DefaultConverter(),
batch_size: int = 1000) -> List[T]:
'''
Sample data from the replay buffer and convert to dataset row.
DefaultConverter return List[DataRow]
'''
sampled_data = sampler.sample(
self._storage, batch_size, query_condition)
return converter.to_dataset_row(sampled_data)
|