support tfrecord for datastore

[youliangtan]

Summary

This PR introduces rlds loader and reader for all existing datastore. It supports both TrajectoryBuffer and ReplayBuffer Datastore. how it is being used is

This requires installation of

• tensorflow and tensorflow_dataset
• oxe_envlogger: https://github.com/rail-berkeley/oxe_envlogger

from oxe_envlogger.rlds_logger import RLDSLogger
from edgeml.data.jaxrl_data_store import ReplayBufferDataStore

rlds_logger =  RLDSLogger(
        observation_space=env.observation_space,
        action_space=env.action_space,
        dataset_name="test_rlds_env",
        directory=log_dir,
        max_episodes_per_file=5,  # TODO: arbitrary number
)
store = ReplayBufferDataStore(obs_space, action_space, cap, rlds_logger)

# insert
store.insert(DATA)   # every insert it will automatically get logged

...OTHER PROCESSES...

# exit
rlds_logger.close()   # explicityly flush and close the logger

• For more info, refer to the test case: edgeml/tests/test_tfds.py.
• an example is also shown in async_learner_actor.py by adding flag of --rlds_log_dir.

Gotchas

Some things might need further considerations, and even changes.

• the default traj buffer only supports array as input (no dict)
• we logged next_observations instead of observations, as it seems that next obs is the obs at the current timestep.
• ReplayBufferDataStore doesnt support truncation, which is sad.
• The TrajectoryBuffer method should require more cleanup.