google / google/brax

Increase modularization of training wrappers

Open
#422 4 comments 0 reactions 1 assignee Claimed by @vyeevani View on GitHub
Dominant language
Jupyter Notebook
Stars
3.2k
Forks
349
PR merge metrics
No merged PRs in 30d

Description

The training wrappers for auto reset and episode wrapper are leaking info to each other.

This is a bigger problem if people want to stack their own wrappers. For example, I'd like to write a meta-episode wrapper that takes multiple episodes and aggregates them into a single meta episode for use in a meta-RL setting. The wrapper that I'm writing needs to track the number of episodes so that it can do a meta episode reset when it reaches some watermark. However, the auto reset would break this since it wouldn't reset the meta wrapper under it.

At a high level, I'd like to separate the state of the auto reset wrapper from the environments that it's wrapping. I propose to do this by caching the initial state of the environment and the current state of the environment in the info, and only working on that.

1. In reset, get the base state and store it in info with two separate keys: initial_base_state, current_base_state
2. In reset, take the initial_base_state's: observation, reward, and done and package it along with the info from step (1)
3. In step, if you are done, then return the same aggregated state as step (2)
4. In step, if not done, then return the evolved state by updating the current base state.

Note, through this process, you never need to return the pipeline state through the state itself.

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.