Hang in end_episode

[jaguilar]

Friends,

I'm using dm-reverb 0.14.0, and I'm getting a hang during end_episode() in my SAC implementation. This implementation is, as far as I can tell, copied verbatim from the one in the SAC minotaur demo on the tf-agents website. The only difference is I'm using a different env of my own design.

As far as I can tell, the sequence of events is:

• In ReverbAddTrajectoryObserver.__call__, detect a boundary.
• reset(False)
• self._writer is not None, so we call end_episode()

That end_episode call never returns and no feedback is given in the logs as to why this might be.

For what it's worth, I do not have this problem when I run the SAC minitaur example from the tutorial. It seems like I can set up my revert client just fine there.

---

As I'm writing this, I found a workaround that appears to resolve the issue at least with respect to the problem I was having. In __call__, I adjusted the trajectory.is_boundary() condition at the end with:

else:
  self._get_writer().flush(10)

After that, I at least no longer hang. I'm still having other problems but at least this one seems to have been resolved.

[jaguilar]

After fixing that issue, I'm now apparently getting a hang when trying to fetch a dataset sample from reverb.

In case it's any help in diagnosing the issue, here's my reverb setup code:

        rate_limiter = reverb.rate_limiters.SampleToInsertRatio(
            samples_per_insert=3.0, min_size_to_sample=3, error_buffer=3.0
        )
        table_name = "uniform_table_asdf"
        table = reverb.Table(
            table_name,
            max_size=replay_buffer_capacity,
            sampler=reverb.selectors.Uniform(),
            remover=reverb.selectors.Fifo(),
            rate_limiter=rate_limiter,
        )
        reverb_server = reverb.Server([table])
        reverb_replay = reverb_replay_buffer.ReverbReplayBuffer(
            tf_agent.collect_data_spec,
            sequence_length=2,
            table_name=table_name,
            local_server=reverb_server,
        )
        dataset = reverb_replay.as_dataset(
            sample_batch_size=batch_size, num_steps=2
        ).prefetch(50)
        experience_dataset_fn = lambda: dataset
        # ...
        agent_learner = learner.Learner(
            tempdir,
            train_step,
            tf_agent,
            experience_dataset_fn,
            triggers=learning_triggers,
            strategy=self._strategy,
        )

I instrumented single_train_step() in tf_agents/train/learner.py like this:

  def single_train_step(self, iterator):
    tf.print('single_train_step')
    sample = next(iterator)
    tf.print('next done')
    # ...

And the end of my logging is like this:

INFO:tensorflow:Allowlisted <function _normalize_outputs at 0x7f00e7d48550>: from cache
INFO:tensorflow:Allowlisted <function _normalize_outputs at 0x7f00e7d48550>: from cache
/home/jaguilar/projects/Two-Wheel-Robot-DeepRL/.venv/lib/python3.10/site-packages/tensorflow/python/saved_model/nested_structure_coder.py:458: UserWarning: Encoding a StructuredValue with type tfp.distributions.Deterministic_ACTTypeSpec; loading this StructuredValue will require that this type be imported and registered.
  warnings.warn("Encoding a StructuredValue with type %s; loading this "
INFO:tensorflow:Assets written to: /tmp/tmpfc1ifzsrtfagents/policies/greedy_policy/assets
INFO:tensorflow:Assets written to: /tmp/tmpfc1ifzsrtfagents/policies/greedy_policy/assets
step 0
step = 0: AverageReturn = -7.179125, AverageEpisodeLength = 0.000000
collect_actor.run
collect_actor.run done
learner.run
in scope
single_train_step
[reverb/cc/client.cc:165] Sampler and server are owned by the same process (319501) so Table uniform_table_asdf is accessed directly without gRPC.
[reverb/cc/client.cc:165] Sampler and server are owned by the same process (319501) so Table uniform_table_asdf is accessed directly without gRPC.
<hang forever>

This is running on WSL2, for what it's worth. I have had Reverb working with tf-agents, possibly at a different Reverb version though (it's hard to know as I've been muddling my way through dependency hell).

What can I do to debug why sampling the dataset is hanging?