Skip to content

tensorized_state

TensorizedState

Mixin holding the shared lifecycle and class-level tensors for tensorized states.

Concrete tensorized states should inherit through one of two front-ends
  • TensorizedAbsoluteState for single-object values, VALUES shape (S, N, *value_shape)
  • TensorizedRelativeState for pairwise values, VALUES shape (S, N, N, *value_shape)

Both front-ends provide their own global_initialize / initialize_view to allocate their respective VALUES tensor; everything else (pre/post update, change detection, wp.array wrapping, graph_dirty flag) lives here and is shared.

Subclasses MUST implement _update_values.

Source code in OmniGibson/omnigibson/object_states/tensorized_state.py 
 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
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
class TensorizedState:
    """
    Mixin holding the shared lifecycle and class-level tensors for tensorized states.

    Concrete tensorized states should inherit through one of two front-ends:
      - ``TensorizedAbsoluteState`` for single-object values, VALUES shape (S, N, *value_shape)
      - ``TensorizedRelativeState`` for pairwise values, VALUES shape (S, N, N, *value_shape)

    Both front-ends provide their own ``global_initialize`` / ``initialize_view`` to allocate
    their respective VALUES tensor; everything else (pre/post update, change detection,
    wp.array wrapping, graph_dirty flag) lives here and is shared.

    Subclasses MUST implement ``_update_values``.
    """

    # Tensor of raw internally tracked values — GPU-resident during computation.
    # Shape is (S, N, ...) for absolute or (S, N, N, ...) for relative.
    VALUES = None

    # Pinned CPU mirror of VALUES — updated via async DMA after each global_update pass.
    # _get_value() always reads from here to avoid GPU stalls for Python callers.
    VALUES_CPU = None

    # CPU tensor to store VALUES_CPU from the previous step — used for change detection.
    PREV_VALUES = None

    # Dictionary mapping relative prim path to index in the N dimension of VALUES
    OBJ_IDXS = None

    # 2-D list: IDX_OBJS[scene_idx][obj_idx] = object instance (or None if absent in that scene)
    IDX_OBJS = None

    # Int representing per-object state size (flattened size of value_shape)
    STATE_SIZE = None

    # wp.array views over VALUES for use inside Warp kernels and graph capture.
    # Wrapped once in initialize_view; never re-created per call. Re-wrapped on every
    # initialize_view because torch storage is reallocated there.
    VALUES_WP = None

    # wp.array view over the (pinned) VALUES_CPU mirror. Used so that the GPU→CPU mirror
    # copy inside the captured wp.graph runs as a wp.copy (graph-safe), not as torch's
    # .copy_ which uses the legacy stream and is forbidden during capture.
    VALUES_CPU_WP = None

    # Set to True any time any subclass's initialize_view runs, signalling that the
    # captured wp.graph holds stale pointers/shapes and must be re-captured. The simulator
    # checks-and-resets this before each step. update_handles() always calls the view APIs'
    # initialize_view before every TensorizedState's initialize_view, so any viewAPI
    # buffer reallocation is automatically covered.
    graph_dirty = True

    # Set to True whenever something has happened that could leave VALUES_CPU one frame
    # behind the live world: a physics step, a direct pose mutation, a joint mutation, etc.
    # Cleared by ``og.sim._refresh_state_caches()``. The lazy-refresh gate in
    # ``TensorizedAbsoluteState._get_value`` / ``TensorizedRelativeState._get_value`` checks
    # this flag and triggers a refresh on first read so callers always observe fresh state
    # without having to know about the cache lifecycle.
    caches_dirty = False

    # Re-entrance guard for the lazy refresh: while `_refresh_state_caches` is running, any
    # nested `_get_value` should read the (now-being-recomputed) cache directly without
    # re-triggering a refresh. Set/cleared inside the simulator helper.
    _refresh_in_progress = False

    # Single-element wp.array on cuda holding the seconds elapsed for the current logical step
    # (the caller-provided dt, see pre_update). Read by time-dependent subclasses' kernels
    # (Temperature, ToggledOn, SlicerActive) inside the captured graph. We use a wp.array (not a
    # wp.float32 scalar passed at launch) so per-frame mutation is visible to graph replays without
    # re-capture — scalars passed via wp.float32(...) bake the value at capture time.
    _dt = None

    @classmethod
    def maybe_refresh_caches(cls):
        """Helper for `get_value`: if caches are dirty (and we're not already
        mid-refresh), run the lazy refresh so the upcoming read sees fresh values.
        """
        if cls.caches_dirty and not cls._refresh_in_progress:
            import omnigibson as og  # local import to avoid module-level cycle

            og.sim._refresh_state_caches()

    @classmethod
    def initialize_view(cls):
        """
        Δt-tracking setup shared across tensorized states. Subclasses' initialize_view should
        call ``super().initialize_view()`` after they finish allocating VALUES.

        Allocates the 1-element wp.array that pre_update fills with the per-step dt each step.
        """
        if cls._dt is None:
            cls._dt = wp.zeros(shape=(1,), dtype=wp.float32, device="cuda")

    def get_value(self, *args, **kwargs):
        self.maybe_refresh_caches()
        return super().get_value(*args, **kwargs)

    @classmethod
    def pre_update(cls, dt=0.0):
        """
        CPU-side prep run BEFORE global_update each step. Snapshots VALUES_CPU into
        PREV_VALUES so post_update() can detect changes after the warp work completes,
        and stores the caller-provided ``dt`` (seconds elapsed for this logical step) into
        cls._dt for time-dependent kernels to read inside the captured graph.

        Lives outside the captured wp.graph.

        Subclasses may override to add per-state CPU prep (e.g. ToggledOn refreshing
        marker world poses).
        """
        if cls.VALUES_CPU is None or cls.VALUES_CPU.numel() == 0:
            return
        cls.PREV_VALUES.copy_(cls.VALUES_CPU)

        if cls._dt is not None:
            cls._dt.fill_(dt)

    @classmethod
    def global_update(cls):
        """
        Globally update all values via _update_values() and async-copy to VALUES_CPU.
        Change detection and post_update() are called separately by the simulator after synchronize().
        Skips if there are no tracked objects.

        Should be capturable inside wp.graph: only emits CUDA work via Warp
        kernel launches and Warp memcpy.
        """
        if cls.VALUES is None or cls.VALUES.numel() == 0:
            return

        cls._update_values(values=cls.VALUES)
        # Mirror VALUES → VALUES_CPU. Use wp.copy when both wp.array handles exist (graph-safe);
        # fall back to torch's non-blocking copy otherwise (e.g. partial init).
        # TODO(vector): Why would one (torch) exist and not the other (wp)?
        if cls.VALUES_WP is not None and cls.VALUES_CPU_WP is not None:
            wp.copy(cls.VALUES_CPU_WP, cls.VALUES_WP)
        else:
            cls.VALUES_CPU.copy_(cls.VALUES, non_blocking=True)

    @classmethod
    def post_update(cls):
        """
        Called by the simulator after th.cuda.synchronize(). Compares VALUES_CPU with PREV_VALUES
        (both CPU) to detect changes, fires state_updated() for affected objects, then updates PREV_VALUES.

        Reduce semantics:
        - Absolute (S, N, *value_shape): reduces over value_shape dims to get per-(scene, obj) row mask.
        - Relative (S, N, N, *value_shape): reduces over (other-obj, *value_shape) dims to get per-(scene, self-obj)
          row mask. A pairwise-cell change fires state_updated() once on the row's self-object.
        """
        if cls.VALUES_CPU is None or cls.VALUES_CPU.numel() == 0:
            return
        S = cls.VALUES_CPU.shape[0]

        diff = cls.VALUES_CPU != cls.PREV_VALUES
        changed_mask = th.any(diff, dim=tuple(range(2, diff.ndim))) if diff.ndim > 2 else diff
        for s_idx in range(S):
            for obj_idx in th.where(changed_mask[s_idx])[0].tolist():
                obj = cls.IDX_OBJS[s_idx][obj_idx]
                if obj is None:
                    # Column vacated in this scene (e.g. mid-play object removal) but kept
                    # alive by another scene sharing the relative prim path — see IDX_OBJS.
                    continue
                obj.state_updated()

    @classmethod
    def _update_values(cls, values):
        """
        Updates all internally tracked @values for this object state. Should be implemented by subclass.
        Mutates @values in-place. Must not return anything.

        Args:
            values (th.tensor): Tensorized value array. Shape depends on the front-end:
                (S, N, *value_shape) for TensorizedAbsoluteState,
                (S, N, N, *value_shape) for TensorizedRelativeState.
        """
        raise NotImplementedError

    @classproperty
    def value_shape(cls):
        """
        Returns:
            tuple: Expected shape of the per-object state instance value. Default is () (scalar).
        """
        return ()

    @classproperty
    def value_type(cls):
        """
        Returns:
            type: Type of the internal value array, e.g., bool, th.uint, th.float32, etc. Default is th.float32
        """
        return th.float32

    @classproperty
    def value_name(cls):
        """
        Returns:
            str: Name of the value key to assign when dumping / loading the state. Should be implemented by subclass
        """
        raise NotImplementedError

    @property
    def state_size(self):
        # This is merely the class state size
        return self.STATE_SIZE

global_update() classmethod

Globally update all values via _update_values() and async-copy to VALUES_CPU. Change detection and post_update() are called separately by the simulator after synchronize(). Skips if there are no tracked objects.

Should be capturable inside wp.graph: only emits CUDA work via Warp kernel launches and Warp memcpy.

Source code in OmniGibson/omnigibson/object_states/tensorized_state.py 
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
@classmethod
def global_update(cls):
    """
    Globally update all values via _update_values() and async-copy to VALUES_CPU.
    Change detection and post_update() are called separately by the simulator after synchronize().
    Skips if there are no tracked objects.

    Should be capturable inside wp.graph: only emits CUDA work via Warp
    kernel launches and Warp memcpy.
    """
    if cls.VALUES is None or cls.VALUES.numel() == 0:
        return

    cls._update_values(values=cls.VALUES)
    # Mirror VALUES → VALUES_CPU. Use wp.copy when both wp.array handles exist (graph-safe);
    # fall back to torch's non-blocking copy otherwise (e.g. partial init).
    # TODO(vector): Why would one (torch) exist and not the other (wp)?
    if cls.VALUES_WP is not None and cls.VALUES_CPU_WP is not None:
        wp.copy(cls.VALUES_CPU_WP, cls.VALUES_WP)
    else:
        cls.VALUES_CPU.copy_(cls.VALUES, non_blocking=True)

initialize_view() classmethod

Δt-tracking setup shared across tensorized states. Subclasses' initialize_view should call super().initialize_view() after they finish allocating VALUES.

Allocates the 1-element wp.array that pre_update fills with the per-step dt each step.

Source code in OmniGibson/omnigibson/object_states/tensorized_state.py 
103
104
105
106
107
108
109
110
111
112
@classmethod
def initialize_view(cls):
    """
    Δt-tracking setup shared across tensorized states. Subclasses' initialize_view should
    call ``super().initialize_view()`` after they finish allocating VALUES.

    Allocates the 1-element wp.array that pre_update fills with the per-step dt each step.
    """
    if cls._dt is None:
        cls._dt = wp.zeros(shape=(1,), dtype=wp.float32, device="cuda")

maybe_refresh_caches() classmethod

Helper for get_value: if caches are dirty (and we're not already mid-refresh), run the lazy refresh so the upcoming read sees fresh values.

Source code in OmniGibson/omnigibson/object_states/tensorized_state.py 
 93
 94
 95
 96
 97
 98
 99
100
101
@classmethod
def maybe_refresh_caches(cls):
    """Helper for `get_value`: if caches are dirty (and we're not already
    mid-refresh), run the lazy refresh so the upcoming read sees fresh values.
    """
    if cls.caches_dirty and not cls._refresh_in_progress:
        import omnigibson as og  # local import to avoid module-level cycle

        og.sim._refresh_state_caches()

post_update() classmethod

Called by the simulator after th.cuda.synchronize(). Compares VALUES_CPU with PREV_VALUES (both CPU) to detect changes, fires state_updated() for affected objects, then updates PREV_VALUES.

Reduce semantics: - Absolute (S, N, value_shape): reduces over value_shape dims to get per-(scene, obj) row mask. - Relative (S, N, N, value_shape): reduces over (other-obj, *value_shape) dims to get per-(scene, self-obj) row mask. A pairwise-cell change fires state_updated() once on the row's self-object.

Source code in OmniGibson/omnigibson/object_states/tensorized_state.py 
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
@classmethod
def post_update(cls):
    """
    Called by the simulator after th.cuda.synchronize(). Compares VALUES_CPU with PREV_VALUES
    (both CPU) to detect changes, fires state_updated() for affected objects, then updates PREV_VALUES.

    Reduce semantics:
    - Absolute (S, N, *value_shape): reduces over value_shape dims to get per-(scene, obj) row mask.
    - Relative (S, N, N, *value_shape): reduces over (other-obj, *value_shape) dims to get per-(scene, self-obj)
      row mask. A pairwise-cell change fires state_updated() once on the row's self-object.
    """
    if cls.VALUES_CPU is None or cls.VALUES_CPU.numel() == 0:
        return
    S = cls.VALUES_CPU.shape[0]

    diff = cls.VALUES_CPU != cls.PREV_VALUES
    changed_mask = th.any(diff, dim=tuple(range(2, diff.ndim))) if diff.ndim > 2 else diff
    for s_idx in range(S):
        for obj_idx in th.where(changed_mask[s_idx])[0].tolist():
            obj = cls.IDX_OBJS[s_idx][obj_idx]
            if obj is None:
                # Column vacated in this scene (e.g. mid-play object removal) but kept
                # alive by another scene sharing the relative prim path — see IDX_OBJS.
                continue
            obj.state_updated()

pre_update(dt=0.0) classmethod

CPU-side prep run BEFORE global_update each step. Snapshots VALUES_CPU into PREV_VALUES so post_update() can detect changes after the warp work completes, and stores the caller-provided dt (seconds elapsed for this logical step) into cls._dt for time-dependent kernels to read inside the captured graph.

Lives outside the captured wp.graph.

Subclasses may override to add per-state CPU prep (e.g. ToggledOn refreshing marker world poses).

Source code in OmniGibson/omnigibson/object_states/tensorized_state.py 
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
@classmethod
def pre_update(cls, dt=0.0):
    """
    CPU-side prep run BEFORE global_update each step. Snapshots VALUES_CPU into
    PREV_VALUES so post_update() can detect changes after the warp work completes,
    and stores the caller-provided ``dt`` (seconds elapsed for this logical step) into
    cls._dt for time-dependent kernels to read inside the captured graph.

    Lives outside the captured wp.graph.

    Subclasses may override to add per-state CPU prep (e.g. ToggledOn refreshing
    marker world poses).
    """
    if cls.VALUES_CPU is None or cls.VALUES_CPU.numel() == 0:
        return
    cls.PREV_VALUES.copy_(cls.VALUES_CPU)

    if cls._dt is not None:
        cls._dt.fill_(dt)

value_name()

Returns:

Type Description
str

Name of the value key to assign when dumping / loading the state. Should be implemented by subclass
Source code in OmniGibson/omnigibson/object_states/tensorized_state.py 
215
216
217
218
219
220
221
@classproperty
def value_name(cls):
    """
    Returns:
        str: Name of the value key to assign when dumping / loading the state. Should be implemented by subclass
    """
    raise NotImplementedError

value_shape()

Returns:

Type Description
tuple

Expected shape of the per-object state instance value. Default is () (scalar).
Source code in OmniGibson/omnigibson/object_states/tensorized_state.py 
199
200
201
202
203
204
205
@classproperty
def value_shape(cls):
    """
    Returns:
        tuple: Expected shape of the per-object state instance value. Default is () (scalar).
    """
    return ()

value_type()

Returns:

Type Description
type

Type of the internal value array, e.g., bool, th.uint, th.float32, etc. Default is th.float32
Source code in OmniGibson/omnigibson/object_states/tensorized_state.py 
207
208
209
210
211
212
213
@classproperty
def value_type(cls):
    """
    Returns:
        type: Type of the internal value array, e.g., bool, th.uint, th.float32, etc. Default is th.float32
    """
    return th.float32