[feature request] Implement goal-parameterized algorithms (HER)

[ccolas]

I'd like to implement Hindsight Experience Replay (HER). This can be based on a whatever goal-parameterized RL off-policy algorithm.

Goal-parameterized architectures: it requires a variable for the current goal and one for the current outcome. By outcome, I mean anything that is requires to compute the current outcome in the process of targeting the goal, e.g. the RL task is to reach a 3D target (the goal) with a robotic hand. The position of the target is the goal, the position of the hand is the outcome. The reward is a function of the distance between the two. Goal and outcome are usually subparts of the state space.

How Gym handles this: In Gym, there is a class called GoalEnv to deal with such environments.

• The variable observation_space is replaced by another class that contains the true observation space observation_space.spaces['observation'], goal space (observation_space.spaces['desired_goal']) and outcome space(observation_space.spaces['achieved_goal']).
• The observation returned first by env.step is now a dictionnary obs['observation'], obs['desired_goal'], obs['achieved_goal'].
• The environment defines a reward function (compute_reward), that takes as argument the goal and the outcome to return the reward
• It also contains a sample_goal function that simply samples a goal uniformly from the goal space.

Stable-baselines does not consider this so far. The replay buffer, BasePolicy, BaseRLModels OffPolicyRLModels only consider observation, and are not made to include a notion of goal or outcome. Two solutions:

1. Adapt these default classes to allow the representation of goals and outcomes in option.
2. Concatenate the goal and outcome to the observation everywhere so that the previously mention classes don't see the difference. This requires to keep track of the indices of goals and outcomes in the full obs_goal_outcome vector. Ashley started to do this from what I understood. However, he did not take into account the GoalEnv class of Gym, which I think we should use as it's kind of neat, and it's used for the Robotic Fetch environment, which are kind of the only one generally used so far.

I think the second is more clear as it separates observation from goals and outcomes, but probably it would make the code less easy to follow, and would require more changes than the first option. So let's go for the first as Ashley started.

First thoughts on how it could be done.

1. we need (as Ashley started to do), a wrapper around the gym environment. GoalEnv are different from usual env because they return a dict in place of the former observation vector. This wrapper would unpack the observation in obs, goal, outcome from the GoalEnv.step. It would return a concatenation of all of those. Ashley considered that the goal was in the observation space, so that the concatenation was twice as long as the observation. This is generally not true. So we would need to keep as attribute the size of the goal and outcome spaces. It would keep the different spaces as attributes, keep the function to sample goals, and the reward function.

2. A multi-goal replay buffer to implement HER replay. It takes the observation from the buffer and redecompose it in obs, goal and outcome before performing replay.

I think it does not require too much work after what Ashley started to do. It would be a few modifications to integrate the GoalEnv of gym, as it is a standard way to use multi-goal environments. Then correct the assumption he made about the dimension of the goal.

If you're all ok, I will start in that direction and test them on the Fetch environments. In the baselines, their performance is achieved with 19 processes in parallel. They basically average the update of the 19 actors. I'll try first without parallelization.

[araffin]

Some questions to clarify my understanding:

• goal and outcome are part of the observation? (Is that what you mean by "Goal and outcome are usually subparts of the state space." ?)
• in a GoalEnv, does observation key contains both outcome and goal or do we need to concatenate them? (I think we need to concatenate them but I'm not sure)

So let's go for the first as Ashley started.

You meant second, no?

I think it does not require too much work after what Ashley started to do.

Good news then =)

, I will start in that direction and test them on the Fetch environments

For debugging and for your tests, Fetch env is ok. However, we will need either to adapt an existing env (e.g. Pendulum-v0) or create an artificial one (e.g. Identity env) in order to write unit tests for it (Mujoco requires a license and we don't have one for travis...).

For an open source Goal Env, you can also look into the parking env by @eleurent in https://github.com/eleurent/highway-env

[ccolas]

• I wanted to say: goal and outcome spaces are usually subspaces of the state space (but not always). Example: FetchReach (place gripper at 3D location). The observation is made of the position and velocities of all the joints, while the goal is only a 3D location, the outcome is the 3D position of the hand. In that case the outcome is a subpart of the state.

• In GoalEnv, the observation return is a dictionary with three keys: observation, achieved_goal and desired_goal. What I'm doing now is wrapping the gym env to concatenate everything into one observation, and having some variable saying which indexes are what, so that HER can do its substitution. After that, I wrap in a DummyVecEnv.

(Yes I meant the second, as Ashley).

I guess we could modify easily Pendulum or MountainCarContinuous yes.

• Last problem I encountered: DDPG considers the return to be the sum of all rewards. It's not absolutely general, sometimes it's only the last reward (in the Fetch env for instance). A solution can be to add a string parameter to the class return_func='last' or 'sum'.

[araffin]

Last problem I encountered: DDPG considers the return to be the sum of all rewards. It's not absolutely general, sometimes it's only the last reward

I thought return meant sum of discounted reward for one episode... In the Fetch env, I think it is a sparse reward (0 almost everywhere except when reaching the goal), so last reward is the same of the sum no?
Does it affect the training of the algorithm (the difference in formulation) or is it only for logging?

[ccolas]

In Fetch, always -1 and 0 when the goal is touched. It's only for logging, the algorithm uses the transition-based rewards. In that case, saying that the sum is -49 for 50 steps for instance, does not indicate whether the goal was reached in the middle of the episode (in that case the episode is not solved), or at the end (episode solved).

[araffin]

I see... then for GoalEnv, it makes sense to have that feature (showing last return only).

[ccolas]

The problem is that it's computed outside of the env, in the learn function of the algorithm. That would require to update all algorithm to allow that.

[araffin]

That would require to update all algorithm to allow that.

I would update only algorithms that can be used with HER (and therefore GoalEnv), I'm not sure if it makes sense for other type of env.

[keshaviyengar]

Is this still being worked on? This feature is something I'm quite interested in.

[ccolas]

Update:
I implemented it but it's still bugged somehow. It runs but it does not learn anything on FetchReach. It's on a fork from stable baselines. I can give you access if you want to check it out.

In the meantime I also tried implementing it using sac's spinning up. It learns perfectly on FetchReach but nothing on FetchPush. I'm quite puzzled because I already tried to reproduce results on FetchPush last year using a TD3 base and it also worked perfectly on FetchReach and not on FetchPush.

Either I do something wrong, or there is some special trick in the OpenAI Baselines that I didn't catch. Their version uses 19 worked in parallel, each doing 2 rollouts, computing an update using a batch of 256 and summing the 19 updates (yes they sum, they don't average). I would say it's roughly equivalent to do 38 collection rollout, then to use a 19 times bigger batch size and 19 times bigger learning rate (the sum of 19 updates). I tried this also but it got even worse.

I don't have much time these days so it's on pause right now.

[araffin]

@ccolas don't hesitate to ping me if you need some help to double check some parts ;)

[araffin]

@ccolas @hill-a I'm taking over for this one.
I have read the paper again (and some implementations) and started to implement it (mostly from scratch).
There is one difference that was not mentioned that I also spotted: the original paper creates new transitions (by sampling goals) online after each rollout (not after each step) but this should be ok (in the sense we can overcome that limitation by using a custom replay buffer)

Here is my current plan (and current progress):

• I created a FlippedBitEnv to iterate and debug quickly (both with continuous and discrete actions)
• I plan to ensure that the environment inherits from GoalEnv
• I started from scratch but in fact, now, my current plan really ressemble what @hill-a did ^^#
• my current implementation requires the user to pass a sampling_strategy and a get_achieved_goal callable. The first argument refers to HER replay strategy, the second is a callable that convert an observation to a goal (so we can use env.compute_reward), I'm not sure if it is needed (still in early stage of development)

To overcome current limitation of stable-baselines (dict obs are not supported), I'll do something similar to @hill-a , using a wrapper over the environment.

Roadmap:

• implement dev env FlippedBitEnv
• have a first limited version working on FlippedBitEnv , i.e. DQN with FINAL strategy
• add support for DDPG and SAC and make it work for the FlippedBitEnv
• implement more replay strategies
• make them work on Highway env
• polish code
• documentation
• add support for more type of env obs spaces (Discrete, Binary, MultiDiscrete?)
• add full support for all env obs spaces (e.g. where the observation and goal space are not of the same type)

Note: I consider the last point to be not a priority

[araffin]

@ccolas
I have a first working draft with DQN (tried with a flipped bit env with n_bits=30 and it worked =) ),
you can have a look at it even if it is not fully polished yet.

In short, I made those choices, which is mostly wrappers:

• HER will be just a wrapper around a RL model (which makes sense)
• instead of creating a new replay buffer class, I created a wrapper for HER replay buffer, that is were all the magic happens
• i have a env wrapper that concatenate dict observation and can do the inverse operation

to sum it up, the implementation now looks pretty simple, it is just a wrapper around a model and an env (I still don't understand why the original baselines were so complicated).
My next step is to make it work with SAC ;)