Add use_latlon support to OlmoEarth fine-tuning wrapper

[pjreddie]

Summary

Adds a use_latlon option to the OlmoEarth fine-tuning wrapper. When enabled, each sample's crop-center lat/lon is computed from its SampleMetadata (crop_bounds + projection CRS, transformed to WGS84) and passed to the encoder via MaskedOlmoEarthSample.latlon, so models pretrained with a geographic (lat/lon) encoding (e.g. the separate/simple-encoding OlmoEarth variants) can use it during fine-tuning.

The latlon encoding's dropout is train-only, and some checkpoints were pretrained with a non-zero latlon_dropout_rate (e.g. rope_simple_v1 used 0.5). To keep the encoding active on every fine-tuning step, the wrapper forces latlon_dropout_rate = 0 when use_latlon is set.

use_latlon defaults to False, so existing configs are unaffected.

Changes

• OlmoEarth.__init__: new use_latlon: bool = False arg; when set (and the loaded encoder exposes composite_encodings.latlon_dropout_rate), force the rate to 0.
• New OlmoEarth._compute_latlon_from_metadata(...): crop-center → WGS84 (lat, lon) per sample.
• OlmoEarth.forward: when use_latlon, inject the computed lat/lon into the sample before the encoder.

Testing

• Loaded a lat/lon-pretrained checkpoint (rope_simple_v1) with use_latlon=True; confirmed latlon_dropout_rate is forced to 0.
• fast_dev_run of a segmentation fine-tune (encoder + UNetDecoder + SegmentationHead) with use_latlon: true, use_legacy_timestamps: false, and UNetDecoder(use_batch_norm: true) completes a train+val step (loss ≈ ln(num_classes) at init, as expected).

🤖 Generated with Claude Code

[favyen2]

We changed band_dropout_rate to be the reverse, where it is initially set to 0.0 and only set to the configured value during training. Can we do the same here, so that users of olmoearth_pretrain don't need to worry about disabling latlon_dropout_rate?