<jupyter_start><jupyter_text>A notebook containing training logic for a robot to lift a cube using the mpo algorithm<jupyter_code># Include all the imports here
from typing import Dict, Sequence

from absl import app
from absl import flags
import acme
from acme import specs
from acme import types
from acme import wrappers
from acme.agents.tf import dmpo
from acme.tf import networks
from acme.tf import utils as tf2_utils

import numpy as np
import sonnet as snt
import robosuite as suite
from robosuite.wrappers import GymWrapper
from robosuite.controllers import load_controller_config
# Prepare the environment
env_config = {
     "control_freq": 20,
    "env_name": "Lift",
    "hard_reset": False,
    "horizon": 500,
    "ignore_done": False,
    "reward_scale": 1.0,
    "camera_names": "frontview",
    "robots": [
      "Panda"
    ]
}
controller_config = load_controller_config(default_controller="OSC_POSITION")

keys = ["object-state"]
for idx in range(1):
    keys.append(f"robot{idx}_proprio-state")
def make_environment(env_config, controller_config, keys):
    env_suite = suite.make(**env_config,
                 has_renderer=False,
                 has_offscreen_renderer=False,
                 use_camera_obs=False,
                 reward_shaping=True,
                 controller_configs=controller_config,
                 )
    env = GymWrapper(env_suite, keys=keys)
    env = wrappers.gym_wrapper.GymWrapper(env)
    env = wrappers.SinglePrecisionWrapper(env)
    
    spec = specs.make_environment_spec(env)
    
    return env, spec
env, spec = make_environment(env_config, controller_config, keys)
# Prepare the agent

def make_networks(
    action_spec: specs.BoundedArray,
    policy_layer_sizes: Sequence[int] = (256, 256, 256),
    critic_layer_sizes: Sequence[int] = (512, 512, 256),
    vmin: float = -150.,
    vmax: float = 150.,
    num_atoms: int = 51,
) -> Dict[str, types.TensorTransformation]:
      """Creates networks used by the agent."""

      # Get total number of action dimensions from action spec.
      num_dimensions = np.prod(action_spec.shape, dtype=int)

      # Create the shared observation network; here simply a state-less operation.
      observation_network = tf2_utils.batch_concat

      # Create the policy network.
      policy_network = snt.Sequential([
          networks.LayerNormMLP(policy_layer_sizes),
          networks.MultivariateNormalDiagHead(num_dimensions)
      ])

      # The multiplexer transforms concatenates the observations/actions.
      multiplexer = networks.CriticMultiplexer(
          critic_network=networks.LayerNormMLP(critic_layer_sizes),
          action_network=networks.ClipToSpec(action_spec))

      # Create the critic network.
      critic_network = snt.Sequential([
          multiplexer,
          networks.DiscreteValuedHead(vmin, vmax, num_atoms),
      ])

      return {
          'policy': policy_network,
          'critic': critic_network,
          'observation': observation_network,
      }
agent_networks = make_networks(spec.actions)
# construct the agent
agent = dmpo.DistributionalMPO(
      environment_spec=spec,
      policy_network=agent_networks['policy'],
      critic_network=agent_networks['critic'],
      observation_network=agent_networks['observation'],  # pytype: disable=wrong-arg-types
  )
# Start the training process
loop = acme.EnvironmentLoop(env, agent)
num_episodes = 100
loop.run(num_episodes=num_episodes)<jupyter_output>WARNING:tensorflow:From /home/mohan/miniconda3/envs/acme/lib/python3.8/site-packages/tensorflow_probability/python/distributions/distribution.py:345: calling Independent.__init__ (from tensorflow_probability.python.distributions.independent) with reinterpreted_batch_ndims=None is deprecated and will be removed after 2022-03-01.
Instructions for updating:
Please pass an integer value for `reinterpreted_batch_ndims`. The current behavior corresponds to `reinterpreted_batch_ndims=tf.size(distribution.batch_shape_tensor()) - 1`.
{'fraction_non_alphanumeric': 0.08608608608608609, 'fraction_numerical': 0.013513513513513514, 'mean_word_length': 3.81566265060241, 'pattern_counts': {'":': 8, '<': 4, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '1', 'dir': 'jupyter-structured-clean-dedup', 'id': '497093', 'license': 'Apache-2.0', 'n_tokens_mistral': 1263, 'n_tokens_neox': 1205, 'n_words': 287, 'path': 'examples/lift_mpo.ipynb', 'repo_name': 'mohanksriram/acme', 'size': '3996'}