Add memory optimization option to ReplayBuffer

- Introduce `optimize_memory` parameter to reduce memory usage in replay buffer
- Implement simplified memory optimization by not storing duplicate next_states
- Update learner server and buffer initialization to use memory optimization by default
This commit is contained in:
AdilZouitine
2025-02-25 19:04:58 +00:00
parent 20fee3d043
commit 24f93c755a
2 changed files with 296 additions and 71 deletions
+3
View File
@@ -154,6 +154,7 @@ def initialize_replay_buffer(
device=device,
state_keys=cfg.policy.input_shapes.keys(),
storage_device=device,
optimize_memory=True,
)
dataset = LeRobotDataset(
@@ -166,6 +167,7 @@ def initialize_replay_buffer(
capacity=cfg.training.online_buffer_capacity,
device=device,
state_keys=cfg.policy.input_shapes.keys(),
optimize_memory=True,
)
@@ -648,6 +650,7 @@ def train(cfg: DictConfig, out_dir: str | None = None, job_name: str | None = No
action_mask=active_action_dims,
action_delta=cfg.env.wrapper.delta_action,
storage_device=device,
optimize_memory=True,
)
batch_size: int = batch_size // 2 # We will sample from both replay buffer