Function chunk_act_obs

data/traj_transforms.py:12–74

Chunks actions and observations into the given window_size. "observation" keys are given a new axis (at index 1) of size `window_size` containing `window_size - 1` observations from the past and the current observation. "action" is given a new axis (at index 1) of size `window_size

(
    traj: Dict,
    backward_windows_size: int = 0,
    backward_delta: int = 1,
    forward_window_size: int = 0,
)

import tensorflow as tf

def chunk_act_obs(
    traj: Dict,
    backward_windows_size: int = 0,
    backward_delta: int = 1,
    forward_window_size: int = 0,
) -> Dict:
    """
    Chunks actions and observations into the given window_size.

    "observation" keys are given a new axis (at index 1) of size `window_size` containing `window_size - 1`
    observations from the past and the current observation. "action" is given a new axis (at index 1) of size
    `window_size + future_action_window_size` containing `window_size - 1` actions from the past, the current
    action, and `future_action_window_size` actions from the future. "pad_mask" is added to "observation" and
    indicates whether an observation should be considered padding (i.e. if it had come from a timestep
    before the start of the trajectory).
    """
    traj_len = tf.shape(traj["action"])[0]
    action_dim = traj["action"].shape[-1]
    chunk_indices = tf.broadcast_to(tf.range(-backward_windows_size, 1) * backward_delta, [traj_len, backward_windows_size + 1]) + tf.broadcast_to(
        tf.range(traj_len)[:, None], [traj_len, backward_windows_size + 1]
    )

    action_chunk_indices = tf.broadcast_to(
        tf.range(0, 1 + forward_window_size),
        [traj_len, 1 + forward_window_size],
    ) + tf.broadcast_to(
        tf.range(traj_len)[:, None],
        [traj_len, 1 + forward_window_size],
    )

    floored_chunk_indices = tf.maximum(chunk_indices, 0)

    if "timestep" in traj["task"]:
        goal_timestep = traj["task"]["timestep"]
    else:
        goal_timestep = tf.fill([traj_len], traj_len - 1)

    floored_action_chunk_indices = tf.minimum(tf.maximum(action_chunk_indices, 0), goal_timestep[:, None])

    traj["observation"] = tf.nest.map_structure(lambda x: tf.gather(x, floored_chunk_indices), traj["observation"])
    traj["action"] = tf.gather(traj["action"], floored_action_chunk_indices)

    # indicates whether an entire observation is padding
    traj["observation"]["pad_mask"] = chunk_indices >= 0

    # if no absolute_action_mask was provided, assume all actions are relative
    if "absolute_action_mask" not in traj and forward_window_size > 0:
        logging.warning(
            "future_action_window_size > 0 but no absolute_action_mask was provided. "
            "Assuming all actions are relative for the purpose of making neutral actions."
        )
    absolute_action_mask = traj.get("absolute_action_mask", tf.zeros([traj_len, action_dim], dtype=tf.bool))
    neutral_actions = tf.where(
        absolute_action_mask[:, None, :],
        traj["action"],  # absolute actions are repeated (already done during chunking)
        tf.zeros_like(traj["action"]),  # relative actions are zeroed
    )