code: def test_integration_with_policy_tasks(self):
      # Integration test for policy + value training and eval.
      optimizer = opt.Adam()
      lr_schedule = lr_schedules.constant(1e-3)
      advantage_estimator = advantages.td_k(
          gamma=self._task.gamma, margin=1
      )
      policy_dist = distributions.create_distribution(
          self._task.action_space
      )
      body = lambda mode: tl.Dense(64)
      train_model = models.PolicyAndValue(policy_dist, body=body)
      eval_model = models.PolicyAndValue(policy_dist, body=body)
  
      head_selector = tl.Select([1])
      value_train_task = value_tasks.ValueTrainTask(
          self._trajectory_batch_stream,
          optimizer,
          lr_schedule,
          advantage_estimator,
          model=train_model,
          target_model=eval_model,
          head_selector=head_selector,
      )
      value_eval_task = value_tasks.ValueEvalTask(
          value_train_task, head_selector=head_selector
      )
  
      # Drop the value head - just tl.Select([0]) would pass it, and it would
      # override the targets.
      head_selector = tl.Select([0], n_in=2)
      policy_train_task = policy_tasks.PolicyTrainTask(
          self._trajectory_batch_stream,
          optimizer,
          lr_schedule,
          policy_dist,
          advantage_estimator,
          # Plug a trained critic as our value estimate.
          value_fn=value_train_task.value,
          head_selector=head_selector,
      )
      policy_eval_task = policy_tasks.PolicyEvalTask(
          policy_train_task, head_selector=head_selector
      )
  
      loop = training.Loop(
          model=train_model,
          eval_model=eval_model,
          tasks=[policy_train_task, value_train_task],
          eval_tasks=[policy_eval_task, value_eval_task],
          eval_at=(lambda _: True),
          # Switch the task every step.
          which_task=(lambda step: step % 2),
      )
      # Run for a couple of steps to make sure there are a few task switches.
      loop.run(n_steps=10)
  
file_path: trax/trax/rl/value_tasks_test.py
imports: from absl.testing import absltest
  import numpy
  from trax import layers
  from trax import models
  from trax import optimizers
  from trax.rl import advantages
  from trax.rl import distributions
  from trax.rl import policy_tasks
  from trax.rl import task
  from trax.rl import value_tasks
  from trax.supervised import lr_schedules
  from trax.supervised import training
label_desc: Use image node_id 6 for calling the ValueTasksTest obj's underlying member method code with example usage: obj.test_integration_with_policy_tasks() without return types