Optimizing simple query that performs a large table migration
An UPDATE with a co-related subquery is typically quite slow, because the sub-query is run once for each row that is updated. It's typically faster to do that only once. To get the latest status for each train, you can use DISTINCT ON and join to the result of that:
update trains t
  set status = ts.status
from (  
  SELECT distinct on (ts.id) ts.id, ts.status 
  FROM train_statuses ts 
  WHERE status in (?) 
  ORDER BY ts.id, updated_at DESC 
) 
WHERE t.id = ts.id;