This is the completion of google's rt-1 project code and can run directly.
# Copyright 2022 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""robotics_transformer API."""

# A new PyPI release will be pushed everytime `__version__` is increased
# When changing this, also update the CHANGELOG.md
__version__ = '0.1.0'

__init__.py
YiyangHuang-work/RT-1
"""读取RLDS数据集，详见：https://github.com/google-research/rlds  数据读取代码参考https://github.com/google-research/language-table"""

import dataclasses
import functools
from typing import Optional, Tuple
from clu import preprocess_spec
import jax
import rlds
import tensorflow as tf
import tensorflow_datasets as tfds
import tree

Features = preprocess_spec.Features


def create_datasets(
        rng,
        dataset_dirs,
        sequence_length,
        global_batch_size,
        target_width=320,
        target_height=180,
        random_crop_factor=None,
        cache=False,
        shuffle=True,
        shuffle_buffer_size=50_000,
        cache_dir=None,
        dataset_episode_num=10000
):
    """创建一个RLDS数据集."""

    builder = tfds.builder_from_directories(dataset_dirs)

    dataset_options = tf.data.Options()
    dataset_options.experimental_optimization.map_parallelization = True
    dataset_options.threading.private_threadpool_size = 48
    dataset_options.threading.max_intra_op_parallelism = 1

    ds = builder.as_dataset(
        split=f'train[{0}:{dataset_episode_num}]',
        decoders={"steps": {"observation": {"rgb": tfds.decode.SkipDecoding()}}},
        shuffle_files=True
    )

    def _pad_episode(episode, padding):
        first_item_tensor = episode["steps"].take(1).get_single_element()
        first_item_ds = tf.data.Dataset.from_tensors(first_item_tensor)

        first_item_mid_tensor = tf.nest.map_structure(
            tf.identity, first_item_tensor
        )
        first_item_mid_tensor[rlds.IS_FIRST] = False
        padding_ds = tf.data.Dataset.from_tensors(first_item_mid_tensor).repeat(
            padding
        )

        full_padding = rlds.transformations.concatenate(first_item_ds, padding_ds)
        episode["steps"] = rlds.transformations.concatenate(
            full_padding, episode["steps"].skip(1)
        )
        return episode

    ds = ds.map(
        functools.partial(_pad_episode, padding=sequence_length - 1),
        tf.data.AUTOTUNE,
    )

    def get_seqlen_pattern(step):
        return {
            rlds.OBSERVATION: tree.map_structure(
                lambda x: x[-sequence_length:], step[rlds.OBSERVATION]
            ),
            rlds.ACTION: tree.map_structure(
                lambda x: x[-sequence_length:], step[rlds.ACTION]
            ),
            rlds.IS_TERMINAL: tree.map_structure(
                lambda x: x[-sequence_length:], step[rlds.IS_TERMINAL]
            ),
        }

    ds = rlds.transformations.pattern_map_from_transform(
        episodes_dataset=ds,
        transform_fn=get_seqlen_pattern,
        respect_episode_boundaries=True,
    )

    if shuffle:
        shuffle_rng, rng = jax.random.split(rng)
        shuffle_rng = shuffle_rng[0]
        ds = ds.shuffle(shuffle_buffer_size, shuffle_rng)

    preprocessors = [
        DecodeAndRandomResizedCrop(
            random_crop_factor=random_crop_factor,
            resize_size=(target_height, target_width),
        ),
        TransformDict(),
    ]
    train_preprocess = preprocess_spec.PreprocessFn(
        preprocessors, only_jax_types=True
    )

    def _preprocess_fn(example_index, features):
        example_index = tf.cast(example_index, tf.int32)
        features[preprocess_spec.SEED_KEY] = (
            tf.random.experimental.stateless_fold_in(
                tf.cast(rng, tf.int64), example_index
            )
        )
        processed = train_preprocess(features)
        return processed

    ds = ds.enumerate().map(_preprocess_fn, num_parallel_calls=tf.data.AUTOTUNE)

    ds = ds.batch(global_batch_size, drop_remainder=True)

    if cache:
        ds = ds.cache(cache_dir)

    return ds


@dataclasses.dataclass(frozen=True)
class DecodeAndRandomResizedCrop(preprocess_spec.RandomMapTransform):
    """解析图像，提取随机crop, resize并归一化"""

    random_crop_factor: Optional[float] = None
    resize_size: Tuple[int, int] = (180, 320)

    def _transform(self, features, seed):
        image = features["observation"]["rgb"]
        shape = tf.io.extract_jpeg_shape(image[0])
        raw_height, raw_width = shape[0], shape[1]
        raw_height = tf.cast(raw_height, tf.float32)
        raw_width = tf.cast(raw_width, tf.float32)

        if self.random_crop_factor is None:
            random_crop_factor = 1.0
            offset_width = 0
            offset_height = 0
            scaled_height = raw_height
            scaled_width = raw_width
        else:
            random_crop_factor = tf.constant(
                self.random_crop_factor, dtype=tf.float32
            )
            scaled_height = raw_height * random_crop_factor
            scaled_width = raw_width * random_crop_factor

            next_rng, rng = tf.unstack(tf.random.experimental.stateless_split(seed))
            offset_height = tf.random.stateless_uniform(
                shape=(),
                seed=next_rng,
                minval=0,
                maxval=tf.cast(raw_height - scaled_height, dtype=tf.int32),
                dtype=tf.int32,
            )

            next_rng, rng = tf.unstack(tf.random.experimental.stateless_split(rng))
            offset_width = tf.random.stateless_uniform(
                shape=(),
                seed=next_rng,
                minval=0,
                maxval=tf.cast(raw_width - scaled_width, dtype=tf.int32),
                dtype=tf.int32,
            )

        def apply_decode_and_crop(image):
            image = tf.image.decode_and_crop_jpeg(
                image,
                [
                    offset_height,
                    offset_width,
                    tf.cast(scaled_height, tf.int32),
                    tf.cast(scaled_width, tf.int32),
                ],
                channels=3,
            )
            return image

        image = tf.map_fn(apply_decode_and_crop, image, dtype=tf.uint8)
        image = tf.cast(image, tf.float32) / 255.0
        image = tf.image.resize(image, self.resize_size)
        features["observation"]["rgb"] = image
        return features


@dataclasses.dataclass(frozen=True)
class TransformDict(preprocess_spec.RandomMapTransform):
    """将数据存放字典格式转换成网络所需数据字典格式."""

    def _transform(self, features, seed):
        """Applies all distortions."""
        action_lable = {
            "terminate_episode": tf.one_hot(tf.cast(features["is_terminal"], dtype=tf.int32), depth=2, dtype=tf.int32),
            "effector_target_translation": features["observation"]["effector_target_translation"]}
        train_observation = {"image": features["observation"]["rgb"],
                             "natural_language_embedding": tf.cast(features['observation']['instruction'],dtype=tf.float32)}
        features = {"action_lable": action_lable, "train_observation": train_observation}
        return features

data_loader/rlds_dataset_loader.py
YiyangHuang-work/RT-1
"""
robotic transformer(https://github.com/google-research/robotics_transformer)的多节点分布式训练代码,
采用tensorflow2的distribute.MultiWorkerMirroredStrategy(https://www.tensorflow.org/api_docs/python/tf/distribute/MultiWorkerMirroredStrategy)进行分布式训练，使用加载rlds(https://github.com/google-research/rlds)数据的方式进行数据的读取
使用方法：
    python distribute_worker_train.py --args = param, 其中args见代码中的get_args()
"""

import os
from robotics_transformer import transformer_network
from tensor2robot.utils import tensorspec_utils
from tf_agents.specs import tensor_spec
import time
from robotics_transformer.data_loader import rlds_dataset_loader
import tensorflow as tf
import jax
import argparse
import json

'''
Description:
    设置分布式训练参数
Parameters:
    None
Return:
    args:包含训练需要的各种参数，参数详情请见代码
'''
def get_args():
    parser = argparse.ArgumentParser(description='获得分布式训练参数')
    parser.add_argument('--single_gpu_batch_size', '-s', help='batch size for single gpu', default=8, type=int)
    parser.add_argument('--training_epoch', '-te', help='training epoch', default=100, type=int)  # 训练epoch
    parser.add_argument('--log_step', '-ls', help='log step', default=10, type=int)
    parser.add_argument('--dataset_dirs', '-d', help='dataset path', default="/adddisk1/dataset/language-table-sim/")
    parser.add_argument('--learning_rate', '-lr', help='learning rate', default=0.00001, type=float)  # 学习率
    parser.add_argument('--vocab_size', '-vs', help='vocab size for discretization', default=256, type=int)  # 离散词典大小
    parser.add_argument('--dataset_episode_num', '-den', help='训练数据量', default=100, type=int)
    parser.add_argument('--loaded_checkpoints_dir', '-lcd', help='模型加载目录', default="~/", type=str)
    parser.add_argument('--save_model', '-sm', help='save model', default=True)
    parser.add_argument('--model_save_epoch', '-mse', help='save model at every num epoch', default=10, type=int)
    parser.add_argument('--checkpoints_saved_dir', '-csd', help='模型保存目录', default="/adddisk1/huangyiyang/model/", type=str)
    args = parser.parse_args()
    return args


time_sequence_length = 6  # 常量，来自论文每次预测使用6张图片


def create_train_dataset(args, global_batch_size):
    '''创建数据集'''
    dataset_dirs = args.dataset_dirs.split("+") # 可以包含多个数据集，输入时用+分割

    workdir = "~/"
    sequence_length = time_sequence_length # 每次图片张数
    data_target_width = 456 # 输入图像的宽度
    data_target_height = 256 # 输入图像的高度
    random_crop_factor = 0.95
    replay_capacity = 5_000
    seed = 42
    rng = jax.random.PRNGKey(seed) #jax可以理解为针对硬件的numpy等的加速包，详见项目pdf
    rng, data_rng = jax.random.split(rng)
    data_rng = jax.random.fold_in(data_rng, jax.process_index())

    '''导入rlds类型的数据集，rlds数据集为google论文中给定的数据格式，详见项目pdf'''
    train_ds = rlds_dataset_loader.create_datasets(
        data_rng,
        dataset_dirs=dataset_dirs,
        sequence_length=sequence_length,
        global_batch_size=global_batch_size,
        target_width=data_target_width,
        target_height=data_target_height,
        random_crop_factor=random_crop_factor,
        cache=False,
        shuffle=True,
        shuffle_buffer_size=replay_capacity,
        cache_dir=workdir,
        dataset_episode_num=args.dataset_episode_num
    )

    return train_ds


def create_model(args):
    '''创建模型'''
    data_target_width = 456
    data_target_height = 256

    state_spec = tensorspec_utils.TensorSpecStruct() #tensorflow的扩展数据类型，详见项目pdf,此处为

    state_spec.image = tensor_spec.BoundedTensorSpec([data_target_height, data_target_width, 3],
                                                     dtype=tf.float32,
                                                     name='image',
                                                     minimum=0.,
                                                     maximum=1.)
    state_spec.natural_language_embedding = tensor_spec.TensorSpec(
        shape=[512], dtype=tf.float32, name='natural_language_embedding')

    action_spec = tensorspec_utils.TensorSpecStruct()

    action_spec.terminate_episode = tensor_spec.BoundedTensorSpec(
        (2,), dtype=tf.int32, minimum=0, maximum=1, name='terminate_episode')

    action_spec.effector_target_translation = tensor_spec.BoundedTensorSpec(
        (2,), dtype=tf.float32, minimum=-0.3, maximum=0.6, name='effector_target_translation')

    network = transformer_network.TransformerNetwork(
        input_tensor_spec=state_spec,
        output_tensor_spec=action_spec,
        vocab_size=int(args.vocab_size),
        token_embedding_size=512,
        num_layers=8,
        layer_size=128,
        num_heads=8,
        feed_forward_size=512,
        dropout_rate=0.1,
        time_sequence_length=time_sequence_length,
        crop_size=236,
        use_token_learner=True,
        action_order=['terminate_episode', 'effector_target_translation'])
    return network

''' 设置 TF_CONFIG 环境变量,为每个集群设置对应的文件'''
def set_env():
    '''设置分布式训练环境变量，详见项目pdf'''
    worker_idx = 0
    swarm = []
    swarm.append("10.20.0.253" + ":" + "10023")
    cluster = {'worker': swarm}
    type = "worker"
    os.environ['TF_CONFIG'] = json.dumps({
        'cluster': cluster,
        'task': {'type': type, 'index': worker_idx}
    })


if __name__ == '__main__':
    os.environ.pop('TF_CONFIG', None) # 清除TF_CONFIG

    args = get_args()

    physical_devices = tf.config.experimental.list_physical_devices('GPU')
    if len(physical_devices) > 0:
        for k in range(len(physical_devices)):
            tf.config.experimental.set_memory_growth(physical_devices[k], True)
    else:
        print("GPU数量不够")
        exit("异常退出")

   # set_env()

   # options = tf.distribute.experimental.CommunicationOptions(
   #     implementation=tf.distribute.experimental.CommunicationImplementation.NCCL
   # )
   # mirrored_strategy = tf.distribute.MultiWorkerMirroredStrategy(
   #     communication_options=options)
    mirrored_strategy = tf.distribute.MirroredStrategy()
    global_batch_size = args.single_gpu_batch_size * mirrored_strategy.num_replicas_in_sync

    global_learning_rate = args.learning_rate * global_batch_size

    with mirrored_strategy.scope():
        print("begin")
        network = create_model(args)
        network.create_variables()
        print("network")
        dataset_dirs = args.dataset_dirs
        train_ds = create_train_dataset(args, global_batch_size)
        print("dataset")
        dist_dataset = mirrored_strategy.experimental_distribute_dataset(train_ds)
        print("dist_dataset")
        network_state = tensor_spec.sample_spec_nest(
            network.state_spec, outer_dims=[args.single_gpu_batch_size])
        optimizer = tf.keras.optimizers.Adam(learning_rate=global_learning_rate)

        ckpt = tf.train.Checkpoint(step=tf.Variable(1), optimizer=optimizer, model=network)
       # if tf.train.latest_checkpoint(args.loaded_checkpoints_dir):
           # ckpt.restore(tf.train.latest_checkpoint(args.loaded_checkpoints_dir))
           # print("从 %s 恢复模型" % (args.loaded_checkpoints_dir))

        current_step = ckpt.step.numpy()
        print("开始训练")
        T1 = time.time()


        @tf.function
        def train_one_step(model, observation_batch, label_batch, network_state, optimizer):
            '''单步训练'''
            with tf.GradientTape() as tape:
                model.set_actions(label_batch)
                model(observation_batch, step_type=None, network_state=network_state, training=True)
                loss = tf.reduce_mean(model.get_actor_loss())
                gradients = tape.gradient(loss, model.trainable_variables,
                                          unconnected_gradients=tf.UnconnectedGradients.ZERO)
                optimizer.apply_gradients(grads_and_vars=zip(gradients, model.trainable_variables))
                logging_info = model.get_aux_info()
                return loss, logging_info


        action_order = network._action_tokenizer.action_order
        for epoch in range(1, args.training_epoch):
            total_loss = 0.0
            step = 0
            T1 = time.time()

            for data in dist_dataset:
                train_observation = data["train_observation"]
                train_labels = data["action_lable"]
                per_replica_losses, logging_info = mirrored_strategy.run(
                    train_one_step, args=(network, train_observation, train_labels, network_state, optimizer))
                step = step + 1
                mean_loss = mirrored_strategy.reduce(tf.distribute.ReduceOp.MEAN, per_replica_losses, axis=None)
                total_loss = total_loss + mean_loss
                ckpt.step.assign_add(1)
                print('训练1个step')
                print(total_loss)

            T2 = time.time()
            print('训练1个epoch 总耗时: ', ((T2 - T1)))

            if epoch % args.model_save_epoch == 0 and args.save_model:
                checkpoint_prefix = os.path.join(args.checkpoints_saved_dir, "ckpt")
                ckpt.save(checkpoint_prefix)
                print("模型保存位置：  %s !" % (checkpoint_prefix))

print("正常退出!")

distribute_train.py
YiyangHuang-work/RT-1
# Copyright 2022 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""ResNet variants model for Keras with Film-Conditioning.

Related papers/blogs:
- https://arxiv.org/abs/1512.03385
- https://arxiv.org/pdf/1603.05027v2.pdf
- http://torch.ch/blog/2016/02/04/resnets.html
- https://arxiv.org/abs/1709.07871
"""
import tensorflow.compat.v2 as tf

layers = tf.keras.layers


class FilmConditioning(tf.keras.layers.Layer):
  """Layer that adds FiLM conditioning.

  This is intended to be applied after a convolutional layer. It will learn a
  multiplicative and an additive factor to be applied to each channel of the
  convolution's output.

  Conv layer can be rank 2 or 4.

  For further details, see: https://arxiv.org/abs/1709.07871
  """

  def __init__(self, num_channels: int):
    """Constructs a FiLM conditioning layer.

    Args:
      num_channels: Number of filter channels to expect in the input.
    """
    super().__init__()
    # Note that we initialize with zeros because empirically we have found
    # this works better than initializing with glorot.
    self._projection_add = layers.Dense(
        num_channels,
        activation=None,
        kernel_initializer='zeros',
        bias_initializer='zeros')
    self._projection_mult = layers.Dense(
        num_channels,
        activation=None,
        kernel_initializer='zeros',
        bias_initializer='zeros')

  def call(self, conv_filters: tf.Tensor, conditioning: tf.Tensor):
    tf.debugging.assert_rank(conditioning, 2)
    projected_cond_add = self._projection_add(conditioning)
    projected_cond_mult = self._projection_mult(conditioning)

    if len(conv_filters.shape) == 4:
      # [B, D] -> [B, 1, 1, D]
      projected_cond_add = projected_cond_add[:, tf.newaxis, tf.newaxis]
      projected_cond_mult = projected_cond_mult[:, tf.newaxis, tf.newaxis]
    else:
      tf.debugging.assert_rank(conv_filters, 2)

    # Original FiLM paper argues that 1 + gamma centers the initialization at
    # identity transform.
    result = (1 + projected_cond_mult) * conv_filters + projected_cond_add
    return result

film_efficientnet/film_conditioning_layer.py
YiyangHuang-work/RT-1
# Copyright 2022 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Tests for film_conditioning_layer."""
from absl.testing import parameterized
import numpy as np
from robotics_transformer.film_efficientnet import film_conditioning_layer
import tensorflow as tf


class FilmConditioningLayerTest(tf.test.TestCase, parameterized.TestCase):

  @parameterized.parameters([2, 4])
  def test_film_conditioning_rank_two_and_four(self, conv_rank):
    batch = 2
    num_channels = 3
    if conv_rank == 2:
      conv_layer = np.random.randn(batch, num_channels)
    elif conv_rank == 4:
      conv_layer = np.random.randn(batch, 1, 1, num_channels)
    else:
      raise ValueError(f'Unexpected conv rank: {conv_rank}')
    context = np.random.rand(batch, num_channels)
    film_layer = film_conditioning_layer.FilmConditioning(num_channels)
    out = film_layer(conv_layer, context)
    tf.debugging.assert_rank(out, conv_rank)


if __name__ == '__main__':
  tf.test.main()

film_efficientnet/film_conditioning_layer_test.py
YiyangHuang-work/RT-1
# Copyright 2022 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# pytype: skip-file
# pylint: skip-file
"""EfficientNet models modified with added film layers.

Mostly copied from third_party/py/keras/applications/efficientnet.py
"""

import copy
import math
import os
import warnings
import json

from absl import logging
import tensorflow.compat.v2 as tf
from tensorflow.keras import layers

from robotics_transformer.film_efficientnet.film_conditioning_layer import FilmConditioning

BASE_WEIGHTS_PATH = 'efficientnet_checkpoints/efficientnet'
IMAGENET_JSON_PATH = 'efficientnet_checkpoints/imagenet_classes.json'
CLASS_INDEX = None

WEIGHTS_PATHS = {
    'efficientnetb3': BASE_WEIGHTS_PATH + 'b3.h5',
    'efficientnetb3_notop': BASE_WEIGHTS_PATH + 'b3_notop.h5',
}

DEFAULT_BLOCKS_ARGS = [{
    'kernel_size': 3,
    'repeats': 1,
    'filters_in': 32,
    'filters_out': 16,
    'expand_ratio': 1,
    'id_skip': True,
    'strides': 1,
    'se_ratio': 0.25
}, {
    'kernel_size': 3,
    'repeats': 2,
    'filters_in': 16,
    'filters_out': 24,
    'expand_ratio': 6,
    'id_skip': True,
    'strides': 2,
    'se_ratio': 0.25
}, {
    'kernel_size': 5,
    'repeats': 2,
    'filters_in': 24,
    'filters_out': 40,
    'expand_ratio': 6,
    'id_skip': True,
    'strides': 2,
    'se_ratio': 0.25
}, {
    'kernel_size': 3,
    'repeats': 3,
    'filters_in': 40,
    'filters_out': 80,
    'expand_ratio': 6,
    'id_skip': True,
    'strides': 2,
    'se_ratio': 0.25
}, {
    'kernel_size': 5,
    'repeats': 3,
    'filters_in': 80,
    'filters_out': 112,
    'expand_ratio': 6,
    'id_skip': True,
    'strides': 1,
    'se_ratio': 0.25
}, {
    'kernel_size': 5,
    'repeats': 4,
    'filters_in': 112,
    'filters_out': 192,
    'expand_ratio': 6,
    'id_skip': True,
    'strides': 2,
    'se_ratio': 0.25
}, {
    'kernel_size': 3,
    'repeats': 1,
    'filters_in': 192,
    'filters_out': 320,
    'expand_ratio': 6,
    'id_skip': True,
    'strides': 1,
    'se_ratio': 0.25
}]

CONV_KERNEL_INITIALIZER = {
    'class_name': 'VarianceScaling',
    'config': {
        'scale': 2.0,
        'mode': 'fan_out',
        'distribution': 'truncated_normal'
    }
}

DENSE_KERNEL_INITIALIZER = {
    'class_name': 'VarianceScaling',
    'config': {
        'scale': 1. / 3.,
        'mode': 'fan_out',
        'distribution': 'uniform'
    }
}

BASE_DOCSTRING = """Instantiates the {name} architecture.

  Reference:
  - [EfficientNet: Rethinking Model Scaling for Convolutional Neural Networks](
      https://arxiv.org/abs/1905.11946) (ICML 2019)

  This function returns a Keras image classification model,
  optionally loaded with weights pre-trained on ImageNet.

  For image classification use cases, see
  [this page for detailed examples](
    https://keras.io/api/applications/#usage-examples-for-image-classification-models).

  For transfer learning use cases, make sure to read the
  [guide to transfer learning & fine-tuning](
    https://keras.io/guides/transfer_learning/).

  Note: each Keras Application expects a specific kind of input preprocessing.
  For EfficientNet, input preprocessing is included as part of the model
  (as a `Rescaling` layer), and thus
  `tf.keras.applications.efficientnet.preprocess_input` is actually a
  pass-through function. EfficientNet models expect their inputs to be float
  tensors of pixels with values in the [0-255] range.

  Args:
    include_top: Whether to include the fully-connected
        layer at the top of the network. Defaults to True.
    weights: One of `None` (random initialization),
          'imagenet' (pre-training on ImageNet),
          or the path to the weights file to be loaded. Defaults to 'imagenet'.
    input_tensor: Optional Keras tensor
        (i.e. output of `layers.Input()`)
        to use as image input for the model.
    input_shape: Optional shape tuple, only to be specified
        if `include_top` is False.
        It should have exactly 3 inputs channels.
    pooling: Optional pooling mode for feature extraction
        when `include_top` is `False`. Defaults to None.
        - `None` means that the output of the model will be
            the 4D tensor output of the
            last convolutional layer.
        - `avg` means that global average pooling
            will be applied to the output of the
            last convolutional layer, and thus
            the output of the model will be a 2D tensor.
        - `max` means that global max pooling will
            be applied.
    classes: Optional number of classes to classify images
        into, only to be specified if `include_top` is True, and
        if no `weights` argument is specified. Defaults to 1000 (number of
        ImageNet classes).
    classifier_activation: A `str` or callable. The activation function to use
        on the "top" layer. Ignored unless `include_top=True`. Set
        `classifier_activation=None` to return the logits of the "top" layer.
        Defaults to 'softmax'.
        When loading pretrained weights, `classifier_activation` can only
        be `None` or `"softmax"`.

  Returns:
    A `keras.Model` instance.
"""

IMAGENET_STDDEV_RGB = [0.229, 0.224, 0.225]


def validate_activation(classifier_activation, weights):
  """validates that the classifier is compatible with the weights.

  Args:
    classifier_activation: str or callable activation function
    weights: The pretrained weights to load.

  Raises:
    ValueError: if an activation other than `None` or `softmax` are used with
      pretrained weights.
  """
  if weights is None:
    return

  classifier_activation = tf.keras.activations.get(classifier_activation)
  if classifier_activation not in {
      tf.keras.activations.get('softmax'),
      tf.keras.activations.get(None)
  }:
    raise ValueError('Only `None` and `softmax` activations are allowed '
                     'for the `classifier_activation` argument when using '
                     'pretrained weights, with `include_top=True`; Received: '
                     f'classifier_activation={classifier_activation}')


def correct_pad(inputs, kernel_size):
  """Returns a tuple for zero-padding for 2D convolution with downsampling.

  Args:
    inputs: Input tensor.
    kernel_size: An integer or tuple/list of 2 integers.

  Returns:
    A tuple.
  """
  img_dim = 2 if tf.keras.backend.image_data_format() == 'channels_first' else 1
  input_size = tf.keras.backend.int_shape(inputs)[img_dim:(img_dim + 2)]
  if isinstance(kernel_size, int):
    kernel_size = (kernel_size, kernel_size)
  if input_size[0] is None:
    adjust = (1, 1)
  else:
    adjust = (1 - input_size[0] % 2, 1 - input_size[1] % 2)
  correct = (kernel_size[0] // 2, kernel_size[1] // 2)
  return ((correct[0] - adjust[0], correct[0]), (correct[1] - adjust[1],
                                                 correct[1]))


def obtain_input_shape(input_shape,
                       default_size,
                       min_size,
                       data_format,
                       require_flatten,
                       weights=None):
  """Internal utility to compute/validate a model's input shape.

  Args:
    input_shape: Either None (will return the default network input shape), or a
      user-provided shape to be validated.
    default_size: Default input width/height for the model.
    min_size: Minimum input width/height accepted by the model.
    data_format: Image data format to use.
    require_flatten: Whether the model is expected to be linked to a classifier
      via a Flatten layer.
    weights: One of `None` (random initialization) or 'imagenet' (pre-training
      on ImageNet). If weights='imagenet' input channels must be equal to 3.

  Returns:
    An integer shape tuple (may include None entries).

  Raises:
    ValueError: In case of invalid argument values.
  """
  if weights != 'imagenet' and input_shape and len(input_shape) == 3:
    if data_format == 'channels_first':
      if input_shape[0] not in {1, 3}:
        warnings.warn(
            'This model usually expects 1 or 3 input channels. '
            'However, it was passed an input_shape with ' +
            str(input_shape[0]) + ' input channels.',
            stacklevel=2)
      default_shape = (input_shape[0], default_size, default_size)
    else:
      if input_shape[-1] not in {1, 3}:
        warnings.warn(
            'This model usually expects 1 or 3 input channels. '
            'However, it was passed an input_shape with ' +
            str(input_shape[-1]) + ' input channels.',
            stacklevel=2)
      default_shape = (default_size, default_size, input_shape[-1])
  else:
    if data_format == 'channels_first':
      default_shape = (3, default_size, default_size)
    else:
      default_shape = (default_size, default_size, 3)
  if weights == 'imagenet' and require_flatten:
    if input_shape is not None:
      if input_shape != default_shape:
        raise ValueError('When setting `include_top=True` '
                         'and loading `imagenet` weights, '
                         f'`input_shape` should be {default_shape}.  '
                         f'Received: input_shape={input_shape}')
    return default_shape
  if input_shape:
    if data_format == 'channels_first':
      if input_shape is not None:
        if len(input_shape) != 3:
          raise ValueError('`input_shape` must be a tuple of three integers.')
        if input_shape[0] != 3 and weights == 'imagenet':
          raise ValueError('The input must have 3 channels; Received '
                           f'`input_shape={input_shape}`')
        if ((input_shape[1] is not None and input_shape[1] < min_size) or
            (input_shape[2] is not None and input_shape[2] < min_size)):
          raise ValueError(f'Input size must be at least {min_size}'
                           f'x{min_size}; Received: '
                           f'input_shape={input_shape}')
    else:
      if input_shape is not None:
        if len(input_shape) != 3:
          raise ValueError('`input_shape` must be a tuple of three integers.')
        if input_shape[-1] != 3 and weights == 'imagenet':
          raise ValueError('The input must have 3 channels; Received '
                           f'`input_shape={input_shape}`')
        if ((input_shape[0] is not None and input_shape[0] < min_size) or
            (input_shape[1] is not None and input_shape[1] < min_size)):
          raise ValueError('Input size must be at least '
                           f'{min_size}x{min_size}; Received: '
                           f'input_shape={input_shape}')
  else:
    if require_flatten:
      input_shape = default_shape
    else:
      if data_format == 'channels_first':
        input_shape = (3, None, None)
      else:
        input_shape = (None, None, 3)
  if require_flatten:
    if None in input_shape:
      raise ValueError('If `include_top` is True, '
                       'you should specify a static `input_shape`. '
                       f'Received: input_shape={input_shape}')
  return input_shape


def EfficientNet(width_coefficient,
                 depth_coefficient,
                 default_size,
                 dropout_rate=0.2,
                 drop_connect_rate=0.2,
                 depth_divisor=8,
                 activation='swish',
                 blocks_args='default',
                 model_name='efficientnet',
                 include_top=True,
                 weights='imagenet',
                 input_tensor=None,
                 input_shape=None,
                 pooling=None,
                 classes=1000,
                 classifier_activation='softmax',
                 include_film=False):
  """Instantiates the EfficientNet architecture using given scaling coefficients.

  Args:
    width_coefficient: float, scaling coefficient for network width.
    depth_coefficient: float, scaling coefficient for network depth.
    default_size: integer, default input image size.
    dropout_rate: float, dropout rate before final classifier layer.
    drop_connect_rate: float, dropout rate at skip connections.
    depth_divisor: integer, a unit of network width.
    activation: activation function.
    blocks_args: list of dicts, parameters to construct block modules.
    model_name: string, model name.
    include_top: whether to include the fully-connected layer at the top of the
      network.
    weights: one of `None` (random initialization), 'imagenet' (pre-training on
      ImageNet), or the path to the weights file to be loaded.
    input_tensor: optional Keras tensor (i.e. output of `layers.Input()`) to use
      as image input for the model.
    input_shape: optional shape tuple, only to be specified if `include_top` is
      False. It should have exactly 3 inputs channels.
    pooling: optional pooling mode for feature extraction when `include_top` is
      `False`. - `None` means that the output of the model will be the 4D tensor
      output of the last convolutional layer. - `avg` means that global average
      pooling will be applied to the output of the last convolutional layer, and
      thus the output of the model will be a 2D tensor. - `max` means that
      global max pooling will be applied.
    classes: optional number of classes to classify images into, only to be
      specified if `include_top` is True, and if no `weights` argument is
      specified.
    classifier_activation: A `str` or callable. The activation function to use
      on the "top" layer. Ignored unless `include_top=True`. Set
      `classifier_activation=None` to return the logits of the "top" layer.
    include_film: bool, whether or not to insert film conditioning layers.

  Returns:
    A `keras.Model` instance.

  Raises:
    ValueError: in case of invalid argument for `weights`,
      or invalid input shape.
    ValueError: if `classifier_activation` is not `softmax` or `None` when
      using a pretrained top layer.
  """
  if blocks_args == 'default':
    blocks_args = DEFAULT_BLOCKS_ARGS

  if not (weights in {'imagenet', None} or tf.io.gfile.exists(weights)):
    raise ValueError('The `weights` argument should be either '
                     '`None` (random initialization), `imagenet` '
                     '(pre-training on ImageNet), '
                     'or the path to the weights file to be loaded.')

  if weights == 'imagenet' and include_top and classes != 1000:
    raise ValueError('If using `weights` as `"imagenet"` with `include_top`'
                     ' as true, `classes` should be 1000')

  # Determine proper input shape
  input_shape = obtain_input_shape(
      input_shape,
      default_size=default_size,
      min_size=32,
      data_format=tf.keras.backend.image_data_format(),
      require_flatten=include_top,
      weights=weights)

  if include_film:
    with tf.compat.v1.variable_scope('context_input'):
      context_input = layers.Input(shape=512)
  if input_tensor is None:
    img_input = layers.Input(shape=input_shape)
  else:
    if not tf.keras.backend.is_keras_tensor(input_tensor):
      img_input = layers.Input(tensor=input_tensor, shape=input_shape)
    else:
      img_input = input_tensor

  bn_axis = 3 if tf.keras.backend.image_data_format() == 'channels_last' else 1

  def round_filters(filters, divisor=depth_divisor):
    """Round number of filters based on depth multiplier."""
    filters *= width_coefficient
    new_filters = max(divisor, int(filters + divisor / 2) // divisor * divisor)
    # Make sure that round down does not go down by more than 10%.
    if new_filters < 0.9 * filters:
      new_filters += divisor
    return int(new_filters)

  def round_repeats(repeats):
    """Round number of repeats based on depth multiplier."""
    return int(math.ceil(depth_coefficient * repeats))

  # Build stem
  x = img_input
  x = layers.Rescaling(1. / 255.)(x)
  x = layers.Normalization(axis=bn_axis)(x)
  # Note that the normaliztion layer uses square value of STDDEV as the
  # variance for the layer: result = (input - mean) / sqrt(var)
  # However, the original implemenetation uses (input - mean) / var to
  # normalize the input, we need to divide another sqrt(var) to match the
  # original implementation.
  # See https://github.com/tensorflow/tensorflow/issues/49930 for more details
  # We always apply this transformation, even when not using imagenet weights,
  # because it needs to be in the graph when grafting weights from imagenet
  # pretrained models.
  x = layers.Rescaling(1. / tf.math.sqrt(IMAGENET_STDDEV_RGB))(x)

  x = layers.ZeroPadding2D(padding=correct_pad(x, 3), name='stem_conv_pad')(x)
  x = layers.Conv2D(
      round_filters(32),
      3,
      strides=2,
      padding='valid',
      use_bias=False,
      kernel_initializer=CONV_KERNEL_INITIALIZER,
      name='stem_conv')(
          x)
  x = layers.BatchNormalization(axis=bn_axis, name='stem_bn')(x)
  x = layers.Activation(activation, name='stem_activation')(x)

  # Build blocks
  blocks_args = copy.deepcopy(blocks_args)

  b = 0
  blocks = float(sum(round_repeats(args['repeats']) for args in blocks_args))
  for (i, args) in enumerate(blocks_args):
    assert args['repeats'] > 0
    # Update block input and output filters based on depth multiplier.
    args['filters_in'] = round_filters(args['filters_in'])
    args['filters_out'] = round_filters(args['filters_out'])

    for j in range(round_repeats(args.pop('repeats'))):
      # The first block needs to take care of stride and filter size increase.
      if j > 0:
        args['strides'] = 1
        args['filters_in'] = args['filters_out']
      x = block(
          x,
          activation,
          drop_connect_rate * b / blocks,
          name='block{}{}_'.format(i + 1, chr(j + 97)),
          **args)
      if include_film:
        with tf.compat.v1.variable_scope('film_conditioning'):
          x = FilmConditioning(num_channels=x.shape[-1])(x, context_input)
      b += 1

  # Build top
  x = layers.Conv2D(
      round_filters(1280),
      1,
      padding='same',
      use_bias=False,
      kernel_initializer=CONV_KERNEL_INITIALIZER,
      name='top_conv')(
          x)
  x = layers.BatchNormalization(axis=bn_axis, name='top_bn')(x)
  x = layers.Activation(activation, name='top_activation')(x)
  if include_top:
    x = layers.GlobalAveragePooling2D(name='avg_pool')(x)
    if dropout_rate > 0:
      x = layers.Dropout(dropout_rate, name='top_dropout')(x)
    validate_activation(classifier_activation, weights)
    x = layers.Dense(
        classes,
        activation=classifier_activation,
        kernel_initializer=DENSE_KERNEL_INITIALIZER,
        name='predictions')(
            x)
  else:
    if pooling == 'avg':
      x = layers.GlobalAveragePooling2D(name='avg_pool')(x)
    elif pooling == 'max':
      x = layers.GlobalMaxPooling2D(name='max_pool')(x)

  # Ensure that the model takes into account
  # any potential predecessors of `input_tensor`.
  if input_tensor is not None:
    inputs = tf.keras.utils.get_source_inputs(input_tensor)
  else:
    inputs = img_input
    if include_film:
      inputs = (img_input, context_input)

  # Create model.
  model = tf.keras.Model(inputs, x, name=model_name)

  # Load weights.
  if weights == 'imagenet':
    if include_top:
      key = model_name
    else:
      key = model_name + '_notop'
    weights_path = os.path.join(os.path.dirname(__file__), WEIGHTS_PATHS[key])
    model.load_weights(weights_path, skip_mismatch=False, by_name=False)
  elif weights is not None:
    model.load_weights(weights, skip_mismatch=False, by_name=False)
  return model


def block(inputs,
          activation='swish',
          drop_rate=0.,
          name='',
          filters_in=32,
          filters_out=16,
          kernel_size=3,
          strides=1,
          expand_ratio=1,
          se_ratio=0.,
          id_skip=True):
  """An inverted residual block.

  Args:
      inputs: input tensor.
      activation: activation function.
      drop_rate: float between 0 and 1, fraction of the input units to drop.
      name: string, block label.
      filters_in: integer, the number of input filters.
      filters_out: integer, the number of output filters.
      kernel_size: integer, the dimension of the convolution window.
      strides: integer, the stride of the convolution.
      expand_ratio: integer, scaling coefficient for the input filters.
      se_ratio: float between 0 and 1, fraction to squeeze the input filters.
      id_skip: boolean.

  Returns:
      output tensor for the block.
  """
  bn_axis = 3 if tf.keras.backend.image_data_format() == 'channels_last' else 1

  # Expansion phase
  filters = filters_in * expand_ratio
  if expand_ratio != 1:
    x = layers.Conv2D(
        filters,
        1,
        padding='same',
        use_bias=False,
        kernel_initializer=CONV_KERNEL_INITIALIZER,
        name=name + 'expand_conv')(
            inputs)
    x = layers.BatchNormalization(axis=bn_axis, name=name + 'expand_bn')(x)
    x = layers.Activation(activation, name=name + 'expand_activation')(x)
  else:
    x = inputs

  # Depthwise Convolution
  if strides == 2:
    x = layers.ZeroPadding2D(
        padding=correct_pad(x, kernel_size), name=name + 'dwconv_pad')(
            x)
    conv_pad = 'valid'
  else:
    conv_pad = 'same'
  x = layers.DepthwiseConv2D(
      kernel_size,
      strides=strides,
      padding=conv_pad,
      use_bias=False,
      depthwise_initializer=CONV_KERNEL_INITIALIZER,
      name=name + 'dwconv')(
          x)
  x = layers.BatchNormalization(axis=bn_axis, name=name + 'bn')(x)
  x = layers.Activation(activation, name=name + 'activation')(x)

  # Squeeze and Excitation phase
  if 0 < se_ratio <= 1:
    filters_se = max(1, int(filters_in * se_ratio))
    se = layers.GlobalAveragePooling2D(name=name + 'se_squeeze')(x)
    if bn_axis == 1:
      se_shape = (filters, 1, 1)
    else:
      se_shape = (1, 1, filters)
    se = layers.Reshape(se_shape, name=name + 'se_reshape')(se)
    se = layers.Conv2D(
        filters_se,
        1,
        padding='same',
        activation=activation,
        kernel_initializer=CONV_KERNEL_INITIALIZER,
        name=name + 'se_reduce')(
            se)
    se = layers.Conv2D(
        filters,
        1,
        padding='same',
        activation='sigmoid',
        kernel_initializer=CONV_KERNEL_INITIALIZER,
        name=name + 'se_expand')(
            se)
    x = layers.multiply([x, se], name=name + 'se_excite')

  # Output phase
  x = layers.Conv2D(
      filters_out,
      1,
      padding='same',
      use_bias=False,
      kernel_initializer=CONV_KERNEL_INITIALIZER,
      name=name + 'project_conv')(
          x)
  x = layers.BatchNormalization(axis=bn_axis, name=name + 'project_bn')(x)
  if id_skip and strides == 1 and filters_in == filters_out:
    if drop_rate > 0:
      x = layers.Dropout(
          drop_rate, noise_shape=(None, 1, 1, 1), name=name + 'drop')(
              x)
    x = layers.add([x, inputs], name=name + 'add')
  return x


def maybe_restore_with_film(
    *args,
    weights='imagenet',
    include_film=False,
    **kwargs,
):
  n1 = EfficientNet(*args, weights=weights, include_film=False, **kwargs)
  if not include_film:
    return n1
  # Copy the model weights over to a new model. This is necessary
  # in case we have inserted early film layers. In this case,
  # the pretrained weights will fail to restore properly
  # unless we do this trick.
  n2 = EfficientNet(*args, weights=None, include_film=True, **kwargs)
  # The layers without the film layers.
  l1 = {l.name: l for l in n1.layers}
  # The layers with the film layers.
  l2 = {l.name: l for l in n2.layers}
  for layer_name, layer in l2.items():
    if layer_name in l1:
      layer.set_weights(l1[layer_name].get_weights())
    # Annoyingly, the rescaling and normalization layers get different names
    # in each graph.
    elif 'rescaling' in layer_name:
      _, num = layer_name.split('_')
      l1_layer_name = 'rescaling_' + str(int(num) - 2 or '')
      l1_layer_name = l1_layer_name.rstrip('_')
      layer.set_weights(l1[l1_layer_name].get_weights())
    elif 'normalization' in layer_name:
      _, num = layer_name.split('_')
      l1_layer_name = 'normalization_' + str(int(num) - 1 or '')
      l1_layer_name = l1_layer_name.rstrip('_')
      layer.set_weights(l1[l1_layer_name].get_weights())
  return n2


def EfficientNetB3(include_top=True,
                   weights='imagenet',
                   input_tensor=None,
                   input_shape=None,
                   pooling=None,
                   classes=1000,
                   classifier_activation='softmax',
                   include_film=False,
                   **kwargs):
  return maybe_restore_with_film(
      1.2,
      1.4,
      300,
      0.3,
      model_name='efficientnetb3',
      include_top=include_top,
      weights=weights,
      input_tensor=input_tensor,
      input_shape=input_shape,
      pooling=pooling,
      classes=classes,
      classifier_activation=classifier_activation,
      include_film=include_film,
      **kwargs)


EfficientNetB3.__doc__ = BASE_DOCSTRING.format(name='EfficientNetB3')


def preprocess_input(x, data_format=None):  # pylint: disable=unused-argument
  """A placeholder method for backward compatibility.

  The preprocessing logic has been included in the efficientnet model
  implementation. Users are no longer required to call this method to normalize
  the input data. This method does nothing and only kept as a placeholder to
  align the API surface between old and new version of model.

  Args:
    x: A floating point `numpy.array` or a `tf.Tensor`.
    data_format: Optional data format of the image tensor/array. Defaults to
      None, in which case the global setting `tf.keras.image_data_format() is
      used (unless you changed it, it defaults to "channels_last").{mode}

  Returns:
    Unchanged `numpy.array` or `tf.Tensor`.
  """
  return x


def decode_predictions(preds, top=5):
  global CLASS_INDEX
  if CLASS_INDEX is None:
    with open(os.path.join(os.path.dirname(__file__), IMAGENET_JSON_PATH)) as f:
      CLASS_INDEX = json.load(f)
  results = []
  for pred in preds:
    top_indices = pred.argsort()[-top:][::-1]
    result = [tuple(CLASS_INDEX[str(i)]) + (pred[i],) for i in top_indices]
    result.sort(key=lambda x: x[2], reverse=True)
    results.append(result)
  return results

film_efficientnet/film_efficientnet_encoder.py
YiyangHuang-work/RT-1
# Copyright 2022 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Tests that film_efficientnet can detect an image of a cat."""

from absl.testing import parameterized
import numpy as np
from robotics_transformer.film_efficientnet import film_efficientnet_encoder
from skimage import data
import tensorflow as tf


class FilmEfficientnetTest(tf.test.TestCase, parameterized.TestCase):

  def _helper(self, include_film, model_variant):
    if model_variant == 'b0':
      size = 224
      fe = film_efficientnet_encoder.EfficientNetB0
    elif model_variant == 'b1':
      size = 240
      fe = film_efficientnet_encoder.EfficientNetB1
    elif model_variant == 'b2':
      size = 260
      fe = film_efficientnet_encoder.EfficientNetB2
    elif model_variant == 'b3':
      size = 300
      fe = film_efficientnet_encoder.EfficientNetB3
    elif model_variant == 'b4':
      size = 380
      fe = film_efficientnet_encoder.EfficientNetB4
    elif model_variant == 'b5':
      size = 456
      fe = film_efficientnet_encoder.EfficientNetB5
    elif model_variant == 'b6':
      size = 528
      fe = film_efficientnet_encoder.EfficientNetB6
    elif model_variant == 'b7':
      size = 600
      fe = film_efficientnet_encoder.EfficientNetB7
    else:
      raise ValueError(f'Unknown variant: {model_variant}')
    fe = fe(include_top=True, weights='imagenet', include_film=include_film)
    image = np.expand_dims(data.chelsea(), axis=0)
    image = tf.image.resize(image, (size, size))
    context = np.random.randn(1, 512)
    if include_film:
      eff_output = fe(
          (film_efficientnet_encoder.preprocess_input(image), context),
          training=False)
    else:
      eff_output = fe(
          film_efficientnet_encoder.preprocess_input(image), training=False)
    film_preds = film_efficientnet_encoder.decode_predictions(
        eff_output.numpy(), top=10)
    self.assertIn('tabby', [f[1] for f in film_preds[0]])

  @parameterized.parameters([True, False])
  def test_keras_equivalence_b3(self, include_film):
    self._helper(include_film, 'b3')


if __name__ == '__main__':
  tf.test.main()

film_efficientnet/film_efficientnet_encoder_test.py
YiyangHuang-work/RT-1
# Copyright 2022 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Preprocessing functions for transforming the image for training."""

from typing import Optional

import gin
import tensorflow.compat.v2 as tf

CROP_SIZE = 472


@gin.configurable(
    denylist=['images', 'crop_size', 'training', 'convert_dtype', 'seed'])
def convert_dtype_and_crop_images(images,
                                  crop_size: int = CROP_SIZE,
                                  training: bool = True,
                                  pad_then_crop: bool = False,
                                  convert_dtype: bool = True,
                                  seed: Optional[tf.Tensor] = None):
  """Convert uint8 [512, 640, 3] images to float32 and square crop.

  Args:
    images: [B, H, W, 3] uint8 tensor of images.
    crop_size: Width of the square crop.
    training: If we are in training (random crop) or not-training (fixed crop).
    pad_then_crop: If True, pads image and then crops the original image size.
      This allows full field of view to be extracted.
    convert_dtype: whether or not to convert the image to float32 in the range
      of (0, 1).
    seed: Optional seed of shape (2,) for giving to tf.random.stateless_uniform

  Returns:
    [B, crop_size, crop_size, 3] images of dtype float32.
  """

  if seed is None:
    seed = tf.random.uniform(shape=(2,), maxval=2**30, dtype=tf.int32)

  seed2 = tf.random.experimental.stateless_split(seed, num=1)[0]

  if convert_dtype:
    images = tf.image.convert_image_dtype(images, tf.float32)
  image_height = images.get_shape().as_list()[-3]
  image_width = images.get_shape().as_list()[-2]

  if pad_then_crop:

    if training:
      if image_height == 512:
        ud_pad = 40
        lr_pad = 100
      elif image_height == 256:
        ud_pad = 20
        lr_pad = 50
      else:
        raise ValueError(
            'convert_dtype_and_crop_images only supports image height 512 or '
            '256.')
      max_y = 2 * ud_pad
      max_x = 2 * lr_pad
      images = tf.image.pad_to_bounding_box(
          images,
          offset_height=ud_pad,
          offset_width=lr_pad,
          target_height=image_height + 2 * ud_pad,
          target_width=image_width + 2 * lr_pad)
      offset_y = tf.random.stateless_uniform((),
                                             maxval=max_y + 1,
                                             dtype=tf.int32,
                                             seed=seed)
      offset_x = tf.random.stateless_uniform((),
                                             maxval=max_x + 1,
                                             dtype=tf.int32,
                                             seed=seed2)
      images = tf.image.crop_to_bounding_box(images, offset_y, offset_x,
                                             image_height, image_width)
  else:
    # Standard cropping.
    max_y = image_height - crop_size
    max_x = image_width - crop_size

    if training:
      offset_y = tf.random.stateless_uniform((),
                                             maxval=max_y + 1,
                                             dtype=tf.int32,
                                             seed=seed)
      offset_x = tf.random.stateless_uniform((),
                                             maxval=max_x + 1,
                                             dtype=tf.int32,
                                             seed=seed2)
      images = tf.image.crop_to_bounding_box(images, offset_y, offset_x,
                                             crop_size, crop_size)
    else:
      images = tf.image.crop_to_bounding_box(images, max_y // 2, max_x // 2,
                                             crop_size, crop_size)
  return images

film_efficientnet/preprocessors.py
YiyangHuang-work/RT-1
# Copyright 2022 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Tests for preprocessors."""
from absl.testing import parameterized
import numpy as np
from robotics_transformer.film_efficientnet import preprocessors
from tensor2robot.utils import tensorspec_utils
import tensorflow.compat.v2 as tf


def _random_image(shape):
  images = tf.random.uniform(
      shape, minval=0, maxval=255, dtype=tf.dtypes.int32, seed=42)
  return tf.cast(images, tf.uint8)


def _get_features(
    image_shape=(2, 512, 640, 3), use_task_image=False, use_goal_image=False):
  # Time-dimension stacking occurs during training but not eval.
  state = tensorspec_utils.TensorSpecStruct(image=_random_image(image_shape))
  if use_task_image:
    state.task_image = _random_image(image_shape)
  if use_goal_image:
    state.goal_image = _random_image(image_shape)
  return state


class PreprocessorsTest(tf.test.TestCase, parameterized.TestCase):

  @parameterized.parameters((True, False, False), (False, True, False),
                            (True, False, True), (False, True, True))
  def testConvertDtypeAndCropImages(self, training, pad_then_crop,
                                    convert_dtype):
    features = _get_features()
    images = preprocessors.convert_dtype_and_crop_images(
        features.image,
        training=training,
        pad_then_crop=pad_then_crop,
        convert_dtype=convert_dtype)
    expected_cropped_shape = ([2, 512, 640, 3]
                              if pad_then_crop else [2, 472, 472, 3])
    tf.ensure_shape(images, expected_cropped_shape)
    if convert_dtype:
      self.assertEqual(images.dtype, tf.float32)
      self.assertLessEqual(images.numpy().max(), 1.)
      self.assertGreaterEqual(images.numpy().min(), 0.)
    else:
      self.assertEqual(images.dtype, tf.uint8)
      self.assertLessEqual(images.numpy().max(), 255)
      self.assertGreaterEqual(images.numpy().min(), 0)
      self.assertGreater(images.numpy().max(), 1)

  def testConvertDtypeAndCropImagesSeeded(self):
    features = _get_features()
    seed = tf.constant([1, 2], tf.int32)
    images1 = preprocessors.convert_dtype_and_crop_images(
        features.image, training=True, pad_then_crop=True, seed=seed)
    images2 = preprocessors.convert_dtype_and_crop_images(
        features.image, training=True, pad_then_crop=True, seed=seed)
    diff = np.sum(np.abs(images1.numpy() - images2.numpy()))
    self.assertAlmostEqual(diff, 0)

  def testConvertDtypeAndCropImagesUnseeded(self):
    features = _get_features()
    seed1 = tf.constant([1, 2], tf.int32)
    images1 = preprocessors.convert_dtype_and_crop_images(
        features.image, training=True, pad_then_crop=True, seed=seed1)
    seed2 = tf.constant([2, 3], tf.int32)
    images2 = preprocessors.convert_dtype_and_crop_images(
        features.image, training=True, pad_then_crop=True, seed=seed2)
    diff = np.sum(np.abs(images1.numpy() - images2.numpy()))
    self.assertNotAlmostEqual(diff, 0)

film_efficientnet/preprocessors_test.py
YiyangHuang-work/RT-1
# Copyright 2022 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Encoder based on Efficientnet."""

from typing import Optional

import gin
from robotics_transformer.film_efficientnet import film_conditioning_layer
from robotics_transformer.film_efficientnet import film_efficientnet_encoder
import tensorflow as tf

_MODELS = {
    'b3': film_efficientnet_encoder.EfficientNetB3,
}

_SIZES = {
    'b3': 300,
}


@gin.configurable
class EfficientNetEncoder(tf.keras.layers.Layer):
  """Applies a pretrained Efficientnet based encoder."""

  def __init__(self,
               model_variant: str = 'b3',
               freeze: bool = False,
               early_film: bool = True,
               weights: Optional[str] = 'imagenet',
               include_top: bool = False,
               pooling: bool = True,
               **kwargs):
    """Initialize the model.

    Args:
      model_variant: One of 'b0-b7' of the efficient encoders. See
        https://arxiv.org/abs/1905.11946 to understand the variants.
      freeze: Whether or not to freeze the pretrained weights (seems to not work
        well).
      early_film: Whether to inject film layers into the efficientnet encoder
        (seems to be essential to getting strong performance).
      weights: Which pretrained weights to use. Either 'imagenet', a path to the
        pretrained weights, or None for from scratch.
      include_top: Whether to add the top fully connected layer. If True, this
        will cause encoding to fail and is used only for unit testing purposes.
      pooling: If false, returns feature map before global average pooling
      **kwargs: Keras specific layer kwargs.
    """
    super(EfficientNetEncoder, self).__init__(**kwargs)
    if model_variant not in _MODELS:
      raise ValueError(f'Unknown variant {model_variant}')
    self.model_variant = model_variant
    self.early_film = early_film
    self.freeze = freeze
    self.conv1x1 = tf.keras.layers.Conv2D(
        filters=512,
        kernel_size=(1, 1),
        strides=(1, 1),
        padding='SAME',
        use_bias=False,
        kernel_initializer=tf.keras.initializers.VarianceScaling())
    self.net = _MODELS[model_variant](
        include_top=include_top,
        weights=weights,
        include_film=early_film,
    )
    self.film_layer = film_conditioning_layer.FilmConditioning(num_channels=512)
    self._pooling = pooling

  def _prepare_image(self, image: tf.Tensor) -> tf.Tensor:
    """Resize the input image and check that the range is correct."""
    if len(image.shape) != 4 or image.shape[-1] != 3:
      raise ValueError('Provided image should have shape (b, h, w, 3).')
    size = _SIZES[self.model_variant]
    if image.shape[1] < size / 4 or image.shape[2] < size / 4:
      raise ValueError('Provided image is too small.')
    if image.shape[1] > size * 4 or image.shape[2] > size * 4:
      raise ValueError('Provided image is too large.')
    image = tf.image.resize(image, (size, size))
    c1 = tf.Assert(tf.reduce_max(image) <= 1, data=[tf.reduce_max(image)])
    c2 = tf.Assert(tf.reduce_min(image) >= 0, data=[tf.reduce_min(image)])
    with tf.control_dependencies([c1, c2]):
      image *= 255  # The image is expected to be in range(0, 255).
      image = film_efficientnet_encoder.preprocess_input(image)
      return image

  def _encode(self, image: tf.Tensor, context: tf.Tensor,
              training: bool) -> tf.Tensor:
    """Run the image through the efficientnet encoder."""
    image = self._prepare_image(image)
    if self.early_film:
      return self.net((image, context), training=training)
    return self.net(image, training=training)

  def call(self,
           image: tf.Tensor,
           context: Optional[tf.Tensor] = None,
           training: bool = True) -> tf.Tensor:
    if self.freeze:
      features = tf.stop_gradient(self._encode(image, context, training))
    else:
      features = self._encode(image, context, training)
    if context is not None:
      features = self.conv1x1(features)
      features = self.film_layer(features, context)

    if not self._pooling:
      return features

    # Global average pool.
    return tf.reduce_mean(features, [1, 2])

film_efficientnet/pretrained_efficientnet_encoder.py
YiyangHuang-work/RT-1
# Copyright 2022 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Tests for pretrained_efficientnet_encoder."""

import numpy as np
from robotics_transformer.film_efficientnet import film_efficientnet_encoder
from robotics_transformer.film_efficientnet import pretrained_efficientnet_encoder as eff
from skimage import data
import tensorflow as tf


class PretrainedEfficientnetEncoderTest(tf.test.TestCase):

  def test_encoding(self):
    """Test that we get a correctly shaped decoding."""
    state = np.random.RandomState(0)
    context = state.uniform(-1, 1, (10, 512))
    model = eff.EfficientNetEncoder()
    image = np.expand_dims(data.chelsea(), axis=0) / 255
    preds = model(image, context, training=False).numpy()
    self.assertEqual(preds.shape, (10, 512))

  def test_imagenet_classification(self):
    """Test that we can correctly classify an image of a cat."""
    state = np.random.RandomState(0)
    context = state.uniform(-1, 1, (10, 512))
    model = eff.EfficientNetEncoder(include_top=True)
    image = np.expand_dims(data.chelsea(), axis=0) / 255
    preds = model._encode(image, context, training=False).numpy()
    predicted_names = [
        n[1]
        for n in film_efficientnet_encoder.decode_predictions(preds, top=3)[0]
    ]
    self.assertIn('tabby', predicted_names)


if __name__ == '__main__':
  tf.test.main()

film_efficientnet/pretrained_efficientnet_encoder_test.py
YiyangHuang-work/RT-1
# Copyright 2022 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Sequence policy and agent that directly output actions via actor network.

These classes are not intended to change as they are generic enough for any
all-neural actor based agent+policy. All new features are intended to be
implemented in `actor_network` and `loss_fn`.
"""
from typing import Optional, Type

from absl import logging
import tensorflow as tf
from tf_agents.agents import data_converter
from tf_agents.agents import tf_agent
from tf_agents.networks import network
from tf_agents.policies import actor_policy
from tf_agents.trajectories import policy_step
from tf_agents.trajectories import time_step as ts
from tf_agents.typing import types
from tf_agents.utils import nest_utils


class SequencePolicy(actor_policy.ActorPolicy):
  """A policy that directly outputs actions via an actor network."""

  def __init__(self, **kwargs):
    self._actions = None
    super().__init__(**kwargs)

  def set_actions(self, actions):
    self._actor_network.set_actions(actions)

  def get_actor_loss(self):
    return self._actor_network.get_actor_loss()

  def get_aux_info(self):
    return self._actor_network.get_aux_info()

  def set_training(self, training):
    self._training = training

  def _action(self,
              time_step: ts.TimeStep,
              policy_state: types.NestedTensor,
              seed: Optional[types.Seed] = None) -> policy_step.PolicyStep:
    del seed
    action, policy_state = self._apply_actor_network(
        time_step.observation,
        step_type=time_step.step_type,
        policy_state=policy_state)
    info = ()
    return policy_step.PolicyStep(action, policy_state, info)

  def _distribution(self, time_step, policy_state):
    current_step = super()._distribution(time_step, policy_state)
    return current_step


class SequenceAgent(tf_agent.TFAgent):
  """A sequence agent that directly outputs actions via an actor network."""

  def __init__(self,
               time_step_spec: ts.TimeStep,
               action_spec: types.NestedTensorSpec,
               actor_network: Type[network.Network],
               actor_optimizer: tf.keras.optimizers.Optimizer,
               policy_cls: Type[actor_policy.ActorPolicy] = SequencePolicy,
               time_sequence_length: int = 6,
               debug_summaries: bool = False,
               **kwargs):
    self._info_spec = ()
    self._actor_network = actor_network(  # pytype: disable=missing-parameter  # dynamic-method-lookup
        input_tensor_spec=time_step_spec.observation,
        output_tensor_spec=action_spec,
        policy_info_spec=self._info_spec,
        train_step_counter=kwargs['train_step_counter'],
        time_sequence_length=time_sequence_length)

    self._actor_optimizer = actor_optimizer
    # Train policy is only used for loss and never exported as saved_model.
    self._train_policy = policy_cls(
        time_step_spec=time_step_spec,
        action_spec=action_spec,
        info_spec=self._info_spec,
        actor_network=self._actor_network,
        training=True)
    collect_policy = policy_cls(
        time_step_spec=time_step_spec,
        action_spec=action_spec,
        info_spec=self._info_spec,
        actor_network=self._actor_network,
        training=False)
    super(SequenceAgent, self).__init__(
        time_step_spec,
        action_spec,
        collect_policy,  # We use the collect_policy as the eval policy.
        collect_policy,
        train_sequence_length=time_sequence_length,
        **kwargs)
    self._data_context = data_converter.DataContext(
        time_step_spec=time_step_spec,
        action_spec=action_spec,
        info_spec=collect_policy.info_spec,
        use_half_transition=True)
    self.as_transition = data_converter.AsHalfTransition(
        self._data_context, squeeze_time_dim=False)
    self._debug_summaries = debug_summaries

    num_params = 0
    for weight in self._actor_network.trainable_weights:
      weight_params = 1
      for dim in weight.shape:
        weight_params *= dim
      logging.info('%s has %s params.', weight.name, weight_params)
      num_params += weight_params
    logging.info('Actor network has %sM params.', round(num_params / 1000000.,
                                                        2))

  def _train(self, experience: types.NestedTensor,
             weights: types.Tensor) -> tf_agent.LossInfo:
    self.train_step_counter.assign_add(1)
    loss_info = self._loss(experience, weights, training=True)
    self._apply_gradients(loss_info.loss)
    return loss_info

  def _apply_gradients(self, loss: types.Tensor):
    variables = self._actor_network.trainable_weights
    gradients = tf.gradients(loss, variables)
    # Skip nan and inf gradients.
    new_gradients = []
    for g in gradients:
      if g is not None:
        new_g = tf.where(
            tf.math.logical_or(tf.math.is_inf(g), tf.math.is_nan(g)),
            tf.zeros_like(g), g)
        new_gradients.append(new_g)
      else:
        new_gradients.append(g)
    grads_and_vars = list(zip(new_gradients, variables))
    self._actor_optimizer.apply_gradients(grads_and_vars)

  def _loss(self, experience: types.NestedTensor, weights: types.Tensor,
            training: bool) -> tf_agent.LossInfo:
    transition = self.as_transition(experience)
    time_steps, policy_steps, _ = transition
    batch_size = nest_utils.get_outer_shape(time_steps, self._time_step_spec)[0]
    policy = self._train_policy
    policy.set_actions(policy_steps.action)
    policy.set_training(training=training)
    with tf.name_scope('actor_loss'):
      policy_state = policy.get_initial_state(batch_size)
      policy.action(time_steps, policy_state=policy_state)
      valid_mask = tf.cast(~time_steps.is_last(), tf.float32)
      loss = valid_mask * policy.get_actor_loss()
      loss = tf.reduce_mean(loss)
      policy.set_actions(None)
      self._actor_network.add_summaries(time_steps.observation,
                                        policy.get_aux_info(),
                                        self._debug_summaries, training)
      return tf_agent.LossInfo(loss=loss, extra=loss)

sequence_agent.py
YiyangHuang-work/RT-1
# Copyright 2022 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Tests for sequence_agent."""
from robotics_transformer.sequence_agent_test_set_up import SequenceAgentTestSetUp
import tensorflow as tf
from tf_agents.agents import data_converter


class SequenceAgentTest(SequenceAgentTestSetUp):

  def testAsTransitionType(self):
    agent = self.create_agent_and_initialize()
    self.assertIsInstance(agent.as_transition, data_converter.AsHalfTransition)


if __name__ == '__main__':
  tf.test.main()

sequence_agent_test.py
YiyangHuang-work/RT-1
# Copyright 2022 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Tests for sequence_agent."""
from typing import Type

import numpy as np
from robotics_transformer import sequence_agent
from tensor2robot.utils import tensorspec_utils
import tensorflow as tf
from tf_agents.networks import network
from tf_agents.policies import policy_saver
from tf_agents.specs import tensor_spec
from tf_agents.trajectories import time_step as ts


class DummyActorNet(network.Network):
  """Used for testing SequenceAgent and its subclass."""

  def __init__(self,
               output_tensor_spec=None,
               train_step_counter=None,
               policy_info_spec=None,
               time_sequence_length=1,
               use_tcl=False,
               **kwargs):
    super().__init__(**kwargs)

  @property
  def tokens_per_action(self):
    return 8

  def set_actions(self, actions):
    self._actions = actions

  def get_actor_loss(self):
    return self._actor_loss

  def call(self,
           observations,
           step_type,
           network_state,
           actions=None,
           training=False):
    del step_type
    image = observations['image']
    tf.expand_dims(tf.reduce_mean(image, axis=-1), -1)
    actions = tensorspec_utils.TensorSpecStruct(
        world_vector=tf.constant(1., shape=[1, 3]),
        rotation_delta=tf.constant(1., shape=[1, 3]),
        terminate_episode=tf.constant(1, shape=[1, 2]),
        gripper_closedness_action=tf.constant(1., shape=[1, 1]),
    )
    return actions, network_state

  @property
  def trainable_weights(self):
    return [tf.Variable(1.0)]


class SequenceAgentTestSetUp(tf.test.TestCase):
  """Defines spec for testing SequenceAgent and its subclass, tests create."""

  def setUp(self):
    super().setUp()
    self._action_spec = tensorspec_utils.TensorSpecStruct()
    self._action_spec.world_vector = tensor_spec.BoundedTensorSpec(
        (3,), dtype=tf.float32, minimum=-1., maximum=1., name='world_vector')

    self._action_spec.rotation_delta = tensor_spec.BoundedTensorSpec(
        (3,),
        dtype=tf.float32,
        minimum=-np.pi / 2,
        maximum=np.pi / 2,
        name='rotation_delta')

    self._action_spec.gripper_closedness_action = tensor_spec.BoundedTensorSpec(
        (1,),
        dtype=tf.float32,
        minimum=-1.,
        maximum=1.,
        name='gripper_closedness_action')
    self._action_spec.terminate_episode = tensor_spec.BoundedTensorSpec(
        (2,), dtype=tf.int32, minimum=0, maximum=1, name='terminate_episode')

    state_spec = tensorspec_utils.TensorSpecStruct()
    state_spec.image = tensor_spec.BoundedTensorSpec([256, 320, 3],
                                                     dtype=tf.float32,
                                                     name='image',
                                                     minimum=0.,
                                                     maximum=1.)
    state_spec.natural_language_embedding = tensor_spec.TensorSpec(
        shape=[512], dtype=tf.float32, name='natural_language_embedding')
    self._time_step_spec = ts.time_step_spec(observation_spec=state_spec)

    self.sequence_agent_cls = sequence_agent.SequenceAgent

  def create_agent_and_initialize(self,
                                  actor_network: Type[
                                      network.Network] = DummyActorNet,
                                  **kwargs):
    """Creates the agent and initialize it."""
    agent = self.sequence_agent_cls(
        time_step_spec=self._time_step_spec,
        action_spec=self._action_spec,
        actor_network=actor_network,
        actor_optimizer=tf.keras.optimizers.Adam(),
        train_step_counter=tf.compat.v1.train.get_or_create_global_step(),
        **kwargs)
    agent.initialize()
    return agent

  def testCreateAgent(self):
    """Creates the Agent and save the agent.policy."""
    agent = self.create_agent_and_initialize()
    self.assertIsNotNone(agent.policy)

    policy_model_saver = policy_saver.PolicySaver(
        agent.policy,
        train_step=tf.compat.v2.Variable(
            0,
            trainable=False,
            dtype=tf.int64,
            aggregation=tf.VariableAggregation.ONLY_FIRST_REPLICA,
            shape=()),
        input_fn_and_spec=None)
    save_options = tf.saved_model.SaveOptions(
        experimental_io_device='/job:localhost',
        experimental_custom_gradients=False)
    policy_model_saver.save('/tmp/unittest/policy/0', options=save_options)


if __name__ == '__main__':
  tf.test.main()

sequence_agent_test_set_up.py
YiyangHuang-work/RT-1
# Copyright 2022 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""A simple action tokenizer used with Robotics Transformer 1.

As an example, if an action is:
terminate = [0, 1]
world_vector = [0.9, 0.8, -0.3]
rotation_delta = [-0.1, 0.2, .6]
gripper_closedness = 0.9

Then we build a sequence of tokens of length 8 [one for each dimension].
The int32 type action dimensions are already assumed discrete and tokenized,
the float dimensions are bucketed according to the specs min and max. Each
dimension has 'vocab_size' buckets.

Currently, this tokenizer assumes one action spec and it is highly recommended
to specify the 'action_order', eg [terminate, world_vector, rotation_delta,
gripper_closedness]. Since after tokenization you lose that information, this
will be useful for debugging. Actions may also be subselected for prediction,
since not all actions are needed in the action_order.
"""
from typing import Optional

from tensor2robot.utils import tensorspec_utils
import tensorflow as tf


class RT1ActionTokenizer:
  """Tokenizes based on vocab size."""

  def __init__(self,
               action_spec: tensorspec_utils.TensorSpecStruct,
               vocab_size: int,
               action_order: Optional[list[str]] = None):
    """Instantiates an RT1ActionTokenizer.

    Args:
      action_spec: Tensor spec of the expected action tensor.
      vocab_size: Number of buckets to discretize action to.
      action_order: Order of the action names, used to discern the order of
        tokenized actions to detokenize and assemble back to action tensor
    """
    self._action_spec = action_spec
    self._vocab_size = vocab_size
    if action_order is None:
      self._action_order = self._action_spec.keys()
    else:
      for action in action_order:
        if action not in self._action_spec.keys():
          raise ValueError('actions: %s not found in action_spec: %s' %
                           (action, action_spec.keys()))
        assert action in self._action_spec.keys()
      self._action_order = action_order
    self._tokens_per_action = 0
    for action in self._action_order:
      action_shape = self._action_spec[action].shape
      if len(action_shape) != 1:
        raise ValueError(
            'Only action shapes with single dimension supported, got %s' %
            action_shape)
      if self._action_spec[action].dtype == tf.int32:
        # Int32 actions are already assumed to be tokens.
        self._tokens_per_action += 1
      else:
        self._tokens_per_action += action_shape[0]

    # We measure # of action tokens in two different way. One is by checking
    # from action_order (above) and the other is by looping through the
    # action spec (below). We aseert the # of action tokens are the same
    # calculated by these two ways. This will assure action_order is correctly
    # configured, otherwise, it will through an error in the assert.
    num_action_token = 0
    for spec in self._action_spec.values():
      if spec.dtype == tf.int32:
        num_action_token += 1
      else:
        num_action_token += spec.shape[-1]
    tf.debugging.assert_equal(num_action_token, self._tokens_per_action)

  @property
  def tokens_per_action(self) -> int:
    return self._tokens_per_action

  @property
  def action_spec(self) -> tensorspec_utils.TensorSpecStruct:
    return self._action_spec

  @property
  def action_order(self) -> list[str]:
    return self._action_order

  def tokenize(self, action: tensorspec_utils.TensorSpecStruct) -> tf.Tensor:
    """Tokenizes an action."""
    action_tokens = []
    for k in self._action_order:
      a = action[k]  # a is [batch, actions_size]
      spec = self._action_spec[k]
      if spec.dtype == tf.int32:
        # Int32 actions are already assumed to be tokens, assume it is smaller
        # than the vocab size, so all we need to do is pad zeros.
        tf.debugging.assert_equal(1, tf.reduce_sum(a, axis=-1))
        # extract the token [batch, 1]
        token = tf.argmax(a, axis=-1, output_type=tf.int32)
        tf.debugging.assert_less(token, self._vocab_size)
        # Add a seq dimension [batch, 1]
        token = tf.expand_dims(token, axis=-1)
      else:
        a = tf.clip_by_value(a, spec.minimum, spec.maximum)
        # Normalize the action [batch, actions_size]
        token = (a - spec.minimum) / (spec.maximum - spec.minimum)
        # Bucket and discretize the action to vocab_size, [batch, actions_size]
        token = tf.cast(token * (self._vocab_size - 1), tf.int32)
      action_tokens.append(token)
    # Append all actions, [batch, all_actions_size]
    action_tokens = tf.concat(action_tokens, axis=-1)
    return action_tokens

  def detokenize(self,
                 action_tokens: tf.Tensor) -> tensorspec_utils.TensorSpecStruct:
    """Detokenizes an action."""
    action = tensorspec_utils.TensorSpecStruct()
    token_index = 0
    for k in self._action_order:
      spec = self._action_spec[k]
      action_dim = spec.shape[0]
      if spec.dtype == tf.int32:
        # Int32 actions are already assumed to be tokens.
        action[k] = action_tokens[..., token_index]
        # A poor model may output tokens outside the allowed range, in that case
        # set them to a default value, the 0 token in this case.
        outside_range = tf.greater_equal(action[k], action_dim)
        action[k] = tf.where(outside_range, tf.zeros_like(action[k]), action[k])
        action[k] = tf.one_hot(
            action[k], depth=action_dim, axis=-1, dtype=tf.int32)
        token_index += 1
      else:
        actions = []
        for _ in range(action_dim):
          a = action_tokens[..., token_index:token_index + 1]
          a = tf.cast(a, tf.float32)
          a = a / (self._vocab_size - 1)
          a = (a * (spec.maximum - spec.minimum)) + spec.minimum
          actions.append(a)
          token_index += 1
        action[k] = tf.concat(actions, axis=-1)
    return action

tokenizers/action_tokenizer.py
YiyangHuang-work/RT-1
# Copyright 2022 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Tests for action_tokenizer."""
import numpy as np
from robotics_transformer.tokenizers import action_tokenizer
from tensor2robot.utils import tensorspec_utils
import tensorflow as tf
from tf_agents.specs import tensor_spec


class ActionTokenizerTest(tf.test.TestCase):

  def testTokenize_int32(self):
    action_spec = tensorspec_utils.TensorSpecStruct()
    action_spec.terminate_episode = tensor_spec.BoundedTensorSpec(
        (2,), dtype=tf.int32, minimum=0, maximum=1, name='terminate_episode')
    tokenizer = action_tokenizer.RT1ActionTokenizer(action_spec, vocab_size=10)
    self.assertEqual(1, tokenizer.tokens_per_action)
    action = tensorspec_utils.TensorSpecStruct(terminate_episode=[0, 1])
    action_tokens = tokenizer.tokenize(action)
    self.assertEqual([1], action_tokens.numpy())

  def testTokenize_int32_not_one_hot(self):
    action_spec = tensorspec_utils.TensorSpecStruct()
    action_spec.terminate_episode = tensor_spec.BoundedTensorSpec(
        (2,), dtype=tf.int32, minimum=0, maximum=1, name='terminate_episode')
    tokenizer = action_tokenizer.RT1ActionTokenizer(action_spec, vocab_size=10)
    self.assertEqual(1, tokenizer.tokens_per_action)
    action = tensorspec_utils.TensorSpecStruct(terminate_episode=[1, 8])
    with self.assertRaises(tf.errors.InvalidArgumentError):
      tokenizer.tokenize(action)

  def testDetokenize_int32(self):
    action_spec = tensorspec_utils.TensorSpecStruct()
    action_spec.terminate_episode = tensor_spec.BoundedTensorSpec(
        (2,), dtype=tf.int32, minimum=0, maximum=1, name='terminate_episode')
    tokenizer = action_tokenizer.RT1ActionTokenizer(action_spec, vocab_size=10)
    # 0 token should become a one hot: [1, 0]
    action = tokenizer.detokenize(tf.constant([0], dtype=tf.int32))
    self.assertSequenceEqual([1, 0], list(action['terminate_episode'].numpy()))
    # 1 token should become a one hot: [0, 1]
    action = tokenizer.detokenize(tf.constant([1], dtype=tf.int32))
    self.assertSequenceEqual([0, 1], list(action['terminate_episode'].numpy()))
    # OOV 3 token should become a default one hot: [1, 0]
    action = tokenizer.detokenize(tf.constant([3], dtype=tf.int32))
    self.assertSequenceEqual([1, 0], list(action['terminate_episode'].numpy()))

  def testTokenize_float(self):
    action_spec = tensorspec_utils.TensorSpecStruct()
    action_spec.world_vector = tensor_spec.BoundedTensorSpec(
        (3,), dtype=tf.float32, minimum=-1., maximum=1., name='world_vector')
    tokenizer = action_tokenizer.RT1ActionTokenizer(action_spec, vocab_size=10)
    self.assertEqual(3, tokenizer.tokens_per_action)
    action = tensorspec_utils.TensorSpecStruct(world_vector=[0.1, 0.5, -0.8])
    action_tokens = tokenizer.tokenize(action)
    self.assertSequenceEqual([4, 6, 0], list(action_tokens.numpy()))

  def testTokenize_float_with_time_dimension(self):
    action_spec = tensorspec_utils.TensorSpecStruct()
    action_spec.world_vector = tensor_spec.BoundedTensorSpec(
        (3,), dtype=tf.float32, minimum=-1., maximum=1., name='world_vector')
    tokenizer = action_tokenizer.RT1ActionTokenizer(action_spec, vocab_size=10)
    self.assertEqual(3, tokenizer.tokens_per_action)
    batch_size = 2
    time_dimension = 3
    action = tensorspec_utils.TensorSpecStruct(
        world_vector=tf.constant(
            [[0.1, 0.5, -0.8], [0.1, 0.5, -0.8], [0.1, 0.5, -0.8],
             [0.1, 0.5, -0.8], [0.1, 0.5, -0.8], [0.1, 0.5, -0.8]],
            shape=[batch_size, time_dimension, tokenizer.tokens_per_action]))
    action_tokens = tokenizer.tokenize(action)
    self.assertSequenceEqual(
        [batch_size, time_dimension, tokenizer.tokens_per_action],
        action_tokens.shape.as_list())

  def testTokenize_float_at_limits(self):
    minimum = -1.
    maximum = 1.
    vocab_size = 10
    action_spec = tensorspec_utils.TensorSpecStruct()
    action_spec.world_vector = tensor_spec.BoundedTensorSpec(
        (2,),
        dtype=tf.float32,
        minimum=minimum,
        maximum=maximum,
        name='world_vector')
    tokenizer = action_tokenizer.RT1ActionTokenizer(
        action_spec, vocab_size=vocab_size)
    self.assertEqual(2, tokenizer.tokens_per_action)
    action = tensorspec_utils.TensorSpecStruct(world_vector=[minimum, maximum])
    action_tokens = tokenizer.tokenize(action)
    # Minimum value will go to 0
    # Maximum value witll go to vocab_size-1
    self.assertSequenceEqual([0, vocab_size - 1], list(action_tokens.numpy()))

  def testTokenize_invalid_action_spec_shape(self):
    action_spec = tensorspec_utils.TensorSpecStruct()
    action_spec.world_vector = tensor_spec.BoundedTensorSpec(
        (2, 2), dtype=tf.float32, minimum=1, maximum=-1, name='world_vector')
    with self.assertRaises(ValueError):
      action_tokenizer.RT1ActionTokenizer(action_spec, vocab_size=10)

  def testTokenizeAndDetokenizeIsEqual(self):
    action_spec = tensorspec_utils.TensorSpecStruct()

    action_spec.world_vector = tensor_spec.BoundedTensorSpec(
        (3,), dtype=tf.float32, minimum=-1., maximum=1., name='world_vector')

    action_spec.rotation_delta = tensor_spec.BoundedTensorSpec(
        (3,),
        dtype=tf.float32,
        minimum=-np.pi / 2.,
        maximum=np.pi / 2.,
        name='rotation_delta')

    action_spec.gripper_closedness_action = tensor_spec.BoundedTensorSpec(
        (1,),
        dtype=tf.float32,
        minimum=-1.,
        maximum=1.,
        name='gripper_closedness_action')

    num_sub_action_space = 2
    action_spec.terminate_episode = tensor_spec.BoundedTensorSpec(
        (num_sub_action_space,),
        dtype=tf.int32,
        minimum=0,
        maximum=1,
        name='terminate_episode')

    tokenizer = action_tokenizer.RT1ActionTokenizer(
        action_spec,
        vocab_size=1024,
        action_order=[
            'terminate_episode', 'world_vector', 'rotation_delta',
            'gripper_closedness_action'
        ])
    self.assertEqual(8, tokenizer.tokens_per_action)

    # Repeat the following test N times with fuzzy inputs.
    n_repeat = 10
    for _ in range(n_repeat):
      action = tensorspec_utils.TensorSpecStruct(
          world_vector=np.random.uniform(low=-1., high=1.0, size=3),
          rotation_delta=np.random.uniform(
              low=-np.pi / 2., high=np.pi / 2., size=3),
          gripper_closedness_action=np.random.uniform(low=0., high=1.0, size=1),
          terminate_episode=[0, 1])
      action_tokens = tokenizer.tokenize(action)
      policy_action = tokenizer.detokenize(action_tokens)

      for k in action:
        self.assertSequenceAlmostEqual(
            action[k], policy_action[k].numpy(), places=2)

      # Repeat the test with batched actions
      batched_action = tensorspec_utils.TensorSpecStruct(
          world_vector=[
              np.random.uniform(low=-1., high=1.0, size=3),
              np.random.uniform(low=-1., high=1.0, size=3)
          ],
          rotation_delta=[
              np.random.uniform(low=-np.pi / 2., high=np.pi / 2., size=3),
              np.random.uniform(low=-np.pi / 2., high=np.pi / 2., size=3)
          ],
          gripper_closedness_action=[
              np.random.uniform(low=0., high=1.0, size=1),
              np.random.uniform(low=0., high=1.0, size=1)
          ],
          terminate_episode=[[0, 1], [1, 0]])
      action_tokens = tokenizer.tokenize(batched_action)
      policy_action = tokenizer.detokenize(action_tokens)

      for k in batched_action:
        for a, policy_a in zip(batched_action[k], policy_action[k].numpy()):
          self.assertSequenceAlmostEqual(a, policy_a, places=2)


if __name__ == '__main__':
  tf.test.main()

tokenizers/action_tokenizer_test.py
YiyangHuang-work/RT-1
# Copyright 2022 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""A FiLM Efficientnet contextual image tokenizer used in Robotics Transformer 1.
"""
from typing import Optional
from robotics_transformer.film_efficientnet import pretrained_efficientnet_encoder
from robotics_transformer.tokenizers import token_learner
import tensorflow as tf


class RT1ImageTokenizer(tf.keras.layers.Layer):
  """Tokenizes based on vocab size."""

  def __init__(self,
               embedding_output_dim: int,
               use_token_learner: bool = False,
               num_tokens: int = 8,
               **kwargs):
    """Instantiates a RT1ImageTokenizer.

    Args:
      embedding_output_dim: The output size of the tokens.
      use_token_learner: Whether to use token learner. See
        https://arxiv.org/abs/2106.11297
      num_tokens: Relevant only for token learner - the number of learned
        tokens.
      **kwargs: Keyword arguments to base class.
    """
    super().__init__(**kwargs)
    self._embedding_output_dim = embedding_output_dim

    self._tokenizer = pretrained_efficientnet_encoder.EfficientNetEncoder(
        pooling=False, early_film=True)

    self._use_token_learner = use_token_learner
    if self._use_token_learner:
      self._num_tokens = num_tokens
      self._token_learner = token_learner.TokenLearnerModule(
          num_tokens=self._num_tokens)

  @property
  def tokens_per_context_image(self) -> int:
    if self._use_token_learner:
      num_tokens = self._num_tokens
    else:
      num_tokens = 81
    return num_tokens

  def __call__(self,
               image: tf.Tensor,
               context: Optional[tf.Tensor] = None,
               training: bool = False) -> tf.Tensor:
    """Gets image tokens.

    Args:
      image: Images of shape (b, t, h, w, 3) to tokenize.
      context: An optional context vector (e.g., a natural language embedding).
        Expected to have shape (b, t, embedding_dim).
      training: Whether or not we are in training mode.

    Returns:
      tokens: has shape (batch, t, num_tokens_per_timestep, embedding_dim)
    """
    image_shape = tf.shape(image)
    b = image_shape[0]
    t = image_shape[1]
    h = image_shape[2]
    w = image_shape[3]
    c = image_shape[4]

    # Fold the time axis into the batch axis.
    image = tf.reshape(image, [b * t, h, w, c])
    if context is not None:
      context_rank = tf.rank(context)
      assertion = tf.Assert(context_rank == 3, data=[context_rank])
      with tf.control_dependencies([assertion]):
        context = tf.reshape(context, [b * t, tf.shape(context)[-1]])
    tokens = self.get_image_embeddings(image, context, training)
    if self._use_token_learner:
      tokens = self._token_learner(tokens, training)
    # Unflatten the time axis, which was previously flattened into the batch.
    tokens = tf.reshape(tokens, [b, t, tf.shape(tokens)[1], -1])
    return tokens

  def get_image_embeddings(self,
                           image: tf.Tensor,
                           context: Optional[tf.Tensor],
                           training: bool = False) -> tf.Tensor:
    """Gets embeddings from image.

    Args:
      image: Expected to be float32 in range [0, 1] with shape (b, h, w, 3).
      context: Expected to be float32 with shape (b, embedding_dim)
      training: Whether or not we are in training mode.

    Returns:
      tokens of shape (b, num_tokens, emedding_dim)
    """
    image_tokens = self._tokenizer(image, context=context, training=training)
    image_tokens = tf.reshape(image_tokens, [-1, 81, 512])
    return image_tokens

tokenizers/image_tokenizer.py
YiyangHuang-work/RT-1
# Copyright 2022 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Tests for image_tokenizer."""
from absl.testing import parameterized
from robotics_transformer.tokenizers import image_tokenizer
import tensorflow as tf


class ImageTokenizerTest(tf.test.TestCase, parameterized.TestCase):

  @parameterized.named_parameters(
      ('sample_image', 512, 224, False, 8),
      ('sample_image_token_learner', 512, 224, True, 8))
  def testTokenize(self, output_dim, image_resolution, use_token_learner,
                   num_tokens):
    batch = 1
    seq = 2
    tokenizer = image_tokenizer.RT1ImageTokenizer(
        embedding_output_dim=output_dim,
        use_token_learner=use_token_learner,
        num_tokens=num_tokens)

    image = tf.random.normal(
        shape=(batch, seq, image_resolution, image_resolution, 3))
    image = tf.clip_by_value(image, 0.0, 1.0)
    context_vector = tf.random.uniform((batch, seq, 512))
    image_tokens = tokenizer(image, context_vector)
    if use_token_learner:
      self.assertEqual(image_tokens.shape, [batch, seq, num_tokens, 512])
    else:
      self.assertEqual(image_tokens.shape, [batch, seq, 81, 512])


if __name__ == '__main__':
  tf.test.main()

tokenizers/image_tokenizer_test.py
YiyangHuang-work/RT-1
# Copyright 2022 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""TF implementation of Token Learner(Ryoo et al 2021)."""

import functools
from typing import Optional, Sequence, Union
import numpy as np
import tensorflow as tf


def gelu(x: float) -> float:
  return 0.5 * x * (1 +
                    tf.tanh(tf.sqrt(2 / np.pi) * (x + 0.044715 * tf.pow(x, 3))))


def _maybe_dropout(rate: float = 0.0, name: str = "dropout"):
  """Helper function to return dropout layer if rate is non zero."""
  if rate:
    return tf.keras.layers.Dropout(rate, name=name)
  return lambda x, *args: x  # Does nothing to x.


class MlpBlock(tf.keras.layers.Layer):
  """Transformer MLP / feed-forward block."""

  def __init__(self,
               *,
               mlp_dim: int,
               out_dim: Optional[int] = None,
               kernel_init: Optional[tf.keras.initializers.Initializer] = tf
               .keras.initializers.glorot_uniform(),
               bias_init: Optional[tf.keras.initializers.Initializer] = tf.keras
               .initializers.RandomNormal(stddev=1e-6),
               dropout_rate: float = 0.1,
               **kwargs):
    """Initializer for the MLP Block.

    This computes outer_dense(gelu(hidden_dense(input))), with dropout
    applied as necessary.

    Note: Especially outside a keras workflow, make sure to call layer.build

    Args:
      mlp_dim: The dimension of the inner representation (output of hidden
        layer). Usually larger than the input/output dim.
      out_dim: The output dimension of the block. If None, the model output dim
        is equal to the input dim (usually desired)
      kernel_init: Initializer for dense kernels, used for both dense layers.
      bias_init: Initializer for dense biases, used for both dense layers.
      dropout_rate: Dropout rate to be applied after dense ( & activation)
      **kwargs: Other keyword args passed to the tf.keras.layers.Layer
        constructor e.g. the name
    """
    super().__init__(**kwargs)
    self._out_dim = out_dim
    self._hidden_dropout = _maybe_dropout(dropout_rate)
    self._output_dropout = _maybe_dropout(dropout_rate)
    self._hidden_layer = tf.keras.layers.Dense(
        mlp_dim,
        activation=gelu,
        kernel_initializer=kernel_init,
        bias_initializer=bias_init,
        name="hidden_dense")

    # If out_dim is None, infer out_dim = input_dim at self.build()
    self._output_layer = functools.partial(
        tf.keras.layers.Dense,
        kernel_initializer=kernel_init,
        bias_initializer=bias_init,
        name="final_dense")

  def build(self, input_shape: Sequence[int]):
    out_dim = self._out_dim or input_shape[-1]
    self._output_layer = self._output_layer(units=out_dim)
    super().build(input_shape)

  def call(self,
           inputs: tf.Tensor,
           *,
           is_training: Union[bool, tf.Tensor] = False) -> tf.Tensor:
    """Applies Transformer MlpBlock module."""
    x = self._hidden_layer(inputs)
    x = self._hidden_dropout(x, is_training)
    x = self._output_layer(x)
    x = self._output_dropout(x, is_training)
    return x


class TokenLearnerModule(tf.keras.layers.Layer):
  """TokenLearner module V1.1 (https://arxiv.org/abs/2106.11297)."""

  def __init__(self,
               num_tokens: int,
               bottleneck_dim: int = 64,
               dropout_rate: float = 0.):
    super().__init__()

    self.mlp = MlpBlock(
        mlp_dim=bottleneck_dim, out_dim=num_tokens, dropout_rate=dropout_rate)
    self.layernorm = tf.keras.layers.LayerNormalization(epsilon=1e-6)

  def call(self, inputs: tf.Tensor, training: bool = False) -> tf.Tensor:
    if len(inputs.shape) == 4:
      bs, h, w, c = inputs.shape
      inputs = tf.reshape(inputs, [bs, h * w, c])

    selected = self.layernorm(inputs)

    selected = self.mlp(
        selected, is_training=training)  # Shape: [bs, h*w, n_token].

    selected = tf.transpose(selected, [0, 2, 1])  # Shape: [bs, n_token, h*w].
    selected = tf.nn.softmax(selected, axis=-1)

    feat = tf.einsum("...si,...id->...sd", selected, inputs)

    return feat  # Shape: [bs, n_token, c]

tokenizers/token_learner.py
YiyangHuang-work/RT-1
# Copyright 2022 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Tests for token_learner."""
from absl.testing import parameterized
from robotics_transformer.tokenizers import token_learner
import tensorflow as tf


class TokenLearnerTest(parameterized.TestCase):

  @parameterized.named_parameters(('sample_input', 512, 8))
  def testTokenLearner(self, embedding_dim, num_tokens):
    batch = 1
    seq = 2
    token_learner_layer = token_learner.TokenLearnerModule(
        num_tokens=num_tokens)

    inputvec = tf.random.normal(shape=(batch * seq, 81, embedding_dim))

    learnedtokens = token_learner_layer(inputvec)
    self.assertEqual(learnedtokens.shape,
                     [batch * seq, num_tokens, embedding_dim])


if __name__ == '__main__':
  tf.test.main()

tokenizers/token_learner_test.py
YiyangHuang-work/RT-1
# Copyright 2022 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""RT1 decoder transformer.

Copied from:
https://www.tensorflow.org/text/tutorials/transformer#decoder
"""
from typing import Tuple, Union

import tensorflow as tf


class _TransformerLayer(tf.keras.layers.Layer):
  """A single transformer block."""

  def __init__(self,
               layer_size: int = 4096,
               num_heads: int = 8,
               feed_forward_size: int = 512,
               dropout_rate: float = 0.1,
               return_attention_scores: bool = False):
    """Creates a Transformer layer.

    Args:
      layer_size: Size of the multiple head attention layer.
      num_heads: Number of heads for the multiple head attention layer.
      feed_forward_size: Dimensionality of the feed_forward layer.
      dropout_rate: Dropout rate.
      return_attention_scores: Return attention scores.
    """
    super(_TransformerLayer, self).__init__()

    self.layernorm1 = tf.keras.layers.LayerNormalization(epsilon=1e-6)
    self.mha1 = tf.keras.layers.MultiHeadAttention(
        key_dim=layer_size, num_heads=num_heads, dropout=dropout_rate)
    self.ff = tf.keras.layers.Dense(feed_forward_size)
    self.layernorm2 = tf.keras.layers.LayerNormalization(epsilon=1e-6)
    self.dropout_ff = tf.keras.layers.Dropout(dropout_rate)
    self._return_attention_scores = return_attention_scores

  def call(self, x: tf.Tensor, attention_mask: tf.Tensor,
           training: bool) -> Tuple[tf.Tensor, Union[tf.Tensor, None]]:
    """Calls the layer.

    Args:
      x: Input Tensor of shape `(B, T, dim)`.
      attention_mask: a boolean mask of shape `(B, T, T)`, that prevents
        attention to certain positions. The boolean mask specifies which query
        elements can attend to which key elements, 1 indicates attention and 0
        indicates no attention. Broadcasting can happen for the missing batch
        dimensions and the head dimension.
      training: Python boolean indicating whether the layer should behave in
        training mode (adding dropout) or in inference mode (no dropout).

    Returns:
      y: Output Tensor of shape `(B, T, dim)`. Also return the attention scores
      of shape `(B, T, dim)` or None.
    """
    x1 = self.layernorm1(x)
    mha_results = self.mha1(
        query=x1,
        key=x1,
        value=x1,
        attention_mask=attention_mask,
        return_attention_scores=self._return_attention_scores,
        training=training)
    if self._return_attention_scores:
      x1, score = mha_results
    else:
      x1, score = mha_results, None

    x = x + x1

    y = self.layernorm2(x)
    ff_y = self.ff(y)
    ff_y = self.dropout_ff(ff_y, training=training)
    x = x + ff_y
    return x, score


class Transformer(tf.keras.layers.Layer):
  """A decoder only transformer."""

  def __init__(self,
               num_layers: int = 1,
               layer_size: int = 4096,
               num_heads: int = 8,
               feed_forward_size: int = 512,
               dropout_rate: float = 0.1,
               vocab_size: int = 256,
               return_attention_scores: bool = False):
    """Creates a transformer.

    Args:
      num_layers: Number of transformer layers.
      layer_size: Size of the multiple head attention layer.
      num_heads: Number of heads for the multiple head attention layer.
      feed_forward_size: Dimensionality of the feed_forward layer.
      dropout_rate: Dropout rate.
      vocab_size: Dimensionality of tokens from the output layer.
      return_attention_scores: Return attention scores.
    """
    super(Transformer, self).__init__()

    self._layers = [
        _TransformerLayer(  # pylint: disable=g-complex-comprehension
            layer_size=layer_size,
            num_heads=num_heads,
            feed_forward_size=feed_forward_size,
            dropout_rate=dropout_rate,
            return_attention_scores=return_attention_scores)
        for _ in range(num_layers)
    ]
    self._token_emb = tf.keras.layers.Dense(feed_forward_size)
    self._position_emb = tf.keras.layers.Dense(feed_forward_size)
    self._output_tokens = tf.keras.layers.Dense(vocab_size)

  def call(
      self,
      x: tf.Tensor,
      training: bool,
      attention_mask: tf.Tensor,
  ) -> Union[tf.Tensor, Tuple[tf.Tensor, list[tf.Tensor]]]:
    """Calls the layer.

    Args:
      x: Input Tensor of shape `(B, T, dim)`.
      training: Python boolean indicating whether the layer should behave in
        training mode (adding dropout) or in inference mode (no dropout).
      attention_mask: a boolean mask of shape `(B, T, T)`, that prevents
        attention to certain positions. The boolean mask specifies which query
        elements can attend to which key elements, 1 indicates attention and 0
        indicates no attention. Broadcasting can happen for the missing batch
        dimensions and the head dimension.

    Returns:
      x: Output Tensor of shape `(B, T, vocab_size)`. If
      `return_attention_scores`, also return attention scores of
      a list of `layer` of elements with shape `(B, T, dim)`.
    """

    seq_len = tf.shape(x)[1]
    batch_size = tf.shape(x)[0]

    positions = tf.one_hot(
        tf.tile(tf.expand_dims(tf.range(0, seq_len, 1), 0), [batch_size, 1]),
        seq_len)

    x = self._token_emb(x)
    x += self._position_emb(positions)
    scores = []

    for layer in self._layers:
      x, score = layer(x, attention_mask=attention_mask, training=training)
      if score is not None:
        scores.append(score)
    x = self._output_tokens(x)
    return x, scores

transformer.py
YiyangHuang-work/RT-1
# Copyright 2022 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Tensorflow based methods for sequence agents."""
from typing import Optional, Tuple, Union, Any

from absl import logging
import numpy as np

from robotics_transformer import transformer
from robotics_transformer.film_efficientnet import preprocessors
from robotics_transformer.tokenizers import action_tokenizer
from robotics_transformer.tokenizers import image_tokenizer

from tensor2robot.utils import tensorspec_utils
import tensorflow as tf
from tf_agents.networks import network
from tf_agents.specs import tensor_spec
from tf_agents.utils import nest_utils


class TransformerNetwork(network.Network):
  """A transformer based actor network."""

  def __init__(
      self,
      input_tensor_spec: tensorspec_utils.TensorSpecStruct,
      output_tensor_spec: tensorspec_utils.TensorSpecStruct,
      train_step_counter: int = 0,
      vocab_size: int = 256,
      token_embedding_size: int = 512,
      num_layers: int = 1,
      layer_size: int = 4096,
      num_heads: int = 8,
      feed_forward_size: int = 512,
      dropout_rate: float = 0.1,
      time_sequence_length: int = 1,
      crop_size: int = 236,
      policy_info_spec: Optional[dict[Any,
                                      tensor_spec.BoundedTensorSpec]] = None,
      action_order: Optional[list[str]] = None,
      use_token_learner: Optional[bool] = True,
      return_attention_scores: bool = False,
      **kwargs):
    """Creates a transformer network.

    Args:
      input_tensor_spec: Nested list/tuple/dict of TensorSpecs, describing the
        shape of input tensor.
      output_tensor_spec: Nested list/tuple/dict of TensorSpecs, describing the
        shape of output tensor.
      train_step_counter: Counter for number of steps.
      vocab_size: Dimensionality of tokens from the output layer.
      token_embedding_size: Dimensionality of tokens from the embedding layer.
      num_layers: Number of transformer layers.
      layer_size: Size of the multiple head attention layer.
      num_heads: Number of heads for the multiple head attention layer.
      feed_forward_size: Dimensionality of the feed_forward layer.
      dropout_rate: Dropout rate.
      time_sequence_length: Length of the time sequence.
      crop_size: Height and width of the square crop, where original image will
        be padded to allow full field of view to be extracted.
      policy_info_spec: Spec on return value given return type of the return
        tokenizer.
      action_order: Order of actions for the action tokenizer.
      use_token_learner: Whether to use token learner. See
        https://arxiv.org/abs/2106.11297
      return_attention_scores: show attention scores in tensorboard.
      **kwargs: Keyword parameter arguments.
    """
    self._input_tensor_spec = input_tensor_spec
    self._output_tensor_spec = output_tensor_spec
    self._train_step_counter = train_step_counter
    self._actions = None
    self._returns = None
    self._vocab_size = vocab_size
    self._token_embedding_size = token_embedding_size
    self._time_sequence_length = time_sequence_length
    self._crop_size = crop_size

    self._transformer = transformer.Transformer(
        num_layers=num_layers,
        layer_size=layer_size,
        num_heads=num_heads,
        feed_forward_size=feed_forward_size,
        dropout_rate=dropout_rate,
        vocab_size=self._vocab_size,
        return_attention_scores=return_attention_scores)

    # create tokenizers
    self._image_tokenizer = image_tokenizer.RT1ImageTokenizer(
        embedding_output_dim=self._token_embedding_size,
        use_token_learner=use_token_learner)
    self._action_tokenizer = action_tokenizer.RT1ActionTokenizer(
        output_tensor_spec,
        vocab_size=self._vocab_size,
        action_order=action_order)

    self._tokens_per_action = self._action_tokenizer.tokens_per_action
    self._tokens_per_context_image = self._image_tokenizer.tokens_per_context_image
    # generate loss and attention masks
    self._generate_masks()

    # define mappings to token embedding size
    self._action_token_emb = tf.keras.layers.Dense(self._token_embedding_size)

    # define loss function
    self._loss_object = tf.keras.losses.SparseCategoricalCrossentropy(
        from_logits=True, reduction=tf.keras.losses.Reduction.NONE)
    self._attention_scores = []
    self._use_token_learner = use_token_learner

    super(TransformerNetwork, self).__init__(
        input_tensor_spec=input_tensor_spec, **kwargs)
    self._state_spec = {
        # Force this to be 4 dimension due to b/254902773.
        # Otherwise can be dimension 3.
        'context_image_tokens':
            tensor_spec.TensorSpec(
                shape=(time_sequence_length, self._tokens_per_context_image, 1,
                       token_embedding_size),
                dtype=tf.float32,
                name='context_image_tokens'),
        'action_tokens':
            tensor_spec.TensorSpec(
                shape=(time_sequence_length, self._tokens_per_action, 1, 1),
                dtype=tf.int32,
                name='action_tokens'),
        # Stores where in the window we are.
        # This value is within range [0, time_sequence_length + 1].
        # When seq_idx == time_sequence_length, context_image_tokens and
        # action_tokens need to be shifted to the left.
        'seq_idx':
            tensor_spec.TensorSpec(
                shape=(1, 1, 1, 1), dtype=tf.int32, name='seq_idx')
    }

  @property
  def attention_scores(self) -> list[tf.Tensor]:
    """Return attention score. This is for debugging/visualization purpose."""
    return self._attention_scores

  def _get_action_index_for_token(self, k):
    """Returns action associated with the token at given position `k`.

    If k is not an action token then it returns -1.
    If k is part of the first action in the sequence then returns 0 etc.

    Args:
        k: an int that represents the position in the sequence.

    Returns:
        The index of the action that this position belongs to, or if this
        position is part of an image token then returns -1.
    """
    if (k < 0 or k >= self._all_num_tokens):
      return -1

    n = k
    if n % self._single_time_step_num_tokens < self._tokens_per_context_image:
      return -1
    return int(n / self._single_time_step_num_tokens)

  def _generate_masks(self):
    """Generate mask for action prediction loss and attention visualization."""
    # each time step = [image, action]
    self._single_time_step_num_tokens = (
        self._tokens_per_action + self._tokens_per_context_image)

    # full sequence = [prefix context + N x timestep + postfix context]
    self._all_num_tokens = (
        self._time_sequence_length * self._single_time_step_num_tokens)

    # create mask for action predition loss
    self._action_tokens_mask = []
    for n in range(0, self._all_num_tokens, self._single_time_step_num_tokens):
      for x in range(0, self._tokens_per_action, 1):
        self._action_tokens_mask.append(x + n + self._tokens_per_context_image)
    self._action_tokens_mask = tf.constant(
        self._action_tokens_mask, dtype=tf.int32)

    # The look ahead mask ensures causality.
    self._default_attention_mask = tf.linalg.band_part(
        tf.ones((self._all_num_tokens, self._all_num_tokens)), -1, 0)

    action_mask = np.ndarray(
        shape=(self._all_num_tokens, self._all_num_tokens), dtype=int)
    for i in range(self._all_num_tokens):
      for j in range(self._all_num_tokens):
        action_i = self._get_action_index_for_token(i)
        action_j = self._get_action_index_for_token(j)
        mask = 0
        if action_i != -1 and action_j != -1:
          # Ignore actions of previous steps.
          if action_j < action_i:
            mask = 1
          # If we're not auto-regression, ignore action dimensions of current
          # step.
          if (action_j == action_i and j <= i):
            mask = 1
        action_mask[i, j] = mask
    self._default_attention_mask -= action_mask

  def _transformer_call(
      self,
      context_image_tokens: tf.Tensor,
      action_tokens: tf.Tensor,
      batch_size: int,
      training: bool,
      attention_mask: tf.Tensor,
  ) -> Union[tf.Tensor, Tuple[tf.Tensor, tf.Tensor]]:
    """Calls the transformer.

    Args:
      context_image_tokens: Tokenized context and image in Tensor of shape `(B,
        T, num token, -1)`.
      action_tokens: Discrete action token sequence of size [8, 256].
      batch_size: Batch size as when reshaping all tokens.
      training: Whether to run the transformer in training mode.
      attention_mask: Optional bool tensor for masking transformer's attention.

    Returns:
      Output tokens in Tensor of shape `(B, T, dim)`. If
      return_attention_scores, also return the attention scores of
      shape `(B, T, dim)`.
    """
    input_token_sequence = self._assemble_input_token_sequence(
        context_image_tokens, action_tokens, batch_size)

    # run transformer
    output_tokens, self._attention_scores = self._transformer(
        input_token_sequence, training, attention_mask)
    return output_tokens

  def _get_tokens_and_mask(self,
                           observations: dict[str, tf.Tensor],
                           network_state: dict[str, tf.Tensor],
                           training: bool = False):
    # tokenize all inputs
    context_image_tokens, network_state = self._tokenize_images(
        observations, network_state, training)
    action_tokens = self._tokenize_actions(observations, network_state)

    # generate transformer attention mask
    attention_mask = self._default_attention_mask

    return (context_image_tokens, action_tokens, attention_mask)

  def _transformer_call_and_slice(self,
                                  *args,
                                  slice_start: int = 0,
                                  slice_length: int = 1,
                                  **kwargs) -> Tuple[tf.Tensor, tf.Tensor]:
    output_tokens = self._transformer_call(*args, **kwargs)

    slice_end = slice_start + slice_length
    token_logits = output_tokens[:, slice_start:slice_end, :]
    token = tf.argmax(token_logits, axis=-1, output_type=tf.int32)

    return token, token_logits

  def call(self,
           observations: dict[str, tf.Tensor],
           network_state: dict[str, tf.Tensor],
           training: bool = False):
    """Calls the transformer network.

    Args:
      observations: Observation data including image and natural language
        embedding in dict of Tensors.
      network_state: Network state data including time step, image, action
        tokens, step number in dict of Tensors.
      training: Whether to call transformer network in training mode.

    Returns:
      A tuple `(Detokenized output actions, network state)`.
    """
    # used to determine training vs inference call
    # outer_rank will be 2 -> [b, t] during training and
    # outer_rank will be 1 -> [b] during inference
    outer_rank = self._get_outer_rank(observations)
    assert outer_rank in (1, 2)

    b, t = self._get_batch_size_and_seq_len(network_state)

    context_image_tokens, action_tokens, attention_mask = self._get_tokens_and_mask(
        observations, network_state, training)

    self._aux_info = {'action_labels': action_tokens}

    if outer_rank == 1:  # This is an inference call
      # run transformer in loop to produce action tokens one-by-one
      seq_idx = tf.reshape(network_state['seq_idx'], [1])[0]
      action_t = tf.minimum(seq_idx, self._time_sequence_length - 1)
      # Transformer shifts all to the left by one step by default (it's usually
      # predicting the next token as default training task...).
      transformer_shift = -1
      # We only want to get the action predicted at time_step.
      start_index = (
          transformer_shift + self._tokens_per_context_image + action_t *
          (self._single_time_step_num_tokens))
      current_action_tokens = []
      action_predictions_logits = []
      for k in range(self._tokens_per_action):
        action_index = start_index + k
        token, token_logits = self._transformer_call_and_slice(
            context_image_tokens,
            action_tokens,
            attention_mask=attention_mask,
            batch_size=b,
            training=training,
            slice_start=action_index  # slicing single action dimension
        )
        action_predictions_logits.append(token_logits)
        current_action_tokens.append(token)
        # action_tokens is [b, t * self._tokens_per_action]
        action_tokens = tf.reshape(action_tokens, [b, -1])
        action_start_index = (action_t * self._tokens_per_action) + k
        action_tokens = tf.concat([
            action_tokens[:, :action_start_index], token,
            action_tokens[:, action_start_index + 1:]
        ],
                                  axis=1)
        # action_tokens is [b, t, self._tokens_per_action]
        action_tokens = tf.reshape(action_tokens,
                                   [b, t, self._tokens_per_action])
      self._aux_info.update({
          # action_predictions_logits is
          # [b, self._tokens_per_action, self._vocab_size]
          'action_predictions_logits': tf.concat(action_predictions_logits, 1)
      })
      # predicted_tokens_for_output is [b, self._tokens_per_action]
      predicted_tokens_for_output = tf.concat(current_action_tokens, 1)
      # state_action_tokens is [b, 1, self._tokens_per_action, 1, 1]
      one_state_action_tokens = predicted_tokens_for_output[:, tf.newaxis, :,
                                                            tf.newaxis,
                                                            tf.newaxis]

      state_action_tokens = network_state['action_tokens']
      network_state['action_tokens'] = tf.concat([
          state_action_tokens[:, :action_t, ...], one_state_action_tokens,
          state_action_tokens[:, action_t + 1:, ...]
      ],
                                                 axis=1)
      # Increment the time_step for the next inference call.
      network_state['seq_idx'] = tf.reshape(
          tf.minimum(seq_idx + 1, self._time_sequence_length), [-1, 1, 1, 1, 1])

      self._loss = tf.constant(0.0)
    else:
      # training call --> simply run one transformer forward pass
      output_tokens = self._transformer_call(
          context_image_tokens,
          action_tokens,
          attention_mask=attention_mask,
          batch_size=b,
          training=training)

      # Gather all predicted actions for the action loss.
      action_logits = tf.gather(
          output_tokens, self._action_tokens_mask - 1, axis=1)
      action_logits_for_training = tf.reshape(
          action_logits, [b, t, self._tokens_per_action, -1])

      # Only take the last action as the action.
      # action_logits_for_output is [b, self._tokens_per_action, emb]
      action_logits_for_output = action_logits_for_training[:, -1]

      # predicted_tokens_for_output is [b, self._tokens_per_action]
      predicted_tokens_for_output = tf.argmax(
          action_logits_for_output, axis=-1, output_type=tf.int32)

      num_items = (
          tf.cast(b * t, tf.float32) * self._single_time_step_num_tokens)
      action_loss = tf.reduce_mean(
          self._loss_object(action_tokens, action_logits_for_training) /
          num_items,
          axis=-1)

      self._loss = action_loss

      # store action labels and predictions for visualization
      self._aux_info.update({
          'action_predictions':
              tf.argmax(
                  action_logits_for_training, axis=-1, output_type=tf.int32),
          'action_loss':
              action_loss,
          'actor_loss_mask':
              tf.ones([b], dtype=tf.float32)
      })

    output_actions = self._action_tokenizer.detokenize(
        predicted_tokens_for_output)
    return output_actions, network_state

  def add_summaries(self, observations: dict[str, tf.Tensor],
                    logging_info: dict[str, tf.Tensor], debug_summaries: bool,
                    training: bool) -> None:
    """Adds summaries.

    Args:
      observations: Observation data including image and natural language
        instruction in dict of Tensors.
      logging_info: Dict with all data stored for logging during training pass.
      debug_summaries: Whether to include debug summaries.
      training: Whether this function is called during training or inference.
    """
    num_params = 0
    for weight in self.trainable_weights:
      weight_params = 1
      for dim in weight.shape:
        weight_params *= dim
      num_params += weight_params
    tf.compat.v2.summary.scalar(name='num_params', data=num_params)
    # debug_summaries are for the non-tpu worker, train_summary.
    if debug_summaries:
      image = observations['image']  # [b, t, h, w, c]
      image_h = image.shape[2]
      image_w = image.shape[3]
      batch_size = image.shape[0]
      num_ts = image.shape[1]
      logging.info('image shape %s', image.shape)
      # Concat images for different timesteps across width.
      image = tf.concat(tf.unstack(image, axis=1), 2)
      # Concat images for different batches (up to 8) across height.
      image = tf.expand_dims(tf.concat(tf.unstack(image, axis=0)[0:8], 0), 0)
      tf.summary.image(
          'observations/image',
          image,
          step=self._train_step_counter,
          # Single output since we have concatenated images along batch.
          max_outputs=1)

      # [b, t], strings
      if 'natural_language_instruction' in observations:
        task = observations['natural_language_instruction'][:, 0]
        tf.summary.text(
            'natural_language_instruction', task, step=self._train_step_counter)
      if self.attention_scores and not self._use_token_learner:
        for l_idx, layer_attention_score in enumerate(self.attention_scores):
          logging.info('Attention score shape: %s, %s', l_idx,
                       layer_attention_score.shape)
          for head_idx in range(layer_attention_score.shape[1]):
            pairwise_attention = tf.expand_dims(
                layer_attention_score[:, head_idx], -1)
            # pairwise attention shape (16, 552, 552, 1)
            # make attention from different time steps comparable
            pairwise_attention = pairwise_attention * np.arange(
                1, pairwise_attention.shape[1] + 1)[None, :, None, None]

            # visualize spatial attention, note this only supports
            # mk1_500tasks_transformer pipeline with no token learner
            img_tf_ts = tf.reshape(
                tf.transpose(
                    tf.reshape(
                        tf.reduce_sum(pairwise_attention, axis=1) / np.arange(
                            pairwise_attention.shape[1], 0, -1)[None, :, None],
                        [batch_size, num_ts, -1]),
                    [0, 2, 1])[:, :-self._tokens_per_action, :],
                [-1, 9, 9, num_ts])

            img_tf_ts = tf.image.resize(
                img_tf_ts, [image_h, image_w],
                method=tf.image.ResizeMethod.NEAREST_NEIGHBOR)
            img_tf_ts_concat = tf.concat(tf.unstack(img_tf_ts, axis=3), 2)
            img_tf_ts_concat_min = tf.reduce_min(
                img_tf_ts_concat, axis=[1, 2], keepdims=True)
            img_tf_ts_concat = (img_tf_ts_concat - img_tf_ts_concat_min) / (
                tf.reduce_max(img_tf_ts_concat, axis=[1, 2], keepdims=True) -
                img_tf_ts_concat_min)
            img_tf_ts_concat = tf.concat(
                tf.unstack(img_tf_ts_concat, axis=0)[:8], 0)
            img_tf_ts_concat = tf.expand_dims(
                tf.expand_dims(img_tf_ts_concat, 0), -1)
            tf.summary.image(
                'attention/layer_{}/head_{}'.format(l_idx, head_idx),
                img_tf_ts_concat,
                step=self._train_step_counter,
                # Single output since we have concatenated images along batch.
                max_outputs=1)

            if img_tf_ts_concat.shape[1] == image.shape[
                1] and img_tf_ts_concat.shape[2] == image.shape[2]:
              # can overlay
              overlay_viz = tf.cast(
                  (tf.cast(image, tf.float32) * (0.2 + img_tf_ts_concat) / 1.2),
                  tf.uint8)
              tf.summary.image(
                  'overlay_attention/layer_{}/head_{}'.format(l_idx, head_idx),
                  overlay_viz,
                  step=self._train_step_counter,
                  # Single output since we have concatenated images along batch.
                  max_outputs=1)

    # log action info
    action_labels = tf.boolean_mask(logging_info['action_labels'],
                                    logging_info['actor_loss_mask'])
    action_predictions = tf.boolean_mask(logging_info['action_predictions'],
                                         logging_info['actor_loss_mask'])
    with tf.name_scope('ActionTokens'):
      token_accuracy = (
          tf.cast(tf.equal(action_labels, action_predictions), tf.float32))
      accuracy = tf.reduce_mean(token_accuracy)
      tf.compat.v2.summary.scalar(
          name='accuracy', data=accuracy, step=self._train_step_counter)
      # Accuracy across timesteps
      for t in range(self._time_sequence_length):
        tf.compat.v2.summary.scalar(
            name='accuracy/time_step/{}'.format(t),
            data=tf.reduce_mean(token_accuracy[:, t, :]),
            step=self._train_step_counter)
      token_index = 0
      for k in self._action_tokenizer.action_order:
        spec = self._action_tokenizer.action_spec[k]
        if spec.dtype == tf.int32:
          n_tokens = 1
        else:
          n_tokens = spec.shape[0]
        action_token_accuracy = tf.reduce_mean(
            token_accuracy[:, :, token_index:token_index + n_tokens])
        tf.compat.v2.summary.scalar(
            name='accuracy/action_type/{}'.format(k),
            data=action_token_accuracy,
            step=self._train_step_counter)
        for n in range(n_tokens):
          tf.summary.histogram(
              'tokens/{}_{}/labels'.format(k, n + 1),
              action_labels[:, :, token_index],
              step=self._train_step_counter)
          tf.summary.histogram(
              'tokens/{}_{}/predictions'.format(k, n + 1),
              action_predictions[:, :, token_index],
              step=self._train_step_counter)
          token_index += 1

    # log loss components
    with tf.name_scope('TokenLosses'):
      tf.compat.v2.summary.scalar(
          name='action_loss',
          data=tf.reduce_mean(logging_info['action_loss']),
          step=self._train_step_counter)

  def _tokenize_images(self, observations, network_state, training):
    image = observations['image']  # [b, t, h, w, c]
    outer_rank = self._get_outer_rank(observations)
    if outer_rank == 1:  # This is an inference call
      seq_idx = tf.reshape(network_state['seq_idx'], [1])[0]
      time_step = tf.minimum(seq_idx, self._time_sequence_length - 1)
      image = tf.expand_dims(image, 1)

    image_shape = tf.shape(image)
    b = image_shape[0]
    input_t = image_shape[1]
    h = image_shape[2]
    w = image_shape[3]
    c = image_shape[4]

    context = self._extract_context_from_observation(observations, input_t)

    image = tf.reshape(image, [b * input_t, h, w, c])
    seed = tf.random.uniform(shape=(2,), maxval=2**30, dtype=tf.int32)
    image = preprocessors.convert_dtype_and_crop_images(
        image,
        crop_size=self._crop_size,
        training=training,
        pad_then_crop=True,
        convert_dtype=True,
        seed=seed)
    image = tf.reshape(image, [b, input_t, h, w, c])
    context_image_tokens = self._image_tokenizer(
        image, context=context, training=training)
    num_tokens = tf.shape(context_image_tokens)[2]
    context_image_tokens = tf.reshape(context_image_tokens,
                                      [b, input_t, num_tokens, 1, -1])
    if outer_rank == 1:  # This is an inference call
      network_state['context_image_tokens'] = tf.reshape(
          network_state['context_image_tokens'], [
              b, self._time_sequence_length, self._tokens_per_context_image, 1,
              -1
          ])
      state_image_tokens = network_state['context_image_tokens']
      # network_state as input for this call is the output from the last call.
      # Therefore, we need to shift all images to the left by 1 in the time axis
      # to align w/ the time dim in this call.
      state_image_tokens = tf.cond(
          seq_idx == self._time_sequence_length,
          lambda: tf.roll(state_image_tokens, -1, axis=1),
          lambda: state_image_tokens)

      context_image_tokens = tf.concat([
          state_image_tokens[:, :time_step, ...], context_image_tokens,
          state_image_tokens[:, time_step + 1:, ...]
      ],
                                       axis=1)
      network_state['context_image_tokens'] = context_image_tokens

    return context_image_tokens, network_state

  def _tokenize_actions(self, observations, network_state):
    outer_rank = self._get_outer_rank(observations)
    if outer_rank == 1:  # This is an inference call
      action_tokens = tf.squeeze(network_state['action_tokens'], [3, 4])
      seq_idx = tf.reshape(network_state['seq_idx'], [1])[0]
      # network_state as input for this call is the output from the last call.
      # Therefore, we need to shift all actions by 1 to the left.
      action_tokens = tf.cond(seq_idx == self._time_sequence_length,
                              lambda: tf.roll(action_tokens, -1, axis=1),
                              lambda: action_tokens)
    else:
      assert outer_rank == 2
      if self._actions is None:
        b, t = self._get_batch_size_and_seq_len(network_state)
        action_tokens = tf.zeros(
            shape=[b, t, self._tokens_per_action], dtype=tf.int32)
      else:
        action_tokens = self._action_tokenizer.tokenize(self._actions)
    return action_tokens

  def _assemble_input_token_sequence(self, context_image_tokens, action_tokens,
                                     batch_size):
    # embed action tokens
    action_tokens = tf.one_hot(action_tokens, self._vocab_size)
    action_tokens = self._action_token_emb(action_tokens)
    action_tokens = tf.zeros_like(action_tokens)  # b/260260205

    # Because of b/254902773, we need to add 1 extra dimension.
    action_tokens = tf.expand_dims(action_tokens, axis=-2)

    # assemble token sequence
    input_token_sequence = tf.concat([context_image_tokens, action_tokens],
                                     axis=2)

    input_token_sequence = tf.reshape(
        input_token_sequence, [batch_size, -1, self._token_embedding_size])
    return input_token_sequence

  def _extract_context_from_observation(self, observations, seq_len):
    """Extract context from observation."""
    context = None
    if 'natural_language_embedding' in observations:
      outer_rank = self._get_outer_rank(observations)
      context = observations['natural_language_embedding']  # [b, t, emb-size]
      if outer_rank == 1:
        context = tf.tile(context[:, None], [1, seq_len, 1])
    return context

  def set_actions(self, actions: tensorspec_utils.TensorSpecStruct):
    """Sets actions that will be tokenized and used in transformer network.

    Args:
      actions: actions to be tokenized and used in transformer network. example
        actions are terminate = [0, 1] world_vector = [0.9, 0.8, -0.3]
        rotation_delta = [-0.1, 0.2, .6] gripper_closedness = 0.9
    """
    self._actions = actions

  def _get_outer_rank(self, observations):
    # used to determine training vs inference call
    # outer_rank will be 2 -> [b, t] during training and
    # outer_rank will be 1 -> [b] during inference
    return nest_utils.get_outer_rank(observations, self._input_tensor_spec)

  def _get_batch_size_and_seq_len(self, network_state):
    image_shape = tf.shape(network_state['context_image_tokens'])
    b = image_shape[0]
    t = image_shape[1]
    return b, t

  def get_actor_loss(self) -> tf.Tensor:
    return self._loss

  def get_aux_info(self) -> dict[str, Any]:
    return self._aux_info

transformer_network.py
YiyangHuang-work/RT-1
# Copyright 2022 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Tests for networks."""

from absl.testing import parameterized

from robotics_transformer import transformer_network
from robotics_transformer.transformer_network_test_set_up import BATCH_SIZE
from robotics_transformer.transformer_network_test_set_up import NAME_TO_INF_OBSERVATIONS
from robotics_transformer.transformer_network_test_set_up import NAME_TO_STATE_SPECS
from robotics_transformer.transformer_network_test_set_up import observations_list
from robotics_transformer.transformer_network_test_set_up import spec_names_list
from robotics_transformer.transformer_network_test_set_up import state_spec_list
from robotics_transformer.transformer_network_test_set_up import TIME_SEQUENCE_LENGTH
from robotics_transformer.transformer_network_test_set_up import TransformerNetworkTestUtils

import tensorflow as tf
from tf_agents.specs import tensor_spec


class TransformerNetworkTest(TransformerNetworkTestUtils):

  # pylint:disable=g-complex-comprehension
  @parameterized.named_parameters([{
      'testcase_name': '_' + name,
      'state_spec': spec,
      'train_observation': obs,
  } for (name, spec,
         obs) in zip(spec_names_list(), state_spec_list(), observations_list())]
                                 )
  # pylint:enable=g-complex-comprehension
  def testTransformerTrainLossCall(self, state_spec, train_observation):
    network = transformer_network.TransformerNetwork(
        input_tensor_spec=state_spec,
        output_tensor_spec=self._action_spec,
        time_sequence_length=TIME_SEQUENCE_LENGTH)

    network.create_variables()
    self.assertNotEmpty(network.variables)

    network.set_actions(self._train_action)
    network_state = tensor_spec.sample_spec_nest(
        network.state_spec, outer_dims=[BATCH_SIZE])
    output_actions, network_state = network(
        train_observation, step_type=None, network_state=network_state)
    expected_shape = [2, 3]
    self.assertEqual(network.get_actor_loss().shape,
                     tf.TensorShape(expected_shape))
    self.assertCountEqual(self._train_action.keys(), output_actions.keys())

  # pylint:disable=g-complex-comprehension
  @parameterized.named_parameters([{
      'testcase_name': '_' + name,
      'spec_name': name,
  } for name in spec_names_list()])
  # pylint:enable=g-complex-comprehension
  def testTransformerInferenceLossCall(self, spec_name):
    state_spec = NAME_TO_STATE_SPECS[spec_name]
    observation = NAME_TO_INF_OBSERVATIONS[spec_name]

    network = transformer_network.TransformerNetwork(
        input_tensor_spec=state_spec,
        output_tensor_spec=self._action_spec,
        time_sequence_length=TIME_SEQUENCE_LENGTH,
        action_order=[
            'terminate_episode', 'world_vector', 'rotation_delta',
            'gripper_closedness_action'
        ])
    network.create_variables()
    self.assertNotEmpty(network.variables)

    network.set_actions(self._inference_action)
    # inference currently only support batch size of 1
    network_state = tensor_spec.sample_spec_nest(
        network.state_spec, outer_dims=[1])

    output_actions, network_state = network(
        observation, step_type=None, network_state=network_state)

    tf.debugging.assert_equal(network.get_actor_loss(), 0.0)
    self.assertCountEqual(self._inference_action.keys(), output_actions.keys())

  # pylint:disable=g-complex-comprehension
  @parameterized.named_parameters([{
      'testcase_name': '_' + name,
      'state_spec': spec,
      'train_observation': obs,
  } for name, spec, obs in zip(spec_names_list(), state_spec_list(),
                               observations_list())])
  # pylint:enable=g-complex-comprehension
  def testTransformerLogging(self, state_spec, train_observation):
    network = transformer_network.TransformerNetwork(
        input_tensor_spec=state_spec,
        output_tensor_spec=self._action_spec,
        time_sequence_length=TIME_SEQUENCE_LENGTH,
        action_order=[
            'terminate_episode', 'world_vector', 'rotation_delta',
            'gripper_closedness_action'
        ])

    network.create_variables()
    self.assertNotEmpty(network.variables)

    network.set_actions(self._train_action)
    network_state = tensor_spec.sample_spec_nest(
        network.state_spec, outer_dims=[BATCH_SIZE])
    _ = network(train_observation, step_type=None, network_state=network_state)
    network.add_summaries(
        train_observation,
        network.get_aux_info(),
        debug_summaries=True,
        training=True)

  # pylint:disable=g-complex-comprehension
  @parameterized.named_parameters([{
      'testcase_name': '_' + name,
      'state_spec': spec,
  } for name, spec in zip(spec_names_list(), state_spec_list())])
  # pylint:enable=g-complex-comprehension
  def testTransformerCausality(self, state_spec):
    """Tests the causality for the transformer.

    Args:
      state_spec: Which state spec to test the transformer with
    """
    network = transformer_network.TransformerNetwork(
        input_tensor_spec=state_spec,
        output_tensor_spec=self._action_spec,
        time_sequence_length=TIME_SEQUENCE_LENGTH)
    network.create_variables()
    self.assertNotEmpty(network.variables)

    time_sequence_length = network._time_sequence_length
    tokens_per_image = network._tokens_per_context_image
    tokens_per_action = network._tokens_per_action

    def _split_image_and_action_tokens(all_tokens):
      image_start_indices = [(tokens_per_image + tokens_per_action) * k
                             for k in range(time_sequence_length)]
      image_tokens = tf.stack(
          [all_tokens[i:i + tokens_per_image] for i in image_start_indices],
          axis=0)
      action_start_indices = [i + tokens_per_image for i in image_start_indices]
      action_tokens = [
          tf.stack([
              all_tokens[i:i + tokens_per_action] for i in action_start_indices
          ], 0)
      ]
      image_tokens = tf.one_hot(image_tokens, network._token_embedding_size)
      # Remove extra dimension before the end once b/254902773 is fixed.
      shape = image_tokens.shape
      # Add batch dimension.
      image_tokens = tf.reshape(image_tokens,
                                [1] + shape[:-1] + [1] + shape[-1:])
      return image_tokens, action_tokens

    # Generate some random tokens for image and actions.
    all_tokens = tf.random.uniform(
        shape=[time_sequence_length * (tokens_per_image + tokens_per_action)],
        dtype=tf.int32,
        maxval=10,
        minval=0)
    context_image_tokens, action_tokens = _split_image_and_action_tokens(
        all_tokens)
    # Get the output tokens without any zeroed out input tokens.
    output_tokens = network._transformer_call(
        context_image_tokens=context_image_tokens,
        action_tokens=action_tokens,
        attention_mask=network._default_attention_mask,
        batch_size=1,
        training=False)[0]

    for t in range(time_sequence_length *
                   (tokens_per_image + tokens_per_action)):
      # Zero out future input tokens.
      all_tokens_at_t = tf.concat(
          [all_tokens[:t + 1],
           tf.zeros_like(all_tokens[t + 1:])], 0)
      context_image_tokens, action_tokens = _split_image_and_action_tokens(
          all_tokens_at_t)
      # Get the output tokens with zeroed out input tokens after t.
      output_tokens_at_t = network._transformer_call(
          context_image_tokens=context_image_tokens,
          action_tokens=action_tokens,
          attention_mask=network._default_attention_mask,
          batch_size=1,
          training=False)[0]
      # The output token is unchanged if future input tokens are zeroed out.
      self.assertAllEqual(output_tokens[:t + 1], output_tokens_at_t[:t + 1])

  def testLossMasks(self):
    self._define_specs()
    self._create_agent()
    image_tokens = 3
    action_tokens = 2
    self._agent._actor_network._time_sequence_length = 2
    self._agent._actor_network._tokens_per_context_image = image_tokens
    self._agent._actor_network._tokens_per_action = action_tokens
    self._agent._actor_network._generate_masks()
    self.assertAllEqual(
        self._agent._actor_network._action_tokens_mask,
        tf.constant([
            image_tokens, image_tokens + 1, 2 * image_tokens + action_tokens,
            2 * image_tokens + action_tokens + 1
        ], tf.int32))
    self._agent._actor_network._generate_masks()
    self.assertAllEqual(
        self._agent._actor_network._action_tokens_mask,
        tf.constant([
            image_tokens, image_tokens + 1, 2 * (image_tokens) + action_tokens,
            2 * (image_tokens) + action_tokens + 1
        ], tf.int32))


if __name__ == '__main__':
  # Useful to enable if running with ipdb.
  tf.config.run_functions_eagerly(True)
  tf.test.main()

transformer_network_test.py
YiyangHuang-work/RT-1
# Copyright 2022 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Tests for networks."""

import copy
from typing import Optional, Tuple, Union

from absl.testing import parameterized
import numpy as np
from robotics_transformer import sequence_agent
from robotics_transformer import transformer_network
from tensor2robot.utils import tensorspec_utils
import tensorflow as tf
from tf_agents.specs import tensor_spec
from tf_agents.trajectories import time_step as ts

BATCH_SIZE = 2
TIME_SEQUENCE_LENGTH = 3
HEIGHT = 256
WIDTH = 320
NUM_IMAGE_TOKENS = 2


def spec_names_list() -> list[str]:
  """Lists the different types of specs accepted by the transformer."""
  return ['default']


def state_spec_list() -> list[tensorspec_utils.TensorSpecStruct]:
  """Lists the different types of state spec accepted by the transformer."""
  state_spec = tensorspec_utils.TensorSpecStruct()
  state_spec.image = tensor_spec.BoundedTensorSpec([HEIGHT, WIDTH, 3],
                                                   dtype=tf.float32,
                                                   name='image',
                                                   minimum=0.,
                                                   maximum=1.)
  state_spec.natural_language_embedding = tensor_spec.TensorSpec(
      shape=[512], dtype=tf.float32, name='natural_language_embedding')

  state_spec_mask = copy.deepcopy(state_spec)
  state_spec_mask.initial_binary_mask = tensor_spec.BoundedTensorSpec(
      [HEIGHT, WIDTH, 1],
      dtype=tf.int32,
      name='initial_binary_mask',
      minimum=0,
      maximum=255)

  state_spec_tcl = copy.deepcopy(state_spec)
  state_spec_tcl.original_image = tensor_spec.BoundedTensorSpec(
      [HEIGHT, WIDTH, 3],
      dtype=tf.float32,
      name='original_image',
      minimum=0.,
      maximum=1.)

  return [
      state_spec,
      state_spec_mask,
      state_spec_tcl,
  ]


def observations_list(training: bool = True) -> list[dict[str, tf.Tensor]]:
  """Lists the different types of observations accepted by the transformer."""
  if training:
    image_shape = [BATCH_SIZE, TIME_SEQUENCE_LENGTH, HEIGHT, WIDTH, 3]
    emb_shape = [BATCH_SIZE, TIME_SEQUENCE_LENGTH, 512]
    mask_shape = [BATCH_SIZE, TIME_SEQUENCE_LENGTH, HEIGHT, WIDTH, 1]
  else:
    # inference currently only support batch size of 1
    image_shape = [1, HEIGHT, WIDTH, 3]
    emb_shape = [1, 512]
    mask_shape = [1, HEIGHT, WIDTH, 1]
  return [
      {
          'image': tf.constant(0.5, shape=image_shape),
          'natural_language_embedding': tf.constant(1., shape=emb_shape),
      },
      {
          'image': tf.constant(0.5, shape=image_shape),
          'natural_language_embedding': tf.constant(1., shape=emb_shape),
          'initial_binary_mask': tf.constant(192, shape=mask_shape),
      },
      {  # This is used for TCL.
          'image': tf.constant(0.5, shape=image_shape),
          'original_image': tf.constant(0.4, shape=image_shape),
          'natural_language_embedding': tf.constant(1., shape=emb_shape),
      },
  ]


NAME_TO_STATE_SPECS = dict(zip(spec_names_list(), state_spec_list()))
NAME_TO_OBSERVATIONS = dict(zip(spec_names_list(), observations_list()))
NAME_TO_INF_OBSERVATIONS = dict(
    zip(spec_names_list(), observations_list(False)))


class FakeImageTokenizer(tf.keras.layers.Layer):
  """Fake Image Tokenizer for testing Transformer."""

  def __init__(self,
               encoder: ...,
               position_embedding: ...,
               embedding_output_dim: int,
               patch_size: int,
               use_token_learner: bool = False,
               num_tokens: int = NUM_IMAGE_TOKENS,
               use_initial_binary_mask: bool = False,
               **kwargs):
    del encoder, position_embedding, patch_size, use_token_learner
    super().__init__(**kwargs)
    self.tokens_per_context_image = num_tokens
    if use_initial_binary_mask:
      self.tokens_per_context_image += 1
    self.embedding_output_dim = embedding_output_dim
    self.use_initial_binary_mask = use_initial_binary_mask

  def __call__(self,
               image: tf.Tensor,
               context: Optional[tf.Tensor] = None,
               initial_binary_mask: Optional[tf.Tensor] = None,
               training: bool = False) -> tf.Tensor:
    if self.use_initial_binary_mask:
      assert initial_binary_mask is not None
    image_shape = tf.shape(image)
    seq_size = image_shape[1]
    batch_size = image_shape[0]
    all_tokens = []
    num_tokens = self.tokens_per_context_image
    for t in range(seq_size):
      tokens = tf.ones([batch_size, 1, num_tokens, self.embedding_output_dim
                       ]) * image[0][t][0][0]
      all_tokens.append(tokens)
    return tf.concat(all_tokens, axis=1)


class TransformerNetworkTestUtils(tf.test.TestCase, parameterized.TestCase):
  """Defines specs, SequenceAgent, and various other testing utilities."""

  def _define_specs(self,
                    train_batch_size=BATCH_SIZE,
                    inference_batch_size=1,
                    time_sequence_length=TIME_SEQUENCE_LENGTH,
                    inference_sequence_length=TIME_SEQUENCE_LENGTH,
                    token_embedding_size=512,
                    image_width=WIDTH,
                    image_height=HEIGHT):
    """Defines specs and observations (both training and inference)."""
    self.train_batch_size = train_batch_size
    self.inference_batch_size = inference_batch_size
    self.time_sequence_length = time_sequence_length
    self.inference_sequence_length = inference_sequence_length
    self.token_embedding_size = token_embedding_size
    action_spec = tensorspec_utils.TensorSpecStruct()
    action_spec.world_vector = tensor_spec.BoundedTensorSpec(
        (3,), dtype=tf.float32, minimum=-1., maximum=1., name='world_vector')

    action_spec.rotation_delta = tensor_spec.BoundedTensorSpec(
        (3,),
        dtype=tf.float32,
        minimum=-np.pi / 2,
        maximum=np.pi / 2,
        name='rotation_delta')

    action_spec.gripper_closedness_action = tensor_spec.BoundedTensorSpec(
        (1,),
        dtype=tf.float32,
        minimum=-1.,
        maximum=1.,
        name='gripper_closedness_action')
    action_spec.terminate_episode = tensor_spec.BoundedTensorSpec(
        (2,), dtype=tf.int32, minimum=0, maximum=1, name='terminate_episode')

    state_spec = tensorspec_utils.TensorSpecStruct()
    state_spec.image = tensor_spec.BoundedTensorSpec(
        [image_height, image_width, 3],
        dtype=tf.float32,
        name='image',
        minimum=0.,
        maximum=1.)
    state_spec.natural_language_embedding = tensor_spec.TensorSpec(
        shape=[self.token_embedding_size],
        dtype=tf.float32,
        name='natural_language_embedding')
    self._policy_info_spec = {
        'return':
            tensor_spec.BoundedTensorSpec((),
                                          dtype=tf.float32,
                                          minimum=0.0,
                                          maximum=1.0,
                                          name='return'),
        'discounted_return':
            tensor_spec.BoundedTensorSpec((),
                                          dtype=tf.float32,
                                          minimum=0.0,
                                          maximum=1.0,
                                          name='discounted_return'),
    }

    self._state_spec = state_spec
    self._action_spec = action_spec

    self._inference_observation = {
        'image':
            tf.constant(
                1,
                shape=[self.inference_batch_size, image_height, image_width, 3],
                dtype=tf.dtypes.float32),
        'natural_language_embedding':
            tf.constant(
                1.,
                shape=[self.inference_batch_size, self.token_embedding_size],
                dtype=tf.dtypes.float32),
    }
    self._train_observation = {
        'image':
            tf.constant(
                0.5,
                shape=[
                    self.train_batch_size, self.time_sequence_length,
                    image_height, image_width, 3
                ]),
        'natural_language_embedding':
            tf.constant(
                1.,
                shape=[
                    self.train_batch_size, self.time_sequence_length,
                    self.token_embedding_size
                ]),
    }
    self._inference_action = {
        'world_vector':
            tf.constant(0.5, shape=[self.inference_batch_size, 3]),
        'rotation_delta':
            tf.constant(0.5, shape=[self.inference_batch_size, 3]),
        'terminate_episode':
            tf.constant(
                [0, 1] * self.inference_batch_size,
                shape=[self.inference_batch_size, 2]),
        'gripper_closedness_action':
            tf.constant(0.5, shape=[self.inference_batch_size, 1]),
    }
    self._train_action = {
        'world_vector':
            tf.constant(
                0.5,
                shape=[self.train_batch_size, self.time_sequence_length, 3]),
        'rotation_delta':
            tf.constant(
                0.5,
                shape=[self.train_batch_size, self.time_sequence_length, 3]),
        'terminate_episode':
            tf.constant(
                [0, 1] * self.train_batch_size * self.time_sequence_length,
                shape=[self.train_batch_size, self.time_sequence_length, 2]),
        'gripper_closedness_action':
            tf.constant(
                0.5,
                shape=[self.train_batch_size, self.time_sequence_length, 1]),
    }

  def _create_agent(self, actor_network=None):
    """Creates SequenceAgent using custom actor_network."""
    time_step_spec = ts.time_step_spec(observation_spec=self._state_spec)
    if actor_network is None:
      actor_network = transformer_network.TransformerNetwork

    self._agent = sequence_agent.SequenceAgent(
        time_step_spec=time_step_spec,
        action_spec=self._action_spec,
        actor_network=actor_network,
        actor_optimizer=tf.keras.optimizers.Adam(),
        train_step_counter=tf.compat.v1.train.get_or_create_global_step(),
        time_sequence_length=TIME_SEQUENCE_LENGTH)
    self._num_action_tokens = (
        # pylint:disable=protected-access
        self._agent._actor_network._action_tokenizer._tokens_per_action)
    # pylint:enable=protected-access

  def setUp(self):
    self._define_specs()
    super().setUp()

  def get_image_value(self, step_idx: int) -> float:
    return float(step_idx) / self.time_sequence_length

  def get_action_logits(self, batch_size: int, value: int,
                        vocab_size: int) -> tf.Tensor:
    return tf.broadcast_to(
        tf.one_hot(value % vocab_size, vocab_size)[tf.newaxis, tf.newaxis, :],
        [batch_size, 1, vocab_size])

  def create_obs(self, value) -> dict[str, tf.Tensor]:
    observations = {}
    observations['image'] = value * self._inference_observation['image']
    observations[
        'natural_language_embedding'] = value * self._inference_observation[
            'natural_language_embedding']
    return observations

  def fake_action_token_emb(self, action_tokens) -> tf.Tensor:
    """Just pad with zeros."""
    shape = action_tokens.shape
    assert self.vocab_size > self.token_embedding_size
    assert len(shape) == 4
    return action_tokens[:, :, :, :self.token_embedding_size]

  def fake_transformer(
      self, all_tokens, training,
      attention_mask) -> Union[tf.Tensor, Tuple[tf.Tensor, list[tf.Tensor]]]:
    """Fakes the call to TransformerNetwork._transformer."""
    del training
    del attention_mask
    # We expect ST00 ST01 A00 A01...
    # Where:
    # * ST01 is token 1 of state 0.
    # * A01 is token 1 of action 0.
    shape = all_tokens.shape.as_list()
    batch_size = shape[0]
    self.assertEqual(batch_size, 1)
    emb_size = self.token_embedding_size

    # transform to [batch_size, num_tokens, token_size]
    all_tokens = tf.reshape(all_tokens, [batch_size, -1, emb_size])
    # Pads tokens to be of vocab_size.
    self.assertGreater(self.vocab_size, self.token_embedding_size)
    all_shape = all_tokens.shape
    self.assertLen(all_shape.as_list(), 3)
    output_tokens = tf.concat([
        all_tokens,
        tf.zeros([
            all_shape[0], all_shape[1],
            self.vocab_size - self.token_embedding_size
        ])
    ],
                              axis=-1)
    num_tokens_per_step = NUM_IMAGE_TOKENS + self._num_action_tokens
    # Check state/action alignment.
    window_range = min(self._step_idx + 1, self.time_sequence_length)
    for j in range(window_range):
      # The index step that is stored in j = 0.
      first_step_idx = max(0, self._step_idx + 1 - self.time_sequence_length)
      image_idx = j * num_tokens_per_step
      action_start_index = image_idx + NUM_IMAGE_TOKENS
      for t in range(NUM_IMAGE_TOKENS):
        self.assertAllEqual(
            self.get_image_value(first_step_idx + j) *
            tf.ones_like(all_tokens[0][image_idx][:self.token_embedding_size]),
            all_tokens[0][image_idx + t][:self.token_embedding_size])
      # if j is not the current step in the window, all action dimensions
      # from previous steps are already infered and thus can be checked.
      action_dims_range = self.action_inf_idx if j == window_range - 1 else self._num_action_tokens
      for t in range(action_dims_range):
        token_idx = action_start_index + t
        action_value = (first_step_idx + j) * self._num_action_tokens + t
        self.assertAllEqual(
            self.get_action_logits(
                batch_size=batch_size,
                value=action_value,
                vocab_size=self.vocab_size)[0][0][:self.token_embedding_size],
            all_tokens[0][token_idx][:self.token_embedding_size])
    # Output the right action dimension value.
    image_token_index = (
        min(self._step_idx, self.time_sequence_length - 1) *
        num_tokens_per_step)
    transformer_shift = -1
    action_index = (
        image_token_index + NUM_IMAGE_TOKENS + self.action_inf_idx +
        transformer_shift)
    action_value = self._step_idx * self._num_action_tokens + self.action_inf_idx
    action_logits = self.get_action_logits(
        batch_size=batch_size, value=action_value, vocab_size=self.vocab_size)
    output_tokens = tf.concat([
        output_tokens[:, :action_index, :], action_logits[:, :, :],
        output_tokens[:, action_index + 1:, :]
    ],
                              axis=1)
    self.action_inf_idx = (self.action_inf_idx + 1) % self._num_action_tokens
    attention_scores = []
    return output_tokens, attention_scores

transformer_network_test_set_up.py
YiyangHuang-work/RT-1
# Copyright 2022 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Tests for transformer."""
from absl.testing import parameterized
from robotics_transformer import transformer
import tensorflow as tf


class TransformerTest(parameterized.TestCase):

  def setUp(self):
    self._vocab_size = 10
    batch_size = 8
    sequence_len = 12
    self._tokens = tf.random.uniform(
        [batch_size, sequence_len, self._vocab_size],
        minval=0,
        maxval=1,
        dtype=tf.dtypes.float32,
    )
    super(TransformerTest, self).setUp()

  @parameterized.parameters(True, False)
  def test_transformer_forwardpass(self, return_attention_scores):
    network = transformer.Transformer(
        num_layers=2,
        layer_size=512,
        num_heads=4,
        feed_forward_size=256,
        dropout_rate=0.1,
        vocab_size=self._vocab_size,
        return_attention_scores=return_attention_scores)

    output_tokens, attention_scores = network(self._tokens, attention_mask=None)
    self.assertSequenceEqual(self._tokens.shape.as_list(),
                             output_tokens.shape.as_list())
    if return_attention_scores:
      self.assertNotEmpty(attention_scores)
    else:
      self.assertEmpty(attention_scores)


if __name__ == '__main__':
  tf.test.main()

transformer_test.py
YiyangHuang-work/RT-1