instadeepai / instadeepai/flashbax

How to load saved buffer properly?

Open
#47 5 comments 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
283
Forks
22
PR merge metrics
No merged PRs in 30d

Description

I tried:

```
class ReplayBufferDataStore():
def __init__(
self,
env: gym.Env,
capacity: int,
sample_batch_size: int = 32,
priority_exponent: float = 0.8,
device: str = "gpu",
name: str = "replay_buffer",
checkpoint_path: str = None,
):
self.sample_batch_size = sample_batch_size
self.priority_exponent = priority_exponent
self.device = jax.devices(device)[0]

self.buffer = fbx.make_prioritised_flat_buffer(
max_length=capacity,
min_length=sample_batch_size,
sample_batch_size=sample_batch_size,
add_sequences=True,
add_batch_size=None,
priority_exponent=priority_exponent,
device=device,
)

# Preprocess the transition once to avoid redundant transformations
single_transition = self._initialize_single_transition(env)
self.state = self.buffer.init(single_transition)
self.state = jax.device_put(self.state, device=self.device)

self.vault = Vault(
vault_name=name,
experience_structure=self.state.experience,
rel_dir=os.path.join(os.path.dirname(checkpoint_path), "vaults"),
)
...

def save(self):
self.vault.write(self.state)

def load(self, vault_path: str):
vault_name = vault_path.split("/")[-2]
vault_uid = vault_path.split("/")[-1]
vault_path = os.path.dirname(os.path.dirname(vault_path))
vault = Vault(
vault_name=vault_name,
experience_structure=self.state.experience,
rel_dir=vault_path,
vault_uid=vault_uid,
)
state = vault.read()

loaded_experience = frozen_dict.freeze(state.experience)
self.state = _insert(self.buffer, self.state, loaded_experience)
```

`experience_structure` of loaded `state` doesn't match with `self.state`. For example, if there are 500 transitions stored in the buffer via vault, the loaded `state` size will be 500, but I initialized the buffer with size 100_000.

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.