| 33 | |
| 34 | @registry.register_task(name="RearrangePredicateTask-v0") |
| 35 | class RearrangePredicateTask(RearrangeTask): |
| 36 | def __init__(self, *args, sim, config, dataset, **kwargs): |
| 37 | print(f"Num episodes {len(dataset.episodes)}") |
| 38 | |
| 39 | self._all_cls, obj_cats, self._name_to_cls = get_category_info( |
| 40 | config.skip_load_receps |
| 41 | ) |
| 42 | |
| 43 | self.pddl = get_pddl(config, self._all_cls, obj_cats) |
| 44 | self._fix_agent_pos = config.fix_agent_pos |
| 45 | |
| 46 | super().__init__( |
| 47 | *args, |
| 48 | sim=sim, |
| 49 | config=config, |
| 50 | dataset=dataset, |
| 51 | should_place_articulated_agent=not self._fix_agent_pos, |
| 52 | **kwargs, |
| 53 | ) |
| 54 | self._tokenizer = get_parser(config.tokenizer_name) |
| 55 | |
| 56 | self._start_template = self._config.start_template |
| 57 | self._goal_template = self._config.goal_template |
| 58 | self._sample_entities = self._config.sample_entities |
| 59 | self._sample_entities_use_constant_sampling = ( |
| 60 | self._config.sample_entities_use_constant_sampling |
| 61 | ) |
| 62 | self._force_scene_per_worker = self._config.force_scene_per_worker |
| 63 | self._goal_expr = None |
| 64 | self._is_first_reset = True |
| 65 | self._is_freeform = False |
| 66 | |
| 67 | @property |
| 68 | def tokenizer(self): |
| 69 | return self._tokenizer |
| 70 | |
| 71 | @add_perf_timing_func() |
| 72 | def _load_start_goal(self, episode): |
| 73 | """ |
| 74 | Setup the start and goal PDDL conditions. Will change the simulator |
| 75 | state to set the start state. |
| 76 | """ |
| 77 | |
| 78 | pddl_entities = self.pddl.all_entities |
| 79 | self.pddl.bind_to_instance(self._sim, self._dataset, self, episode) |
| 80 | |
| 81 | self._setup_pddl_entities(episode) |
| 82 | |
| 83 | if self._is_first_reset or not self._force_scene_per_worker: |
| 84 | self.pddl.bind_actions() |
| 85 | self._is_first_reset = False |
| 86 | |
| 87 | self._sim.internal_step(-1) |
| 88 | self._load_sampled_names() |
| 89 | |
| 90 | self._goal_expr = self._load_goal_preds(episode) |
| 91 | if self._goal_expr is not None: |
| 92 | t_start = time.time() |
nothing calls this directly
no outgoing calls
no test coverage detected