# Best way to have custom value state + LSTM

**URL:** <https://discuss.ray.io/t/best-way-to-have-custom-value-state-lstm/2029>\
**Category:** RLlib\
**Created:** [May 3, 2021, 4:08pm UTC](https://discuss.ray.io/t/best-way-to-have-custom-value-state-lstm/2029 "2021-05-03T16:08:02Z")\
**Posts on this page:** 10\
**Page:** 1

<div class="post-metadata">

**Author:** ![nathanlct](https://avatars.discourse-cdn.com/v4/letter/n/919ad9/32.png) [@nathanlct](https://discuss.ray.io/u/nathanlct)\
**Post date:** [May 3, 2021, 4:08pm UTC](https://discuss.ray.io/t/best-way-to-have-custom-value-state-lstm/2029/1 "2021-05-03T16:08:02Z")

</div>

Hi,

I’m doing some training using PPO, and I would like the value function to have additional states that the policy doesn’t have.

By default, the FullyConnectedNetwork looks like this (14 should be 7 here):

 ![image](https://us1.discourse-cdn.com/flex020/uploads/ray/original/1X/7e35cd844f79c67e17ecccb73f8fbba1115b8ff4.jpeg)

I slightly modify it by splitting the input layer, this way I can give additional states to the value function input. I have a state space of size 14 and split it in 2, passing the first 7 to policy\_obs and the 7 last to value\_obs.

 ![image](https://us1.discourse-cdn.com/flex020/uploads/ray/original/1X/afe1c04b76ecc78f9ea6b0a5014b031f56f47bb2.jpeg)

This seems to work fine, however I’d like to add an lstm on top of it, and the `use_lstm` wrapper doesn’t work by default (cf. error below). I was wondering if there is an easier way to do this (giving additional states to the VF) that I didn’t find? That would save me from having to add an lstm manually to my custom net. Or did I miss something with the `use_lstm`?

`ValueError: Input 0 of layer fc_value_1 is incompatible with the layer: expected axis -1 of input shape to have value 14 but received input with shape [None, 7]`

Thanks!

---

<div class="post-metadata">

**Author:** ![mannyv](https://sea2.discourse-cdn.com/flex020/user_avatar/discuss.ray.io/mannyv/32/606_2.png) [@mannyv](https://discuss.ray.io/u/mannyv)\
**Post date:** [May 3, 2021, 4:55pm UTC](https://discuss.ray.io/t/best-way-to-have-custom-value-state-lstm/2029/2 "2021-05-03T16:55:38Z")

</div>

Hi @nathanlct,

The value function look different when using the lstm wrapper than when not.

When you use the lstm wrapper there is only one set of layers going into the lstm model and the value layers use the final state coming from the lstm as the input. So instead of two heads going into the lstm you have two heads coming out of the lstm. You can see images I created of the two cases in this post [What is the intended architecture of PPO vf\_share\_layers=False when using an LSTM](https://discuss.ray.io/t/what-is-the-intended-architecture-of-ppo-vf-share-layers-false-when-using-an-lstm/1579).

I think this means you are going to have to add your own lstm. Luckily this is pretty straightforward with rllib. You do have to think how and where you want to do this seperation though. Are you going to add a layer before the lstm then feed the lstm a concatenation of a policy and value embedding layer? Are you going to treat them as pass through inputs and feed them after the lstm? Are you going to have two lstms? Will only your policy use an lstm and have the value function be just fc layers?

---

<div class="post-metadata">

**Author:** ![smorad](https://sea2.discourse-cdn.com/flex020/user_avatar/discuss.ray.io/smorad/32/272_2.png) [@smorad](https://discuss.ray.io/u/smorad)\
**Post date:** [May 11, 2021, 10:03am UTC](https://discuss.ray.io/t/best-way-to-have-custom-value-state-lstm/2029/3 "2021-05-11T10:03:25Z")

</div>

It’s quite straightforward, but you will have to implement your own model. See [ray/rnn\_model.py at 4795048f1b3779658e8b0ffaa05b1eb61914bc60 · ray-project/ray · GitHub](https://github.com/ray-project/ray/blob/4795048f1b3779658e8b0ffaa05b1eb61914bc60/rllib/examples/models/rnn_model.py#L77)

---

<div class="post-metadata">

**Author:** ![mickelliu](https://sea2.discourse-cdn.com/flex020/user_avatar/discuss.ray.io/mickelliu/32/1347_2.png) [@mickelliu](https://discuss.ray.io/u/mickelliu)\
**Post date:** [January 15, 2022, 12:34pm UTC](https://discuss.ray.io/t/best-way-to-have-custom-value-state-lstm/2029/4 "2022-01-15T12:34:59Z")

</div>

hi @mannyv, if I want two separate LSTMs one for the value and one for the actor, how do I manage the hidden states of the value branch? I only saw examples showing managing hidden states of one LSTM.

Thanks.

---

<div class="post-metadata">

**Author:** ![mannyv](https://sea2.discourse-cdn.com/flex020/user_avatar/discuss.ray.io/mannyv/32/606_2.png) [@mannyv](https://discuss.ray.io/u/mannyv)\
**Post date:** [January 15, 2022, 1:04pm UTC](https://discuss.ray.io/t/best-way-to-have-custom-value-state-lstm/2029/5 "2022-01-15T13:04:44Z")

</div>

Hi @mickelliu,

I saw you also posted a question on this topic

> [@What is the intended architecture of PPO vf\_share\_layers=False when using an LSTM](https://discuss.ray.io/t/what-is-the-intended-architecture-of-ppo-vf-share-layers-false-when-using-an-lstm/1579/3):
>
> Thinking about this more, this would actually take a lot of changes to make this work: our Models have a separate value\_function method, which would - in this case - also have to take a state input list (it currently doesn’t). the LSTM wrapper would need to know about the “location” of the separate value branch in the wrapped model (as it would have to call that branch separately). This would be ok for default models (under RLlib’s own control), but would become quite difficult for custom mode…

You are right that there was no resolution to that issue and the reason is because of what you brought up here. There is no good way that I (and I think @sven1977, though I don’t want to speak for him) have found to carry around state for the value network.

The value network is not used during rollouts so you will not have states for each transition to use during training.

Two thoughts,

1. You could construct your model to retrieve the value on every call to forward. Then you would have states for rollouts too but you would also need to modify get\_initial\_state to return the combined states for both branches and break the states up in your forward method. I am not certain if this would be compatible with the default handling of states in the library.

2. You could implement something like burnin used in R2D2 and pass in some number of zeroed states to warm up the lstm.

If you find something that works please do let us know.

* * *

> **[Manny is helping people do cool things with RL](https://www.buymeacoffee.com/mannyv)**
>
> Hey, so glad we got to connect. If you found my suggestions useful feel free to buy me a coffee!

---

<div class="post-metadata">

**Author:** ![mickelliu](https://sea2.discourse-cdn.com/flex020/user_avatar/discuss.ray.io/mickelliu/32/1347_2.png) [@mickelliu](https://discuss.ray.io/u/mickelliu)\
**Post date:** [January 15, 2022, 1:12pm UTC](https://discuss.ray.io/t/best-way-to-have-custom-value-state-lstm/2029/6 "2022-01-15T13:12:25Z")

</div>

Hi @mannyv, thanks for your quick response!

I have an idea and I am just about to run my custom model to see if this will work.  
Would it be ok if I just concat my value\_hidden\_states with action\_hidden\_states and send them together away upon each `forward_run` call?

More concretely:

```auto
# Two separate LSTMs with two separate branches...

            self.actor_layers[-1]._model[0].weight.new(1, self.lstm_state_size).zero_().squeeze(0),
            self.actor_layers[-1]._model[0].weight.new(1, self.lstm_state_size).zero_().squeeze(0),
            self.value_layers[-1]._model[0].weight.new(1, self.lstm_state_size).zero_().squeeze(0),
            self.value_layers[-1]._model[0].weight.new(1, self.lstm_state_size).zero_().squeeze(0)

    @override(ModelV2)
    def value_function(self):
        assert self._values is not None, "must call forward() first"
        return torch.reshape(self.value_branch(self._values), [-1])

    @override(TorchRNN)
    def forward_rnn(self, inputs, state, seq_lens):

        self._features, [h1, c1] = self.actor_lstm(
            self.actor_layers(inputs), [torch.unsqueeze(state[0], 0),
                torch.unsqueeze(state[1], 0)])
        action_out = self.action_branch(self._features)

        self._values, [h2, c2] = self.value_lstm(
            self.value_layers(inputs), [torch.unsqueeze(state[2], 0),
                torch.unsqueeze(state[3], 0)])

        return action_out, [torch.squeeze(h1, 0), torch.squeeze(c1, 0), torch.squeeze(h2, 0), torch.squeeze(c2, 0)]

```

If this doesn’t work out, should I change my view\_requirement or anything?

---

<div class="post-metadata">

**Author:** ![mannyv](https://sea2.discourse-cdn.com/flex020/user_avatar/discuss.ray.io/mannyv/32/606_2.png) [@mannyv](https://discuss.ray.io/u/mannyv)\
**Post date:** [January 15, 2022, 1:49pm UTC](https://discuss.ray.io/t/best-way-to-have-custom-value-state-lstm/2029/7 "2022-01-15T13:49:14Z")

</div>

@mickelliu,

That is what I was trying to describe as approach 1.

You will need to save the value states as a member in forward for them to be available in the call to value. Same as you do with \_values.

I do not think you will need to modify the ViewRequirements.

---

<div class="post-metadata">

**Author:** ![mickelliu](https://sea2.discourse-cdn.com/flex020/user_avatar/discuss.ray.io/mickelliu/32/1347_2.png) [@mickelliu](https://discuss.ray.io/u/mickelliu)\
**Post date:** [January 15, 2022, 1:51pm UTC](https://discuss.ray.io/t/best-way-to-have-custom-value-state-lstm/2029/8 "2022-01-15T13:51:56Z")

</div>

Thanks @mannyv, it seems that the code didn’t pop any error. I will provide you with an update later.

---

<div class="post-metadata">

**Author:** ![mickelliu](https://sea2.discourse-cdn.com/flex020/user_avatar/discuss.ray.io/mickelliu/32/1347_2.png) [@mickelliu](https://discuss.ray.io/u/mickelliu)\
**Post date:** [January 16, 2022, 7:46am UTC](https://discuss.ray.io/t/best-way-to-have-custom-value-state-lstm/2029/9 "2022-01-16T07:46:56Z")

</div>

Hi @mannyv,

I ran the separate branch LSTM model for 10+ hours and it works pretty well, I didn’t see any problem yet. For our custom env it’s always been the case that separate value and actor network will yield better performance.

---

<div class="post-metadata">

**Author:** ![dezhi](https://sea2.discourse-cdn.com/flex020/user_avatar/discuss.ray.io/dezhi/32/2189_2.png) [@dezhi](https://discuss.ray.io/u/dezhi)\
**Post date:** [April 10, 2022, 7:05pm UTC](https://discuss.ray.io/t/best-way-to-have-custom-value-state-lstm/2029/10 "2022-04-10T19:05:20Z")

</div>

Hi @mickelliu,

Can you please share your separate branch LSTM model here? I think i am missing something in writing the code…