From c6d218d8365752af9546e3ca8b43ca7e88f59b82 Mon Sep 17 00:00:00 2001 From: Vikash Kumar Date: Wed, 12 Aug 2026 20:19:46 -0400 Subject: [PATCH] FEATURE: Add resize/crop visual encoders Camera capture resolution and encoder output resolution were previously conflated (the same HxW drove both). Adds and encoder_ids so a camera can capture at one resolution (hardware-native, or shared across tasks) while the observation is resized/center-cropped to whatever a downstream model needs, using torchvision to stay consistent with the existing r3m/rrl/vc1 transform pipelines. Documents the visual_keys format and design rationale in robohive/envs/VISUAL_ENCODERS.md. --- robohive/envs/VISUAL_ENCODERS.md | 77 ++++++++++++++++++++++++++++++++ robohive/envs/env_base.py | 20 +++++++++ 2 files changed, 97 insertions(+) create mode 100644 robohive/envs/VISUAL_ENCODERS.md diff --git a/robohive/envs/VISUAL_ENCODERS.md b/robohive/envs/VISUAL_ENCODERS.md new file mode 100644 index 00000000..d3f51ce1 --- /dev/null +++ b/robohive/envs/VISUAL_ENCODERS.md @@ -0,0 +1,77 @@ +# Visual Encoders + +RoboHive can render a camera and pass it through an encoder before it becomes part of the +observation, via the `visual_keys` mechanism (`MujocoEnv._setup(visual_keys=...)`, or set +through an env registration's `variants={'visual_keys': [...]}`, e.g. in +`robohive/envs/arms/__init__.py`). Implementation: `robohive/envs/env_base.py`, +`_setup_rgb_encoders()` (builds the encoder) and `get_visuals()` (applies it every step). + +## Key format + +``` +rgb::x: +``` + +- `` — a camera defined in the robot config. +- `x` — the **capture resolution**: what's rendered (sim) or streamed (hardware) for + that camera. This is independent of whatever size the encoder ultimately outputs. +- `` — which encoder to apply to the captured image (see table below). + +A companion depth key `d::x:` returns the raw depth map for the +same camera/resolution if also present in `visual_keys`. + +All `rgb:` keys on a single env instance currently share **one** capture resolution and +**one** encoder — this is enforced by an assert in `_setup_rgb_encoders`. That's a scope +limitation of the current implementation, not a fundamental requirement (see Roadmap below). + +## Supported encoders + +| `encoder_id` | Output | Needs | +|---|---|---| +| `1d` | flattened raw pixels, `(H*W*3,)` uint8 | nothing | +| `2d` | raw pixels, `(H,W,3)` uint8 | nothing | +| `resizex` | bilinear-resized pixels, `(H,W,3)` uint8 | torch + torchvision | +| `cropx` | center-cropped pixels, `(H,W,3)` uint8 | torch + torchvision | +| `r3m18` / `r3m34` / `r3m50` | R3M embedding, `(D,)` | torch + torchvision + R3M | +| `rrl18` / `rrl34` / `rrl50` | ImageNet ResNet embedding, `(D,)` | torch + torchvision | +| `vc1s` / `vc1l` | VC-1 embedding, `(D,)` | torch + vc_models | + +`resize`/`crop` follow torchvision's `T.Resize`/`T.CenterCrop` semantics exactly. Note +`cropx` **zero-pads** (rather than raising) if the requested crop is larger than the +captured `x` — that's `T.CenterCrop`'s native behavior, kept as-is for this minimal +implementation. + +## Why a separate `resize`/`crop` encoder, instead of just setting capture `x`? + +The capture `x` already lets you render at (almost) any size, so it's natural to ask +why a post-capture resize/crop step is needed at all. A few reasons it doesn't collapse into +"just pick the right capture size": + +1. **Hardware cameras don't support arbitrary capture resolutions.** `Robot.get_visual_sensors` + asserts a live camera's stream size exactly matches the requested capture size — real + camera drivers/lenses typically support only a fixed set of native resolutions (often tied + to calibration). If you want anything other than the native stream size in your + observation, something has to resize it after capture. +2. **Resize and crop are different operations, not two spellings of the same thing.** + Capturing at a smaller `x` resamples the *same field of view* at lower pixel density. + It cannot select a sub-region of the scene (e.g. "zoom to where the hand operates"). A true + crop requires capturing at the full resolution/FOV and cutting afterward — there's no + capture-time parameter that does that. +3. **Decoupling capture config from per-task/per-model needs.** Capture resolution is often + fixed by calibration or shared across many envs/tasks; different downstream policies want + different input sizes from the same physical camera. `resize`/`crop` let one capture config + serve many encoder targets without touching the capture-side config per task. +4. **Re-deriving multiple sizes from one recorded dataset.** If frames are recorded once at a + native/high resolution, a resize/crop step (used offline) can derive different observation + sizes after the fact without re-rendering or re-recording. + +## Why torchvision, not OpenCV, for resize/crop + +`torchvision` is already a lazily-guarded optional dependency for this exact feature (used by +the `r3m`/`rrl`/`vc1` transform pipelines already), so `resize`/`crop` add no new dependency. +It also operates on tensors, which keeps the door open for batched/GPU execution and for +chaining directly into a pretrained encoder's tensor pipeline with no extra numpy round-trips. +OpenCV would be a new, fairly heavy dependency whose default interpolation doesn't numerically +match torchvision/PIL's bilinear-with-antialias on downscale — for `rrl`/`r3m`, whose existing +`Resize(256)+CenterCrop(224)` preprocessing was tuned against torchvision, swapping in cv2 +resize could silently shift the input distribution their checkpoints were trained on. \ No newline at end of file diff --git a/robohive/envs/env_base.py b/robohive/envs/env_base.py index a0d36988..bc3cc13f 100644 --- a/robohive/envs/env_base.py +++ b/robohive/envs/env_base.py @@ -6,6 +6,7 @@ ================================================= """ import os +import re import time as timer from sys import platform @@ -212,6 +213,11 @@ def forward(self, x): import_utils.vc_isavailable() from vc_models.models.vit import model_utils as vc + resize_crop_match = re.match(r'^(resize|crop)(\d+)x(\d+)$', id_encoder) + if resize_crop_match: + import_utils.torchvision_isavailable() + import torchvision.transforms as T + # Load encoder prompt("Using {} visual inputs with {} encoder".format(wxh, id_encoder), type=Prompt.INFO) if id_encoder == "1d": @@ -240,6 +246,13 @@ def forward(self, x): model,embd_size,model_transforms,model_info = vc.load_model(vc.VC1_LARGE_NAME) self.rgb_encoder = model self.rgb_transform = model_transforms + elif resize_crop_match: + self.rgb_encoder = IdentityEncoder() + target_h, target_w = int(resize_crop_match.group(2)), int(resize_crop_match.group(3)) + if resize_crop_match.group(1) == "resize": + self.rgb_transform = T.Resize((target_h, target_w), antialias=True) + else: + self.rgb_transform = T.CenterCrop((target_h, target_w)) else: raise ValueError("Unsupported visual encoder: {}".format(id_encoder)) self.rgb_encoder.eval() @@ -363,6 +376,8 @@ def get_visuals(self, sim=None, visual_keys:list=None, device_id:int=None)->dict - 'rgb:cam_name:HxW:r3m18' - 'rgb:cam_name:HxW:r3m34' - 'rgb:cam_name:HxW:r3m50' + - 'rgb:cam_name:HxW:resizex' (e.g. resize128x128) + - 'rgb:cam_name:HxW:cropx' (e.g. crop224x224) """ # return if no visual configured if self.visual_keys == None: @@ -426,6 +441,11 @@ def get_visuals(self, sim=None, visual_keys:list=None, device_id:int=None)->dict rgb_encoded = rgb_encoded.to(self.device_encoder) rgb_encoded = self.rgb_encoder(rgb_encoded).cpu().numpy() rgb_encoded = np.squeeze(rgb_encoded) + elif re.match(r'^(resize|crop)(\d+)x(\d+)$', rgb_encoder_id): + with torch.no_grad(): + img_t = torch.from_numpy(img[0]).permute(2, 0, 1).unsqueeze(0).float() # 1x3xHxW + img_t = self.rgb_transform(img_t) + rgb_encoded = img_t.squeeze(0).permute(1, 2, 0).clamp(0, 255).byte().numpy() # HxWx3 uint8 else: raise ValueError("Unsupported visual encoder: {}".format(rgb_encoder_id))