Skip to content

tensorized_relative_state

TensorizedRelativeState

Bases: TensorizedState, RelativeObjectState

Tensorized state-mixin for RELATIVE (pairwise) values.

All pairwise values across all object state instances of this type are updated at once via _update_values(), rather than per individual instance update() call.

Multi-scene layout

VALUES (S, N, N, *value_shape) — [scene, self_obj, other_obj, ...] OBJ_IDXS {relative_prim_path: int} — shared N index across both pairwise dims IDX_OBJS list[list[obj|None]] — IDX_OBJS[scene_idx][obj_idx] = object instance

Conventions

  • Diagonal cells VALUES[s, i, i, ...] should be False / zero by convention; subclasses enforce this in their _update_values kernel.
  • Cross-scene cells (object a exists in scene s but b does not) should be False / zero; subclasses enforce via existence-mask checks in their kernel.

Subclasses are derived states (computed from simulator state each step). The default _dump_state / _load_state / serialize / deserialize / state_size are no-ops because pairwise values are recomputed and don't need to round-trip through saves. Subclasses may override if persistence is needed.

_set_value(self, other, new_value) raises NotImplementedError by default. Subclasses that need a setter should override it (typically by copying the original non-tensorized setter logic and writing into the (s, idx_self, idx_other) cell).

Source code in OmniGibson/omnigibson/object_states/tensorized_relative_state.py 
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
class TensorizedRelativeState(TensorizedState, RelativeObjectState):
    """
    Tensorized state-mixin for RELATIVE (pairwise) values.

    All pairwise values across all object state instances of this type are updated at once
    via _update_values(), rather than per individual instance update() call.

    Multi-scene layout
    ------------------
    VALUES      (S, N, N, *value_shape)  — [scene, self_obj, other_obj, ...]
    OBJ_IDXS    {relative_prim_path: int}  — shared N index across both pairwise dims
    IDX_OBJS    list[list[obj|None]]  — IDX_OBJS[scene_idx][obj_idx] = object instance

    Conventions
    -----------
    - Diagonal cells `VALUES[s, i, i, ...]` should be False / zero by convention; subclasses
      enforce this in their `_update_values` kernel.
    - Cross-scene cells (object `a` exists in scene s but `b` does not) should be False / zero;
      subclasses enforce via existence-mask checks in their kernel.

    Subclasses are derived states (computed from simulator state each step). The default
    `_dump_state` / `_load_state` / `serialize` / `deserialize` / `state_size` are no-ops
    because pairwise values are recomputed and don't need to round-trip through saves.
    Subclasses may override if persistence is needed.

    `_set_value(self, other, new_value)` raises NotImplementedError by default. Subclasses
    that need a setter should override it (typically by copying the original non-tensorized
    setter logic and writing into the (s, idx_self, idx_other) cell).
    """

    @classmethod
    def global_initialize(cls):
        """Initialize the global class-level tensors and indices for the (S, N, N, ...) shape."""
        cls.VALUES = th.empty(0, dtype=cls.value_type, device="cuda").reshape(0, 0, 0, *cls.value_shape)
        cls.VALUES_CPU = th.empty(0, dtype=cls.value_type).pin_memory().reshape(0, 0, 0, *cls.value_shape)
        cls.PREV_VALUES = th.empty(0, dtype=cls.value_type).reshape(0, 0, 0, *cls.value_shape)
        cls.OBJ_IDXS = {}  # {relative_prim_path: int}
        cls.IDX_OBJS = []  # list[list[obj|None]]

        # Per-pair state size (flattened size of value_shape). Full per-instance state size
        # depends on N; see `state_size` property below.
        cls.STATE_SIZE = math.prod(cls.value_shape)

    @classmethod
    def initialize_view(cls):
        """
        Rebuild all class-level tensors by scanning current objects across all scenes.
        Called from ``simulator.py`` after scene changes.

        Two-pass to avoid the O(N^3) grow-by-cat pattern the naive implementation has:
        pass 1 only touches Python dicts/lists to settle OBJ_IDXS / IDX_OBJS at their
        final size; pass 2 allocates VALUES once at (S, N, N, *value_shape) and fills
        carry-over values directly. For a scene with N objects across S scenes this is
        O(S*N^2) instead of O(N^3).

        Carry-over: for each (a, b) pair where both relative_prim_paths still exist, the
        previous VALUES[s, a, b] is preserved. New rows/columns are zero-initialized.
        """
        # Snapshot for carry-over (OBJ_IDXS / VALUES are None on the very first call)
        prev_obj_idxs = dict(cls.OBJ_IDXS) if cls.OBJ_IDXS is not None else {}
        prev_values = cls.VALUES.clone() if cls.VALUES is not None and cls.VALUES.numel() > 0 else None

        # Reset
        cls.global_initialize()

        # Pass 1: scan all scenes, settle OBJ_IDXS / IDX_OBJS at final shape (Python only).
        for scene_idx, scene in enumerate(og.sim.scenes):
            for obj in scene.objects:
                if cls not in obj.states:
                    continue
                rel_path = obj.relative_prim_path

                # Extend scene dimension if this is a new scene index
                while len(cls.IDX_OBJS) <= scene_idx:
                    cls.IDX_OBJS.append([None] * len(cls.OBJ_IDXS))

                # Register new relative path if first seen across all scenes
                if rel_path not in cls.OBJ_IDXS:
                    cls.OBJ_IDXS[rel_path] = len(cls.OBJ_IDXS)
                    for s_row in cls.IDX_OBJS:
                        s_row.append(None)

                cls.IDX_OBJS[scene_idx][cls.OBJ_IDXS[rel_path]] = obj

        # Pass 2: allocate VALUES once at the final shape, then fill carry-over values.
        S = len(cls.IDX_OBJS)
        N = len(cls.OBJ_IDXS)
        if S > 0 and N > 0:
            cls.VALUES = th.zeros((S, N, N, *cls.value_shape), dtype=cls.value_type, device="cuda")

            if prev_values is not None and prev_values.numel() > 0:
                # Fast path if we have identical object layout across scenes.
                # Used for topology changes that adds/removes no objects and preserves ordering — e.g. an assisted-grasp/attachment joint.
                if prev_obj_idxs == cls.OBJ_IDXS and prev_values.shape == cls.VALUES.shape:
                    cls.VALUES.copy_(prev_values)
                else:
                    # General carry-over, vectorized: for each scene, block-copy the surviving pairs'
                    # previous values.
                    # "surviving" = objects present both before and after this re-init, as
                    # (old_index, new_index) pairs — old_index into prev_values, new_index into VALUES.
                    surviving = [
                        (old_index, cls.OBJ_IDXS[rel_path])
                        for rel_path, old_index in prev_obj_idxs.items()
                        if rel_path in cls.OBJ_IDXS
                    ]
                    for scene_idx in range(min(prev_values.shape[0], S)):
                        # Restrict to survivors that actually exist in THIS scene.
                        in_scene = [
                            (old_i, new_i) for (old_i, new_i) in surviving if cls.IDX_OBJS[scene_idx][new_i] is not None
                        ]
                        if in_scene:
                            old_idx = th.tensor(
                                [old_i for old_i, _ in in_scene], dtype=th.long, device=cls.VALUES.device
                            )
                            new_idx = th.tensor(
                                [new_i for _, new_i in in_scene], dtype=th.long, device=cls.VALUES.device
                            )
                            # VALUES[scene, new_a, new_b] = prev_values[scene, old_a, old_b] for every (a, b)
                            # survivor pair, in one indexed assignment.
                            cls.VALUES[scene_idx][new_idx[:, None], new_idx[None, :]] = prev_values[scene_idx][
                                old_idx[:, None], old_idx[None, :]
                            ]

        # Rebuild pinned CPU mirror — synchronous copy so _get_value() is valid before first async copy
        cls.VALUES_CPU = th.zeros(cls.VALUES.shape, dtype=cls.value_type).pin_memory()
        if cls.VALUES.numel() > 0:
            cls.VALUES_CPU.copy_(cls.VALUES)

        cls.PREV_VALUES = cls.VALUES_CPU.clone()

        # Wrap as wp.array for kernel consumption
        if cls.VALUES.numel() > 0:
            cls.VALUES_WP = _wp_from_torch(cls.VALUES)
            cls.VALUES_CPU_WP = _wp_from_torch(cls.VALUES_CPU)
        else:
            cls.VALUES_WP = None
            cls.VALUES_CPU_WP = None

        # Mark the captured wp.graph as stale — the simulator will re-capture before the next step.
        TensorizedState.graph_dirty = True

        super().initialize_view()

    def _get_value(self, other):
        # Read from the pinned CPU mirror — no GPU stall for Python callers.
        s = self.obj.scene.idx
        idx_self = self.OBJ_IDXS[self.obj.relative_prim_path]
        idx_other = self.OBJ_IDXS[other.relative_prim_path]
        val = self.VALUES_CPU[s, idx_self, idx_other].to(self.value_type)
        if isinstance(val, th.Tensor) and val.numel() == 1:
            val = val.item()
        return val

    def _set_value(self, other, new_value):
        """Default: not implemented for tensorized relative states.

        Subclasses may override to support pairwise setters by copying the original
        non-tensorized setter logic — typically writing the value into both
        ``self.VALUES[s, idx_self, idx_other]`` and ``self.VALUES_CPU[s, idx_self, idx_other]``.
        Most tensorized relative states (e.g. Adjacency, Touching) are derived from
        simulator state and don't expose setters.
        """
        raise NotImplementedError(
            f"_set_value not implemented for {self.__class__.__name__}. "
            "Override in the subclass if a pairwise setter is needed."
        )

    # Pairwise values are derived state; persistence is a no-op by default.
    # Subclasses override if round-tripping through saves is required.

    def _dump_state(self):
        return {}

    def _load_state(self, state):
        pass

    def serialize(self, state):
        return th.empty(0)

    def deserialize(self, state):
        return {}, 0

    @property
    def state_size(self):
        return 0

    @classproperty
    def _do_not_register_classes(cls):
        # Don't register this class since it's an abstract template
        classes = super()._do_not_register_classes
        classes.add("TensorizedRelativeState")
        return classes

global_initialize() classmethod

Initialize the global class-level tensors and indices for the (S, N, N, ...) shape.

Source code in OmniGibson/omnigibson/object_states/tensorized_relative_state.py 
41
42
43
44
45
46
47
48
49
50
51
52
@classmethod
def global_initialize(cls):
    """Initialize the global class-level tensors and indices for the (S, N, N, ...) shape."""
    cls.VALUES = th.empty(0, dtype=cls.value_type, device="cuda").reshape(0, 0, 0, *cls.value_shape)
    cls.VALUES_CPU = th.empty(0, dtype=cls.value_type).pin_memory().reshape(0, 0, 0, *cls.value_shape)
    cls.PREV_VALUES = th.empty(0, dtype=cls.value_type).reshape(0, 0, 0, *cls.value_shape)
    cls.OBJ_IDXS = {}  # {relative_prim_path: int}
    cls.IDX_OBJS = []  # list[list[obj|None]]

    # Per-pair state size (flattened size of value_shape). Full per-instance state size
    # depends on N; see `state_size` property below.
    cls.STATE_SIZE = math.prod(cls.value_shape)

initialize_view() classmethod

Rebuild all class-level tensors by scanning current objects across all scenes. Called from simulator.py after scene changes.

Two-pass to avoid the O(N^3) grow-by-cat pattern the naive implementation has: pass 1 only touches Python dicts/lists to settle OBJ_IDXS / IDX_OBJS at their final size; pass 2 allocates VALUES once at (S, N, N, value_shape) and fills carry-over values directly. For a scene with N objects across S scenes this is O(SN^2) instead of O(N^3).

Carry-over: for each (a, b) pair where both relative_prim_paths still exist, the previous VALUES[s, a, b] is preserved. New rows/columns are zero-initialized.

Source code in OmniGibson/omnigibson/object_states/tensorized_relative_state.py 
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
@classmethod
def initialize_view(cls):
    """
    Rebuild all class-level tensors by scanning current objects across all scenes.
    Called from ``simulator.py`` after scene changes.

    Two-pass to avoid the O(N^3) grow-by-cat pattern the naive implementation has:
    pass 1 only touches Python dicts/lists to settle OBJ_IDXS / IDX_OBJS at their
    final size; pass 2 allocates VALUES once at (S, N, N, *value_shape) and fills
    carry-over values directly. For a scene with N objects across S scenes this is
    O(S*N^2) instead of O(N^3).

    Carry-over: for each (a, b) pair where both relative_prim_paths still exist, the
    previous VALUES[s, a, b] is preserved. New rows/columns are zero-initialized.
    """
    # Snapshot for carry-over (OBJ_IDXS / VALUES are None on the very first call)
    prev_obj_idxs = dict(cls.OBJ_IDXS) if cls.OBJ_IDXS is not None else {}
    prev_values = cls.VALUES.clone() if cls.VALUES is not None and cls.VALUES.numel() > 0 else None

    # Reset
    cls.global_initialize()

    # Pass 1: scan all scenes, settle OBJ_IDXS / IDX_OBJS at final shape (Python only).
    for scene_idx, scene in enumerate(og.sim.scenes):
        for obj in scene.objects:
            if cls not in obj.states:
                continue
            rel_path = obj.relative_prim_path

            # Extend scene dimension if this is a new scene index
            while len(cls.IDX_OBJS) <= scene_idx:
                cls.IDX_OBJS.append([None] * len(cls.OBJ_IDXS))

            # Register new relative path if first seen across all scenes
            if rel_path not in cls.OBJ_IDXS:
                cls.OBJ_IDXS[rel_path] = len(cls.OBJ_IDXS)
                for s_row in cls.IDX_OBJS:
                    s_row.append(None)

            cls.IDX_OBJS[scene_idx][cls.OBJ_IDXS[rel_path]] = obj

    # Pass 2: allocate VALUES once at the final shape, then fill carry-over values.
    S = len(cls.IDX_OBJS)
    N = len(cls.OBJ_IDXS)
    if S > 0 and N > 0:
        cls.VALUES = th.zeros((S, N, N, *cls.value_shape), dtype=cls.value_type, device="cuda")

        if prev_values is not None and prev_values.numel() > 0:
            # Fast path if we have identical object layout across scenes.
            # Used for topology changes that adds/removes no objects and preserves ordering — e.g. an assisted-grasp/attachment joint.
            if prev_obj_idxs == cls.OBJ_IDXS and prev_values.shape == cls.VALUES.shape:
                cls.VALUES.copy_(prev_values)
            else:
                # General carry-over, vectorized: for each scene, block-copy the surviving pairs'
                # previous values.
                # "surviving" = objects present both before and after this re-init, as
                # (old_index, new_index) pairs — old_index into prev_values, new_index into VALUES.
                surviving = [
                    (old_index, cls.OBJ_IDXS[rel_path])
                    for rel_path, old_index in prev_obj_idxs.items()
                    if rel_path in cls.OBJ_IDXS
                ]
                for scene_idx in range(min(prev_values.shape[0], S)):
                    # Restrict to survivors that actually exist in THIS scene.
                    in_scene = [
                        (old_i, new_i) for (old_i, new_i) in surviving if cls.IDX_OBJS[scene_idx][new_i] is not None
                    ]
                    if in_scene:
                        old_idx = th.tensor(
                            [old_i for old_i, _ in in_scene], dtype=th.long, device=cls.VALUES.device
                        )
                        new_idx = th.tensor(
                            [new_i for _, new_i in in_scene], dtype=th.long, device=cls.VALUES.device
                        )
                        # VALUES[scene, new_a, new_b] = prev_values[scene, old_a, old_b] for every (a, b)
                        # survivor pair, in one indexed assignment.
                        cls.VALUES[scene_idx][new_idx[:, None], new_idx[None, :]] = prev_values[scene_idx][
                            old_idx[:, None], old_idx[None, :]
                        ]

    # Rebuild pinned CPU mirror — synchronous copy so _get_value() is valid before first async copy
    cls.VALUES_CPU = th.zeros(cls.VALUES.shape, dtype=cls.value_type).pin_memory()
    if cls.VALUES.numel() > 0:
        cls.VALUES_CPU.copy_(cls.VALUES)

    cls.PREV_VALUES = cls.VALUES_CPU.clone()

    # Wrap as wp.array for kernel consumption
    if cls.VALUES.numel() > 0:
        cls.VALUES_WP = _wp_from_torch(cls.VALUES)
        cls.VALUES_CPU_WP = _wp_from_torch(cls.VALUES_CPU)
    else:
        cls.VALUES_WP = None
        cls.VALUES_CPU_WP = None

    # Mark the captured wp.graph as stale — the simulator will re-capture before the next step.
    TensorizedState.graph_dirty = True

    super().initialize_view()